From 0bcecffc4c52c52e0502e53b59715ebaba88f998 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Thu, 12 Jun 2025 09:09:13 +0100 Subject: [PATCH 01/14] set timeout for sse in httpx_client_factory (#943) --- src/mcp/client/sse.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index 1ce6a6da4..0c05c6def 100644 --- a/src/mcp/client/sse.py +++ b/src/mcp/client/sse.py @@ -54,12 +54,13 @@ async def sse_client( async with anyio.create_task_group() as tg: try: logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}") - async with httpx_client_factory(headers=headers, auth=auth) as client: + async with httpx_client_factory( + headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout) + ) 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") From df15e9566d36f819762234c96a0e1a3e62037b06 Mon Sep 17 00:00:00 2001 From: Felix Weinberger <3823880+felixweinberger@users.noreply.github.com> Date: Thu, 12 Jun 2025 18:01:53 +0100 Subject: [PATCH 02/14] feat: implement MCP-Protocol-Version header requirement for HTTP transport (#898) --- src/mcp/client/auth.py | 10 +- src/mcp/client/streamable_http.py | 55 +++++++-- src/mcp/server/auth/routes.py | 3 +- src/mcp/server/streamable_http.py | 42 ++++++- src/mcp/types.py | 8 ++ tests/shared/test_streamable_http.py | 171 ++++++++++++++++++++++++++- 6 files changed, 268 insertions(+), 21 deletions(-) diff --git a/src/mcp/client/auth.py b/src/mcp/client/auth.py index 7782022ce..4e777d600 100644 --- a/src/mcp/client/auth.py +++ b/src/mcp/client/auth.py @@ -17,7 +17,13 @@ import anyio import httpx -from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthMetadata, OAuthToken +from mcp.client.streamable_http import MCP_PROTOCOL_VERSION +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthMetadata, + OAuthToken, +) from mcp.types import LATEST_PROTOCOL_VERSION logger = logging.getLogger(__name__) @@ -121,7 +127,7 @@ async def _discover_oauth_metadata(self, server_url: str) -> OAuthMetadata | Non # Extract base URL per MCP spec auth_base_url = self._get_authorization_base_url(server_url) url = urljoin(auth_base_url, "/.well-known/oauth-authorization-server") - headers = {"MCP-Protocol-Version": LATEST_PROTOCOL_VERSION} + headers = {MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION} async with httpx.AsyncClient() as client: try: diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index 471870533..39ac34d8a 100644 --- a/src/mcp/client/streamable_http.py +++ b/src/mcp/client/streamable_http.py @@ -22,6 +22,7 @@ from mcp.shared.message import ClientMessageMetadata, SessionMessage from mcp.types import ( ErrorData, + InitializeResult, JSONRPCError, JSONRPCMessage, JSONRPCNotification, @@ -39,6 +40,7 @@ GetSessionIdCallback = Callable[[], str | None] MCP_SESSION_ID = "mcp-session-id" +MCP_PROTOCOL_VERSION = "mcp-protocol-version" LAST_EVENT_ID = "last-event-id" CONTENT_TYPE = "content-type" ACCEPT = "Accept" @@ -97,17 +99,20 @@ def __init__( ) self.auth = auth self.session_id = None + self.protocol_version = None self.request_headers = { ACCEPT: f"{JSON}, {SSE}", CONTENT_TYPE: JSON, **self.headers, } - def _update_headers_with_session(self, base_headers: dict[str, str]) -> dict[str, str]: - """Update headers with session ID if available.""" + def _prepare_request_headers(self, base_headers: dict[str, str]) -> dict[str, str]: + """Update headers with session ID and protocol version if available.""" headers = base_headers.copy() if self.session_id: headers[MCP_SESSION_ID] = self.session_id + if self.protocol_version: + headers[MCP_PROTOCOL_VERSION] = self.protocol_version return headers def _is_initialization_request(self, message: JSONRPCMessage) -> bool: @@ -128,12 +133,28 @@ def _maybe_extract_session_id_from_response( self.session_id = new_session_id logger.info(f"Received session ID: {self.session_id}") + def _maybe_extract_protocol_version_from_message( + self, + message: JSONRPCMessage, + ) -> None: + """Extract protocol version from initialization response message.""" + if isinstance(message.root, JSONRPCResponse) and message.root.result: + try: + # Parse the result as InitializeResult for type safety + init_result = InitializeResult.model_validate(message.root.result) + self.protocol_version = str(init_result.protocolVersion) + logger.info(f"Negotiated protocol version: {self.protocol_version}") + except Exception as exc: + logger.warning(f"Failed to parse initialization response as InitializeResult: {exc}") + logger.warning(f"Raw result: {message.root.result}") + async def _handle_sse_event( self, sse: ServerSentEvent, read_stream_writer: StreamWriter, original_request_id: RequestId | None = None, resumption_callback: Callable[[str], Awaitable[None]] | None = None, + is_initialization: bool = False, ) -> bool: """Handle an SSE event, returning True if the response is complete.""" if sse.event == "message": @@ -141,6 +162,10 @@ async def _handle_sse_event( message = JSONRPCMessage.model_validate_json(sse.data) logger.debug(f"SSE message: {message}") + # Extract protocol version from initialization response + if is_initialization: + self._maybe_extract_protocol_version_from_message(message) + # If this is a response and we have original_request_id, replace it if original_request_id is not None and isinstance(message.root, JSONRPCResponse | JSONRPCError): message.root.id = original_request_id @@ -174,7 +199,7 @@ async def handle_get_stream( if not self.session_id: return - headers = self._update_headers_with_session(self.request_headers) + headers = self._prepare_request_headers(self.request_headers) async with aconnect_sse( client, @@ -194,7 +219,7 @@ async def handle_get_stream( async def _handle_resumption_request(self, ctx: RequestContext) -> None: """Handle a resumption request using GET with SSE.""" - headers = self._update_headers_with_session(ctx.headers) + headers = self._prepare_request_headers(ctx.headers) if ctx.metadata and ctx.metadata.resumption_token: headers[LAST_EVENT_ID] = ctx.metadata.resumption_token else: @@ -227,7 +252,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None: async def _handle_post_request(self, ctx: RequestContext) -> None: """Handle a POST request with response processing.""" - headers = self._update_headers_with_session(ctx.headers) + headers = self._prepare_request_headers(ctx.headers) message = ctx.session_message.message is_initialization = self._is_initialization_request(message) @@ -256,9 +281,9 @@ async def _handle_post_request(self, ctx: RequestContext) -> None: content_type = response.headers.get(CONTENT_TYPE, "").lower() if content_type.startswith(JSON): - await self._handle_json_response(response, ctx.read_stream_writer) + await self._handle_json_response(response, ctx.read_stream_writer, is_initialization) elif content_type.startswith(SSE): - await self._handle_sse_response(response, ctx) + await self._handle_sse_response(response, ctx, is_initialization) else: await self._handle_unexpected_content_type( content_type, @@ -269,18 +294,29 @@ async def _handle_json_response( self, response: httpx.Response, read_stream_writer: StreamWriter, + is_initialization: bool = False, ) -> None: """Handle JSON response from the server.""" try: content = await response.aread() message = JSONRPCMessage.model_validate_json(content) + + # Extract protocol version from initialization response + if is_initialization: + self._maybe_extract_protocol_version_from_message(message) + session_message = SessionMessage(message) await read_stream_writer.send(session_message) except Exception as exc: logger.error(f"Error parsing JSON response: {exc}") await read_stream_writer.send(exc) - async def _handle_sse_response(self, response: httpx.Response, ctx: RequestContext) -> None: + async def _handle_sse_response( + self, + response: httpx.Response, + ctx: RequestContext, + is_initialization: bool = False, + ) -> None: """Handle SSE response from the server.""" try: event_source = EventSource(response) @@ -289,6 +325,7 @@ async def _handle_sse_response(self, response: httpx.Response, ctx: RequestConte sse, ctx.read_stream_writer, resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None), + is_initialization=is_initialization, ) # If the SSE event indicates completion, like returning respose/error # break the loop @@ -385,7 +422,7 @@ async def terminate_session(self, client: httpx.AsyncClient) -> None: return try: - headers = self._update_headers_with_session(self.request_headers) + headers = self._prepare_request_headers(self.request_headers) response = await client.delete(self.url, headers=headers) if response.status_code == 405: diff --git a/src/mcp/server/auth/routes.py b/src/mcp/server/auth/routes.py index dff468ebd..8647334e0 100644 --- a/src/mcp/server/auth/routes.py +++ b/src/mcp/server/auth/routes.py @@ -16,6 +16,7 @@ from mcp.server.auth.middleware.client_auth import ClientAuthenticator from mcp.server.auth.provider import OAuthAuthorizationServerProvider from mcp.server.auth.settings import ClientRegistrationOptions, RevocationOptions +from mcp.server.streamable_http import MCP_PROTOCOL_VERSION_HEADER from mcp.shared.auth import OAuthMetadata @@ -55,7 +56,7 @@ def cors_middleware( app=request_response(handler), allow_origins="*", allow_methods=allow_methods, - allow_headers=["mcp-protocol-version"], + allow_headers=[MCP_PROTOCOL_VERSION_HEADER], ) return cors_app diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py index 9356a9948..13ee27b32 100644 --- a/src/mcp/server/streamable_http.py +++ b/src/mcp/server/streamable_http.py @@ -25,7 +25,9 @@ from starlette.types import Receive, Scope, Send from mcp.shared.message import ServerMessageMetadata, SessionMessage +from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS from mcp.types import ( + DEFAULT_NEGOTIATED_VERSION, INTERNAL_ERROR, INVALID_PARAMS, INVALID_REQUEST, @@ -45,6 +47,7 @@ # Header names MCP_SESSION_ID_HEADER = "mcp-session-id" +MCP_PROTOCOL_VERSION_HEADER = "mcp-protocol-version" LAST_EVENT_ID_HEADER = "last-event-id" # Content types @@ -293,7 +296,7 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re has_json, has_sse = self._check_accept_headers(request) if not (has_json and has_sse): response = self._create_error_response( - ("Not Acceptable: Client must accept both application/json and " "text/event-stream"), + ("Not Acceptable: Client must accept both application/json and text/event-stream"), HTTPStatus.NOT_ACCEPTABLE, ) await response(scope, receive, send) @@ -353,8 +356,7 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re ) await response(scope, receive, send) return - # For non-initialization requests, validate the session - elif not await self._validate_session(request, send): + elif not await self._validate_request_headers(request, send): return # For notifications and responses only, return 202 Accepted @@ -513,8 +515,9 @@ async def _handle_get_request(self, request: Request, send: Send) -> None: await response(request.scope, request.receive, send) return - if not await self._validate_session(request, send): + if not await self._validate_request_headers(request, send): return + # Handle resumability: check for Last-Event-ID header if last_event_id := request.headers.get(LAST_EVENT_ID_HEADER): await self._replay_events(last_event_id, request, send) @@ -593,7 +596,7 @@ async def _handle_delete_request(self, request: Request, send: Send) -> None: await response(request.scope, request.receive, send) return - if not await self._validate_session(request, send): + if not await self._validate_request_headers(request, send): return await self._terminate_session() @@ -653,6 +656,13 @@ async def _handle_unsupported_request(self, request: Request, send: Send) -> Non ) await response(request.scope, request.receive, send) + async def _validate_request_headers(self, request: Request, send: Send) -> bool: + if not await self._validate_session(request, send): + return False + if not await self._validate_protocol_version(request, send): + return False + return True + async def _validate_session(self, request: Request, send: Send) -> bool: """Validate the session ID in the request.""" if not self.mcp_session_id: @@ -682,6 +692,28 @@ async def _validate_session(self, request: Request, send: Send) -> bool: return True + async def _validate_protocol_version(self, request: Request, send: Send) -> bool: + """Validate the protocol version header in the request.""" + # Get the protocol version from the request headers + protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER) + + # If no protocol version provided, assume default version + if protocol_version is None: + protocol_version = DEFAULT_NEGOTIATED_VERSION + + # Check if the protocol version is supported + if protocol_version not in SUPPORTED_PROTOCOL_VERSIONS: + supported_versions = ", ".join(SUPPORTED_PROTOCOL_VERSIONS) + response = self._create_error_response( + f"Bad Request: Unsupported protocol version: {protocol_version}. " + + f"Supported versions: {supported_versions}", + HTTPStatus.BAD_REQUEST, + ) + await response(request.scope, request.receive, send) + return False + + return True + async def _replay_events(self, last_event_id: str, request: Request, send: Send) -> None: """ Replays events that would have been sent after the specified event ID. diff --git a/src/mcp/types.py b/src/mcp/types.py index 2949ed8e7..824cee73c 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -24,6 +24,14 @@ LATEST_PROTOCOL_VERSION = "2025-03-26" +""" +The default negotiated version of the Model Context Protocol when no version is specified. +We need this to satisfy the MCP specification, which requires the server to assume a +specific version if none is provided by the client. See section "Protocol Version Header" at +https://modelcontextprotocol.io/specification +""" +DEFAULT_NEGOTIATED_VERSION = "2025-03-26" + ProgressToken = str | int Cursor = str Role = Literal["user", "assistant"] diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index 615e68efc..c4604e3fa 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -26,6 +26,7 @@ from mcp.client.streamable_http import streamablehttp_client from mcp.server import Server from mcp.server.streamable_http import ( + MCP_PROTOCOL_VERSION_HEADER, MCP_SESSION_ID_HEADER, SESSION_ID_PATTERN, EventCallback, @@ -64,6 +65,17 @@ } +# Helper functions +def extract_protocol_version_from_sse(response: requests.Response) -> str: + """Extract the negotiated protocol version from an SSE initialization response.""" + assert response.headers.get("Content-Type") == "text/event-stream" + for line in response.text.splitlines(): + if line.startswith("data: "): + init_data = json.loads(line[6:]) + return init_data["result"]["protocolVersion"] + raise ValueError("Could not extract protocol version from SSE response") + + # Simple in-memory event store for testing class SimpleEventStore(EventStore): """Simple in-memory event store for testing.""" @@ -560,11 +572,17 @@ def test_session_termination(basic_server, basic_server_url): ) assert response.status_code == 200 + # Extract negotiated protocol version from SSE response + negotiated_version = extract_protocol_version_from_sse(response) + # Now terminate the session session_id = response.headers.get(MCP_SESSION_ID_HEADER) response = requests.delete( f"{basic_server_url}/mcp", - headers={MCP_SESSION_ID_HEADER: session_id}, + headers={ + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, + }, ) assert response.status_code == 200 @@ -595,16 +613,20 @@ def test_response(basic_server, basic_server_url): ) assert response.status_code == 200 - # Now terminate the session + # Extract negotiated protocol version from SSE response + negotiated_version = extract_protocol_version_from_sse(response) + + # Now get the session ID session_id = response.headers.get(MCP_SESSION_ID_HEADER) - # Try to use the terminated session + # Try to use the session with proper headers tools_response = requests.post( mcp_url, headers={ "Accept": "application/json, text/event-stream", "Content-Type": "application/json", MCP_SESSION_ID_HEADER: session_id, # Use the session ID we got earlier + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, }, json={"jsonrpc": "2.0", "method": "tools/list", "id": "tools-1"}, stream=True, @@ -646,12 +668,23 @@ def test_get_sse_stream(basic_server, basic_server_url): session_id = init_response.headers.get(MCP_SESSION_ID_HEADER) assert session_id is not None + # Extract negotiated protocol version from SSE response + init_data = None + assert init_response.headers.get("Content-Type") == "text/event-stream" + for line in init_response.text.splitlines(): + if line.startswith("data: "): + init_data = json.loads(line[6:]) + break + assert init_data is not None + negotiated_version = init_data["result"]["protocolVersion"] + # Now attempt to establish an SSE stream via GET get_response = requests.get( mcp_url, headers={ "Accept": "text/event-stream", MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, }, stream=True, ) @@ -666,6 +699,7 @@ def test_get_sse_stream(basic_server, basic_server_url): headers={ "Accept": "text/event-stream", MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, }, stream=True, ) @@ -694,11 +728,22 @@ def test_get_validation(basic_server, basic_server_url): session_id = init_response.headers.get(MCP_SESSION_ID_HEADER) assert session_id is not None + # Extract negotiated protocol version from SSE response + init_data = None + assert init_response.headers.get("Content-Type") == "text/event-stream" + for line in init_response.text.splitlines(): + if line.startswith("data: "): + init_data = json.loads(line[6:]) + break + assert init_data is not None + negotiated_version = init_data["result"]["protocolVersion"] + # Test without Accept header response = requests.get( mcp_url, headers={ MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, }, stream=True, ) @@ -711,6 +756,7 @@ def test_get_validation(basic_server, basic_server_url): headers={ "Accept": "application/json", MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, }, ) assert response.status_code == 406 @@ -1004,6 +1050,7 @@ async def test_streamablehttp_client_resumption(event_server): captured_resumption_token = None captured_notifications = [] tool_started = False + captured_protocol_version = None async def message_handler( message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception, @@ -1032,6 +1079,8 @@ async def on_resumption_token_update(token: str) -> None: assert isinstance(result, InitializeResult) captured_session_id = get_session_id() assert captured_session_id is not None + # Capture the negotiated protocol version + captured_protocol_version = result.protocolVersion # Start a long-running tool in a task async with anyio.create_task_group() as tg: @@ -1064,10 +1113,12 @@ async def run_tool(): captured_notifications_pre = captured_notifications.copy() captured_notifications = [] - # Now resume the session with the same mcp-session-id + # Now resume the session with the same mcp-session-id and protocol version headers = {} if captured_session_id: headers[MCP_SESSION_ID_HEADER] = captured_session_id + if captured_protocol_version: + headers[MCP_PROTOCOL_VERSION_HEADER] = captured_protocol_version async with streamablehttp_client(f"{server_url}/mcp", headers=headers) as ( read_stream, @@ -1358,3 +1409,115 @@ async def test_streamablehttp_request_context_isolation(context_aware_server: No assert ctx["headers"].get("x-request-id") == f"request-{i}" assert ctx["headers"].get("x-custom-value") == f"value-{i}" assert ctx["headers"].get("authorization") == f"Bearer token-{i}" + + +@pytest.mark.anyio +async def test_client_includes_protocol_version_header_after_init(context_aware_server, basic_server_url): + """Test that client includes mcp-protocol-version header after initialization.""" + async with streamablehttp_client(f"{basic_server_url}/mcp") as ( + read_stream, + write_stream, + _, + ): + async with ClientSession(read_stream, write_stream) as session: + # Initialize and get the negotiated version + init_result = await session.initialize() + negotiated_version = init_result.protocolVersion + + # Call a tool that echoes headers to verify the header is present + tool_result = await session.call_tool("echo_headers", {}) + + assert len(tool_result.content) == 1 + assert isinstance(tool_result.content[0], TextContent) + headers_data = json.loads(tool_result.content[0].text) + + # Verify protocol version header is present + assert "mcp-protocol-version" in headers_data + assert headers_data[MCP_PROTOCOL_VERSION_HEADER] == negotiated_version + + +def test_server_validates_protocol_version_header(basic_server, basic_server_url): + """Test that server returns 400 Bad Request version if header unsupported or invalid.""" + # First initialize a session to get a valid session ID + init_response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + session_id = init_response.headers.get(MCP_SESSION_ID_HEADER) + + # Test request with invalid protocol version (should fail) + response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: "invalid-version", + }, + json={"jsonrpc": "2.0", "method": "tools/list", "id": "test-2"}, + ) + assert response.status_code == 400 + assert MCP_PROTOCOL_VERSION_HEADER in response.text or "protocol version" in response.text.lower() + + # Test request with unsupported protocol version (should fail) + response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: "1999-01-01", # Very old unsupported version + }, + json={"jsonrpc": "2.0", "method": "tools/list", "id": "test-3"}, + ) + assert response.status_code == 400 + assert MCP_PROTOCOL_VERSION_HEADER in response.text or "protocol version" in response.text.lower() + + # Test request with valid protocol version (should succeed) + negotiated_version = extract_protocol_version_from_sse(init_response) + + response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: session_id, + MCP_PROTOCOL_VERSION_HEADER: negotiated_version, + }, + json={"jsonrpc": "2.0", "method": "tools/list", "id": "test-4"}, + ) + assert response.status_code == 200 + + +def test_server_backwards_compatibility_no_protocol_version(basic_server, basic_server_url): + """Test server accepts requests without protocol version header.""" + # First initialize a session to get a valid session ID + init_response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + json=INIT_REQUEST, + ) + assert init_response.status_code == 200 + session_id = init_response.headers.get(MCP_SESSION_ID_HEADER) + + # Test request without mcp-protocol-version header (backwards compatibility) + response = requests.post( + f"{basic_server_url}/mcp", + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + MCP_SESSION_ID_HEADER: session_id, + }, + json={"jsonrpc": "2.0", "method": "tools/list", "id": "test-backwards-compat"}, + stream=True, + ) + assert response.status_code == 200 # Should succeed for backwards compatibility + assert response.headers.get("Content-Type") == "text/event-stream" From 1bb16fa96ebb7440f481c47798f2e51878b78bae Mon Sep 17 00:00:00 2001 From: Dipesh Rathore <60871161+DarkDk123@users.noreply.github.com> Date: Fri, 13 Jun 2025 14:11:18 +0530 Subject: [PATCH 03/14] clean-up: removed unused ci file (#950) --- .github/workflows/check-lock.yml | 25 ------------------------- 1 file changed, 25 deletions(-) delete mode 100644 .github/workflows/check-lock.yml diff --git a/.github/workflows/check-lock.yml b/.github/workflows/check-lock.yml deleted file mode 100644 index 805b0f3cc..000000000 --- a/.github/workflows/check-lock.yml +++ /dev/null @@ -1,25 +0,0 @@ -name: Check uv.lock - -on: - pull_request: - paths: - - "pyproject.toml" - - "uv.lock" - push: - paths: - - "pyproject.toml" - - "uv.lock" - -jobs: - check-lock: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Install uv - run: | - curl -LsSf https://astral.sh/uv/install.sh | sh - echo "$HOME/.cargo/bin" >> $GITHUB_PATH - - - name: Check uv.lock is up to date - run: uv lock --check From bac27054265a4a71de2da1c4e540e837281202f6 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Fri, 13 Jun 2025 16:24:28 +0200 Subject: [PATCH 04/14] ci: add timeout on the test job (#955) --- .github/workflows/shared.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/shared.yml b/.github/workflows/shared.yml index b5e880aff..499871ca1 100644 --- a/.github/workflows/shared.yml +++ b/.github/workflows/shared.yml @@ -25,6 +25,7 @@ jobs: test: runs-on: ${{ matrix.os }} + timeout-minutes: 10 strategy: matrix: python-version: ["3.10", "3.11", "3.12", "3.13"] From f3cd20c9200de001fb58f83a130ab83c6a6ed5fd Mon Sep 17 00:00:00 2001 From: ihrpr Date: Fri, 13 Jun 2025 18:25:36 +0100 Subject: [PATCH 05/14] Rename ResourceReference to ResourceTemplateReference (#947) --- src/mcp/client/session.py | 2 +- src/mcp/server/lowlevel/server.py | 2 +- src/mcp/types.py | 10 ++++++++-- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 02a4e8e01..0ba46fe5b 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -302,7 +302,7 @@ async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) - async def complete( self, - ref: types.ResourceReference | types.PromptReference, + ref: types.ResourceTemplateReference | types.PromptReference, argument: dict[str, str], ) -> types.CompleteResult: """Send a completion/complete request.""" diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 7a24781ab..b077339fb 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -431,7 +431,7 @@ def completion(self): def decorator( func: Callable[ [ - types.PromptReference | types.ResourceReference, + types.PromptReference | types.ResourceTemplateReference, types.CompletionArgument, ], Awaitable[types.Completion | None], diff --git a/src/mcp/types.py b/src/mcp/types.py index 824cee73c..e1f8ea52f 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -3,6 +3,7 @@ from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel from pydantic.networks import AnyUrl, UrlConstraints +from typing_extensions import deprecated """ Model Context Protocol bindings for Python @@ -961,7 +962,7 @@ class CreateMessageResult(Result): """The reason why sampling stopped, if known.""" -class ResourceReference(BaseModel): +class ResourceTemplateReference(BaseModel): """A reference to a resource or resource template definition.""" type: Literal["ref/resource"] @@ -970,6 +971,11 @@ class ResourceReference(BaseModel): model_config = ConfigDict(extra="allow") +@deprecated("`ResourceReference` is deprecated, you should use `ResourceTemplateReference`.") +class ResourceReference(ResourceTemplateReference): + pass + + class PromptReference(BaseModel): """Identifies a prompt.""" @@ -992,7 +998,7 @@ class CompletionArgument(BaseModel): class CompleteRequestParams(RequestParams): """Parameters for completion requests.""" - ref: ResourceReference | PromptReference + ref: ResourceTemplateReference | PromptReference argument: CompletionArgument model_config = ConfigDict(extra="allow") From 1eb1bba83c70c3121bce7fc0263e5fac2c3f0520 Mon Sep 17 00:00:00 2001 From: Felix Weinberger <3823880+felixweinberger@users.noreply.github.com> Date: Mon, 16 Jun 2025 11:43:54 -0400 Subject: [PATCH 06/14] feat: add _meta to more objects (#963) --- src/mcp/types.py | 49 ++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 43 insertions(+), 6 deletions(-) diff --git a/src/mcp/types.py b/src/mcp/types.py index e1f8ea52f..bf467e37c 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -69,8 +69,7 @@ class Meta(BaseModel): meta: Meta | None = Field(alias="_meta", default=None) """ - This parameter name is reserved by MCP to allow clients and servers to attach - additional metadata to their notifications. + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. """ @@ -105,13 +104,11 @@ class Notification(BaseModel, Generic[NotificationParamsT, MethodT]): class Result(BaseModel): """Base class for JSON-RPC results.""" - model_config = ConfigDict(extra="allow") - meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - This result property is reserved by the protocol to allow clients and servers to - attach additional metadata to their responses. + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. """ + model_config = ConfigDict(extra="allow") class PaginatedResult(Result): @@ -394,6 +391,10 @@ class Resource(BaseModel): This can be used by Hosts to display file sizes and estimate context window usage. """ annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -415,6 +416,10 @@ class ResourceTemplate(BaseModel): included if all resources matching this template have the same type. """ annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -461,6 +466,10 @@ class ResourceContents(BaseModel): """The URI of this resource.""" mimeType: str | None = None """The MIME type of this resource, if known.""" + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -588,6 +597,10 @@ class Prompt(BaseModel): """An optional description of what this prompt provides.""" arguments: list[PromptArgument] | None = None """A list of arguments to use for templating the prompt.""" + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -621,6 +634,10 @@ class TextContent(BaseModel): text: str """The text content of the message.""" annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -636,6 +653,10 @@ class ImageContent(BaseModel): image types. """ annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -651,6 +672,10 @@ class AudioContent(BaseModel): audio types. """ annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -673,6 +698,10 @@ class EmbeddedResource(BaseModel): type: Literal["resource"] resource: TextResourceContents | BlobResourceContents annotations: Annotations | None = None + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -772,6 +801,10 @@ class Tool(BaseModel): """A JSON Schema object defining the expected parameters for the tool.""" annotations: ToolAnnotations | None = None """Optional additional tool information.""" + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") @@ -1064,6 +1097,10 @@ class Root(BaseModel): identifier for the root, which may be useful for display purposes or for referencing the root in other parts of the application. """ + meta: dict[str, Any] | None = Field(alias="_meta", default=None) + """ + Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + """ model_config = ConfigDict(extra="allow") From 7b420656de48cfdb90b39eb582e60b6d55c2f891 Mon Sep 17 00:00:00 2001 From: David Dworken Date: Mon, 16 Jun 2025 17:04:51 -0700 Subject: [PATCH 07/14] Fix uncaught exception in MCP server (#967) --- src/mcp/shared/session.py | 183 +++++++++++++++------------ tests/shared/test_streamable_http.py | 37 ++++++ 2 files changed, 137 insertions(+), 83 deletions(-) diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py index 294986acb..6536272d9 100644 --- a/src/mcp/shared/session.py +++ b/src/mcp/shared/session.py @@ -333,90 +333,107 @@ async def _receive_loop(self) -> None: self._read_stream, self._write_stream, ): - async for message in self._read_stream: - if isinstance(message, Exception): - await self._handle_incoming(message) - elif isinstance(message.message.root, JSONRPCRequest): - try: - validated_request = self._receive_request_type.model_validate( - message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - responder = RequestResponder( - request_id=message.message.root.id, - request_meta=validated_request.root.params.meta if validated_request.root.params else None, - request=validated_request, - session=self, - on_complete=lambda r: self._in_flight.pop(r.request_id, None), - message_metadata=message.metadata, - ) - self._in_flight[responder.request_id] = responder - await self._received_request(responder) - - if not responder._completed: # type: ignore[reportPrivateUsage] - await self._handle_incoming(responder) - except Exception as e: - # For request validation errors, send a proper JSON-RPC error - # response instead of crashing the server - logging.warning(f"Failed to validate request: {e}") - logging.debug(f"Message that failed validation: {message.message.root}") - error_response = JSONRPCError( - jsonrpc="2.0", - id=message.message.root.id, - error=ErrorData( - code=INVALID_PARAMS, - message="Invalid request parameters", - data="", - ), - ) - session_message = SessionMessage(message=JSONRPCMessage(error_response)) - await self._write_stream.send(session_message) - - elif isinstance(message.message.root, JSONRPCNotification): - try: - notification = self._receive_notification_type.model_validate( - message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True) - ) - # Handle cancellation notifications - if isinstance(notification.root, CancelledNotification): - cancelled_id = notification.root.params.requestId - if cancelled_id in self._in_flight: - await self._in_flight[cancelled_id].cancel() + try: + async for message in self._read_stream: + if isinstance(message, Exception): + await self._handle_incoming(message) + elif isinstance(message.message.root, JSONRPCRequest): + try: + validated_request = self._receive_request_type.model_validate( + message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + responder = RequestResponder( + request_id=message.message.root.id, + request_meta=validated_request.root.params.meta + if validated_request.root.params + else None, + request=validated_request, + session=self, + on_complete=lambda r: self._in_flight.pop(r.request_id, None), + message_metadata=message.metadata, + ) + self._in_flight[responder.request_id] = responder + await self._received_request(responder) + + if not responder._completed: # type: ignore[reportPrivateUsage] + await self._handle_incoming(responder) + except Exception as e: + # For request validation errors, send a proper JSON-RPC error + # response instead of crashing the server + logging.warning(f"Failed to validate request: {e}") + logging.debug(f"Message that failed validation: {message.message.root}") + error_response = JSONRPCError( + jsonrpc="2.0", + id=message.message.root.id, + error=ErrorData( + code=INVALID_PARAMS, + message="Invalid request parameters", + data="", + ), + ) + session_message = SessionMessage(message=JSONRPCMessage(error_response)) + await self._write_stream.send(session_message) + + elif isinstance(message.message.root, JSONRPCNotification): + try: + notification = self._receive_notification_type.model_validate( + message.message.root.model_dump(by_alias=True, mode="json", exclude_none=True) + ) + # Handle cancellation notifications + if isinstance(notification.root, CancelledNotification): + cancelled_id = notification.root.params.requestId + if cancelled_id in self._in_flight: + await self._in_flight[cancelled_id].cancel() + else: + # Handle progress notifications callback + if isinstance(notification.root, ProgressNotification): + progress_token = notification.root.params.progressToken + # If there is a progress callback for this token, + # call it with the progress information + if progress_token in self._progress_callbacks: + callback = self._progress_callbacks[progress_token] + await callback( + notification.root.params.progress, + notification.root.params.total, + notification.root.params.message, + ) + await self._received_notification(notification) + await self._handle_incoming(notification) + except Exception as e: + # For other validation errors, log and continue + logging.warning( + f"Failed to validate notification: {e}. " f"Message was: {message.message.root}" + ) + else: # Response or error + stream = self._response_streams.pop(message.message.root.id, None) + if stream: + await stream.send(message.message.root) else: - # Handle progress notifications callback - if isinstance(notification.root, ProgressNotification): - progress_token = notification.root.params.progressToken - # If there is a progress callback for this token, - # call it with the progress information - if progress_token in self._progress_callbacks: - callback = self._progress_callbacks[progress_token] - await callback( - notification.root.params.progress, - notification.root.params.total, - notification.root.params.message, - ) - await self._received_notification(notification) - await self._handle_incoming(notification) - except Exception as e: - # For other validation errors, log and continue - logging.warning( - f"Failed to validate notification: {e}. " f"Message was: {message.message.root}" - ) - else: # Response or error - stream = self._response_streams.pop(message.message.root.id, None) - if stream: - await stream.send(message.message.root) - else: - await self._handle_incoming( - RuntimeError("Received response with an unknown " f"request ID: {message}") - ) - - # after the read stream is closed, we need to send errors - # to any pending requests - for id, stream in self._response_streams.items(): - error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") - await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error)) - await stream.aclose() - self._response_streams.clear() + await self._handle_incoming( + RuntimeError("Received response with an unknown " f"request ID: {message}") + ) + + except anyio.ClosedResourceError: + # This is expected when the client disconnects abruptly. + # Without this handler, the exception would propagate up and + # crash the server's task group. + logging.debug("Read stream closed by client") + except Exception as e: + # Other exceptions are not expected and should be logged. We purposefully + # catch all exceptions here to avoid crashing the server. + logging.exception(f"Unhandled exception in receive loop: {e}") + finally: + # after the read stream is closed, we need to send errors + # to any pending requests + for id, stream in self._response_streams.items(): + error = ErrorData(code=CONNECTION_CLOSED, message="Connection closed") + try: + await stream.send(JSONRPCError(jsonrpc="2.0", id=id, error=error)) + await stream.aclose() + except Exception: + # Stream might already be closed + pass + self._response_streams.clear() async def _received_request(self, responder: RequestResponder[ReceiveRequestT, SendResultT]) -> None: """ diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index c4604e3fa..68d58fc91 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -1521,3 +1521,40 @@ def test_server_backwards_compatibility_no_protocol_version(basic_server, basic_ ) assert response.status_code == 200 # Should succeed for backwards compatibility assert response.headers.get("Content-Type") == "text/event-stream" + + +@pytest.mark.anyio +async def test_client_crash_handled(basic_server, basic_server_url): + """Test that cases where the client crashes are handled gracefully.""" + + # Simulate bad client that crashes after init + async def bad_client(): + """Client that triggers ClosedResourceError""" + async with streamablehttp_client(f"{basic_server_url}/mcp") as ( + read_stream, + write_stream, + _, + ): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + raise Exception("client crash") + + # Run bad client a few times to trigger the crash + for _ in range(3): + try: + await bad_client() + except Exception: + pass + await anyio.sleep(0.1) + + # Try a good client, it should still be able to connect and list tools + async with streamablehttp_client(f"{basic_server_url}/mcp") as ( + read_stream, + write_stream, + _, + ): + async with ClientSession(read_stream, write_stream) as session: + result = await session.initialize() + assert isinstance(result, InitializeResult) + tools = await session.list_tools() + assert tools.tools From a3bcabdce27763f214faf5f6e0a42e67dee6b278 Mon Sep 17 00:00:00 2001 From: Inna Harper Date: Tue, 17 Jun 2025 09:33:25 +0100 Subject: [PATCH 08/14] Include context into completions (#966) --- README.md | 63 +++++++ src/mcp/client/session.py | 6 + src/mcp/server/fastmcp/server.py | 18 ++ src/mcp/server/lowlevel/server.py | 3 +- src/mcp/types.py | 10 ++ tests/server/fastmcp/test_integration.py | 90 +++++++++- tests/server/test_completion_with_context.py | 180 +++++++++++++++++++ 7 files changed, 363 insertions(+), 7 deletions(-) create mode 100644 tests/server/test_completion_with_context.py diff --git a/README.md b/README.md index d76d3d267..38154c878 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ - [Prompts](#prompts) - [Images](#images) - [Context](#context) + - [Completions](#completions) - [Running Your Server](#running-your-server) - [Development Mode](#development-mode) - [Claude Desktop Integration](#claude-desktop-integration) @@ -310,6 +311,68 @@ async def long_task(files: list[str], ctx: Context) -> str: return "Processing complete" ``` +### Completions + +MCP supports providing completion suggestions for prompt arguments and resource template parameters. With the context parameter, servers can provide completions based on previously resolved values: + +Client usage: +```python +from mcp.client.session import ClientSession +from mcp.types import ResourceTemplateReference + + +async def use_completion(session: ClientSession): + # Complete without context + result = await session.complete( + ref=ResourceTemplateReference( + type="ref/resource", uri="github://repos/{owner}/{repo}" + ), + argument={"name": "owner", "value": "model"}, + ) + + # Complete with context - repo suggestions based on owner + result = await session.complete( + ref=ResourceTemplateReference( + type="ref/resource", uri="github://repos/{owner}/{repo}" + ), + argument={"name": "repo", "value": "test"}, + context_arguments={"owner": "modelcontextprotocol"}, + ) +``` + +Server implementation: +```python +from mcp.server import Server +from mcp.types import ( + Completion, + CompletionArgument, + CompletionContext, + PromptReference, + ResourceTemplateReference, +) + +server = Server("example-server") + + +@server.completion() +async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, +) -> Completion | None: + if isinstance(ref, ResourceTemplateReference): + if ref.uri == "github://repos/{owner}/{repo}" and argument.name == "repo": + # Use context to provide owner-specific repos + if context and context.arguments: + owner = context.arguments.get("owner") + if owner == "modelcontextprotocol": + repos = ["python-sdk", "typescript-sdk", "specification"] + # Filter based on partial input + filtered = [r for r in repos if r.startswith(argument.value)] + return Completion(values=filtered) + return None +``` + ### Authentication Authentication can be used by servers that want to expose tools accessing protected resources. diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 0ba46fe5b..8b819ad6d 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -304,8 +304,13 @@ async def complete( self, ref: types.ResourceTemplateReference | types.PromptReference, argument: dict[str, str], + context_arguments: dict[str, str] | None = None, ) -> types.CompleteResult: """Send a completion/complete request.""" + context = None + if context_arguments is not None: + context = types.CompletionContext(arguments=context_arguments) + return await self.send_request( types.ClientRequest( types.CompleteRequest( @@ -313,6 +318,7 @@ async def complete( params=types.CompleteRequestParams( ref=ref, argument=types.CompletionArgument(**argument), + context=context, ), ) ), diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index f74566303..a62974bc9 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -364,6 +364,24 @@ def decorator(fn: AnyFunction) -> AnyFunction: return decorator + def completion(self): + """Decorator to register a completion handler. + + The completion handler receives: + - ref: PromptReference or ResourceTemplateReference + - argument: CompletionArgument with name and partial value + - context: Optional CompletionContext with previously resolved arguments + + Example: + @mcp.completion() + async def handle_completion(ref, argument, context): + if isinstance(ref, ResourceTemplateReference): + # Return completions based on ref, argument, and context + return Completion(values=["option1", "option2"]) + return None + """ + return self._mcp_server.completion() + def add_resource(self, resource: Resource) -> None: """Add a resource to the server. diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index b077339fb..617a4813d 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -433,6 +433,7 @@ def decorator( [ types.PromptReference | types.ResourceTemplateReference, types.CompletionArgument, + types.CompletionContext | None, ], Awaitable[types.Completion | None], ], @@ -440,7 +441,7 @@ def decorator( logger.debug("Registering handler for CompleteRequest") async def handler(req: types.CompleteRequest): - completion = await func(req.params.ref, req.params.argument) + completion = await func(req.params.ref, req.params.argument, req.params.context) return types.ServerResult( types.CompleteResult( completion=completion diff --git a/src/mcp/types.py b/src/mcp/types.py index bf467e37c..be8d7326f 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -1028,11 +1028,21 @@ class CompletionArgument(BaseModel): model_config = ConfigDict(extra="allow") +class CompletionContext(BaseModel): + """Additional, optional context for completions.""" + + arguments: dict[str, str] | None = None + """Previously-resolved variables in a URI template or prompt.""" + model_config = ConfigDict(extra="allow") + + class CompleteRequestParams(RequestParams): """Parameters for completion requests.""" ref: ResourceTemplateReference | PromptReference argument: CompletionArgument + context: CompletionContext | None = None + """Additional, optional context for completions""" model_config = ConfigDict(extra="allow") diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index 84c59cbda..3caf994e0 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -18,7 +18,6 @@ from starlette.applications import Starlette from starlette.requests import Request -import mcp.types as types from mcp.client.session import ClientSession from mcp.client.sse import sse_client from mcp.client.streamable_http import streamablehttp_client @@ -26,14 +25,24 @@ from mcp.server.fastmcp.resources import FunctionResource from mcp.shared.context import RequestContext from mcp.types import ( + Completion, + CompletionArgument, + CompletionContext, CreateMessageRequestParams, CreateMessageResult, GetPromptResult, InitializeResult, + LoggingMessageNotification, + ProgressNotification, + PromptReference, ReadResourceResult, + ResourceListChangedNotification, + ResourceTemplateReference, SamplingMessage, + ServerNotification, TextContent, TextResourceContents, + ToolListChangedNotification, ) @@ -191,6 +200,40 @@ def complex_prompt(user_query: str, context: str = "general") -> str: # Since FastMCP doesn't support system messages in the same way return f"Context: {context}. Query: {user_query}" + # Resource template with completion support + @mcp.resource("github://repos/{owner}/{repo}") + def github_repo_resource(owner: str, repo: str) -> str: + return f"Repository: {owner}/{repo}" + + # Add completion handler for the server + @mcp.completion() + async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, + ) -> Completion | None: + # Handle GitHub repository completion + if isinstance(ref, ResourceTemplateReference): + if ref.uri == "github://repos/{owner}/{repo}" and argument.name == "repo": + if context and context.arguments and context.arguments.get("owner") == "modelcontextprotocol": + # Return repos for modelcontextprotocol org + return Completion(values=["python-sdk", "typescript-sdk", "specification"], total=3, hasMore=False) + elif context and context.arguments and context.arguments.get("owner") == "test-org": + # Return repos for test-org + return Completion(values=["test-repo1", "test-repo2"], total=2, hasMore=False) + + # Handle prompt completions + if isinstance(ref, PromptReference): + if ref.name == "complex_prompt" and argument.name == "context": + # Complete context values + contexts = ["general", "technical", "business", "academic"] + return Completion( + values=[c for c in contexts if c.startswith(argument.value)], total=None, hasMore=False + ) + + # Default: no completion available + return Completion(values=[], total=0, hasMore=False) + # Tool that echoes request headers from context @mcp.tool(description="Echo request headers from context") def echo_headers(ctx: Context[Any, Any, Request]) -> str: @@ -597,15 +640,15 @@ async def handle_tool_list_changed(self, params) -> None: async def handle_generic_notification(self, message) -> None: # Check if this is a ServerNotification - if isinstance(message, types.ServerNotification): + if isinstance(message, ServerNotification): # Check the specific notification type - if isinstance(message.root, types.ProgressNotification): + if isinstance(message.root, ProgressNotification): await self.handle_progress(message.root.params) - elif isinstance(message.root, types.LoggingMessageNotification): + elif isinstance(message.root, LoggingMessageNotification): await self.handle_log(message.root.params) - elif isinstance(message.root, types.ResourceListChangedNotification): + elif isinstance(message.root, ResourceListChangedNotification): await self.handle_resource_list_changed(message.root.params) - elif isinstance(message.root, types.ToolListChangedNotification): + elif isinstance(message.root, ToolListChangedNotification): await self.handle_tool_list_changed(message.root.params) @@ -781,6 +824,41 @@ async def progress_callback(progress: float, total: float | None, message: str | if context_data["method"]: assert context_data["method"] == "POST" + # Test completion functionality + # 1. Test resource template completion with context + repo_result = await session.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), + argument={"name": "repo", "value": ""}, + context_arguments={"owner": "modelcontextprotocol"}, + ) + assert repo_result.completion.values == ["python-sdk", "typescript-sdk", "specification"] + assert repo_result.completion.total == 3 + assert repo_result.completion.hasMore is False + + # 2. Test with different context + repo_result2 = await session.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), + argument={"name": "repo", "value": ""}, + context_arguments={"owner": "test-org"}, + ) + assert repo_result2.completion.values == ["test-repo1", "test-repo2"] + assert repo_result2.completion.total == 2 + + # 3. Test prompt argument completion + context_result = await session.complete( + ref=PromptReference(type="ref/prompt", name="complex_prompt"), + argument={"name": "context", "value": "tech"}, + ) + assert "technical" in context_result.completion.values + + # 4. Test completion without context (should return empty) + no_context_result = await session.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="github://repos/{owner}/{repo}"), + argument={"name": "repo", "value": "test"}, + ) + assert no_context_result.completion.values == [] + assert no_context_result.completion.total == 0 + async def sampling_callback( context: RequestContext[ClientSession, None], diff --git a/tests/server/test_completion_with_context.py b/tests/server/test_completion_with_context.py new file mode 100644 index 000000000..f0d154587 --- /dev/null +++ b/tests/server/test_completion_with_context.py @@ -0,0 +1,180 @@ +""" +Tests for completion handler with context functionality. +""" + +import pytest + +from mcp.server.lowlevel import Server +from mcp.shared.memory import create_connected_server_and_client_session +from mcp.types import ( + Completion, + CompletionArgument, + CompletionContext, + PromptReference, + ResourceTemplateReference, +) + + +@pytest.mark.anyio +async def test_completion_handler_receives_context(): + """Test that the completion handler receives context correctly.""" + server = Server("test-server") + + # Track what the handler receives + received_args = {} + + @server.completion() + async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, + ) -> Completion | None: + received_args["ref"] = ref + received_args["argument"] = argument + received_args["context"] = context + + # Return test completion + return Completion(values=["test-completion"], total=1, hasMore=False) + + async with create_connected_server_and_client_session(server) as client: + # Test with context + result = await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="test://resource/{param}"), + argument={"name": "param", "value": "test"}, + context_arguments={"previous": "value"}, + ) + + # Verify handler received the context + assert received_args["context"] is not None + assert received_args["context"].arguments == {"previous": "value"} + assert result.completion.values == ["test-completion"] + + +@pytest.mark.anyio +async def test_completion_backward_compatibility(): + """Test that completion works without context (backward compatibility).""" + server = Server("test-server") + + context_was_none = False + + @server.completion() + async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, + ) -> Completion | None: + nonlocal context_was_none + context_was_none = context is None + + return Completion(values=["no-context-completion"], total=1, hasMore=False) + + async with create_connected_server_and_client_session(server) as client: + # Test without context + result = await client.complete( + ref=PromptReference(type="ref/prompt", name="test-prompt"), argument={"name": "arg", "value": "val"} + ) + + # Verify context was None + assert context_was_none + assert result.completion.values == ["no-context-completion"] + + +@pytest.mark.anyio +async def test_dependent_completion_scenario(): + """Test a real-world scenario with dependent completions.""" + server = Server("test-server") + + @server.completion() + async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, + ) -> Completion | None: + # Simulate database/table completion scenario + if isinstance(ref, ResourceTemplateReference): + if ref.uri == "db://{database}/{table}": + if argument.name == "database": + # Complete database names + return Completion(values=["users_db", "products_db", "analytics_db"], total=3, hasMore=False) + elif argument.name == "table": + # Complete table names based on selected database + if context and context.arguments: + db = context.arguments.get("database") + if db == "users_db": + return Completion(values=["users", "sessions", "permissions"], total=3, hasMore=False) + elif db == "products_db": + return Completion(values=["products", "categories", "inventory"], total=3, hasMore=False) + + return Completion(values=[], total=0, hasMore=False) + + async with create_connected_server_and_client_session(server) as client: + # First, complete database + db_result = await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"), + argument={"name": "database", "value": ""}, + ) + assert "users_db" in db_result.completion.values + assert "products_db" in db_result.completion.values + + # Then complete table with database context + table_result = await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"), + argument={"name": "table", "value": ""}, + context_arguments={"database": "users_db"}, + ) + assert table_result.completion.values == ["users", "sessions", "permissions"] + + # Different database gives different tables + table_result2 = await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"), + argument={"name": "table", "value": ""}, + context_arguments={"database": "products_db"}, + ) + assert table_result2.completion.values == ["products", "categories", "inventory"] + + +@pytest.mark.anyio +async def test_completion_error_on_missing_context(): + """Test that server can raise error when required context is missing.""" + server = Server("test-server") + + @server.completion() + async def handle_completion( + ref: PromptReference | ResourceTemplateReference, + argument: CompletionArgument, + context: CompletionContext | None, + ) -> Completion | None: + if isinstance(ref, ResourceTemplateReference): + if ref.uri == "db://{database}/{table}": + if argument.name == "table": + # Check if database context is provided + if not context or not context.arguments or "database" not in context.arguments: + # Raise an error instead of returning error as completion + raise ValueError("Please select a database first to see available tables") + # Normal completion if context is provided + db = context.arguments.get("database") + if db == "test_db": + return Completion(values=["users", "orders", "products"], total=3, hasMore=False) + + return Completion(values=[], total=0, hasMore=False) + + async with create_connected_server_and_client_session(server) as client: + # Try to complete table without database context - should raise error + with pytest.raises(Exception) as exc_info: + await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"), + argument={"name": "table", "value": ""}, + ) + + # Verify error message + assert "Please select a database first" in str(exc_info.value) + + # Now complete with proper context - should work normally + result_with_context = await client.complete( + ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"), + argument={"name": "table", "value": ""}, + context_arguments={"database": "test_db"}, + ) + + # Should get normal completions + assert result_with_context.completion.values == ["users", "orders", "products"] From 69e6572d489bac65060f4dac470cb732fa2ab3de Mon Sep 17 00:00:00 2001 From: Inna Harper Date: Tue, 17 Jun 2025 11:54:11 +0100 Subject: [PATCH 09/14] Add support for Elicitation (#625) Co-authored-by: David Soria Parra --- README.md | 48 +++++- src/mcp/client/session.py | 30 ++++ src/mcp/server/elicitation.py | 111 ++++++++++++ src/mcp/server/fastmcp/server.py | 32 ++++ src/mcp/server/session.py | 33 ++++ src/mcp/shared/memory.py | 11 +- src/mcp/types.py | 50 +++++- tests/server/fastmcp/test_elicitation.py | 210 +++++++++++++++++++++++ tests/server/fastmcp/test_integration.py | 126 +++++++++++++- 9 files changed, 643 insertions(+), 8 deletions(-) create mode 100644 src/mcp/server/elicitation.py create mode 100644 tests/server/fastmcp/test_elicitation.py diff --git a/README.md b/README.md index 38154c878..1bdae1d20 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ - [Images](#images) - [Context](#context) - [Completions](#completions) + - [Elicitation](#elicitation) + - [Authentication](#authentication) - [Running Your Server](#running-your-server) - [Development Mode](#development-mode) - [Claude Desktop Integration](#claude-desktop-integration) @@ -74,7 +76,7 @@ The Model Context Protocol allows applications to provide context for LLMs in a ### Adding MCP to your python project -We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects. +We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects. If you haven't created a uv-managed project yet, create one: @@ -372,6 +374,50 @@ async def handle_completion( return Completion(values=filtered) return None ``` +### Elicitation + +Request additional information from users during tool execution: + +```python +from mcp.server.fastmcp import FastMCP, Context +from mcp.server.elicitation import ( + AcceptedElicitation, + DeclinedElicitation, + CancelledElicitation, +) +from pydantic import BaseModel, Field + +mcp = FastMCP("Booking System") + + +@mcp.tool() +async def book_table(date: str, party_size: int, ctx: Context) -> str: + """Book a table with confirmation""" + + # Schema must only contain primitive types (str, int, float, bool) + class ConfirmBooking(BaseModel): + confirm: bool = Field(description="Confirm booking?") + notes: str = Field(default="", description="Special requests") + + result = await ctx.elicit( + message=f"Confirm booking for {party_size} on {date}?", schema=ConfirmBooking + ) + + match result: + case AcceptedElicitation(data=data): + if data.confirm: + return f"Booked! Notes: {data.notes or 'None'}" + return "Booking cancelled" + case DeclinedElicitation(): + return "Booking declined" + case CancelledElicitation(): + return "Booking cancelled" +``` + +The `elicit()` method returns an `ElicitationResult` with: +- `action`: "accept", "decline", or "cancel" +- `data`: The validated response (only when accepted) +- `validation_error`: Any validation error message ### Authentication diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 8b819ad6d..948817140 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -22,6 +22,14 @@ async def __call__( ) -> types.CreateMessageResult | types.ErrorData: ... +class ElicitationFnT(Protocol): + async def __call__( + self, + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, + ) -> types.ElicitResult | types.ErrorData: ... + + class ListRootsFnT(Protocol): async def __call__( self, context: RequestContext["ClientSession", Any] @@ -58,6 +66,16 @@ async def _default_sampling_callback( ) +async def _default_elicitation_callback( + context: RequestContext["ClientSession", Any], + params: types.ElicitRequestParams, +) -> types.ElicitResult | types.ErrorData: + return types.ErrorData( + code=types.INVALID_REQUEST, + message="Elicitation not supported", + ) + + async def _default_list_roots_callback( context: RequestContext["ClientSession", Any], ) -> types.ListRootsResult | types.ErrorData: @@ -91,6 +109,7 @@ def __init__( write_stream: MemoryObjectSendStream[SessionMessage], read_timeout_seconds: timedelta | None = None, sampling_callback: SamplingFnT | None = None, + elicitation_callback: ElicitationFnT | None = None, list_roots_callback: ListRootsFnT | None = None, logging_callback: LoggingFnT | None = None, message_handler: MessageHandlerFnT | None = None, @@ -105,12 +124,16 @@ def __init__( ) self._client_info = client_info or DEFAULT_CLIENT_INFO self._sampling_callback = sampling_callback or _default_sampling_callback + self._elicitation_callback = elicitation_callback or _default_elicitation_callback self._list_roots_callback = list_roots_callback or _default_list_roots_callback self._logging_callback = logging_callback or _default_logging_callback self._message_handler = message_handler or _default_message_handler async def initialize(self) -> types.InitializeResult: sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None + elicitation = ( + types.ElicitationCapability() if self._elicitation_callback is not _default_elicitation_callback else None + ) roots = ( # TODO: Should this be based on whether we # _will_ send notifications, or only whether @@ -128,6 +151,7 @@ async def initialize(self) -> types.InitializeResult: protocolVersion=types.LATEST_PROTOCOL_VERSION, capabilities=types.ClientCapabilities( sampling=sampling, + elicitation=elicitation, experimental=None, roots=roots, ), @@ -362,6 +386,12 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques client_response = ClientResponse.validate_python(response) await responder.respond(client_response) + case types.ElicitRequest(params=params): + with responder: + response = await self._elicitation_callback(ctx, params) + client_response = ClientResponse.validate_python(response) + await responder.respond(client_response) + case types.ListRootsRequest(): with responder: response = await self._list_roots_callback(ctx) diff --git a/src/mcp/server/elicitation.py b/src/mcp/server/elicitation.py new file mode 100644 index 000000000..1e48738c8 --- /dev/null +++ b/src/mcp/server/elicitation.py @@ -0,0 +1,111 @@ +"""Elicitation utilities for MCP servers.""" + +from __future__ import annotations + +import types +from typing import Generic, Literal, TypeVar, Union, get_args, get_origin + +from pydantic import BaseModel +from pydantic.fields import FieldInfo + +from mcp.server.session import ServerSession +from mcp.types import RequestId + +ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel) + + +class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]): + """Result when user accepts the elicitation.""" + + action: Literal["accept"] = "accept" + data: ElicitSchemaModelT + + +class DeclinedElicitation(BaseModel): + """Result when user declines the elicitation.""" + + action: Literal["decline"] = "decline" + + +class CancelledElicitation(BaseModel): + """Result when user cancels the elicitation.""" + + action: Literal["cancel"] = "cancel" + + +ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation + + +# Primitive types allowed in elicitation schemas +_ELICITATION_PRIMITIVE_TYPES = (str, int, float, bool) + + +def _validate_elicitation_schema(schema: type[BaseModel]) -> None: + """Validate that a Pydantic model only contains primitive field types.""" + for field_name, field_info in schema.model_fields.items(): + if not _is_primitive_field(field_info): + raise TypeError( + f"Elicitation schema field '{field_name}' must be a primitive type " + f"{_ELICITATION_PRIMITIVE_TYPES} or Optional of these types. " + f"Complex types like lists, dicts, or nested models are not allowed." + ) + + +def _is_primitive_field(field_info: FieldInfo) -> bool: + """Check if a field is a primitive type allowed in elicitation schemas.""" + annotation = field_info.annotation + + # Handle None type + if annotation is types.NoneType: + return True + + # Handle basic primitive types + if annotation in _ELICITATION_PRIMITIVE_TYPES: + return True + + # Handle Union types + origin = get_origin(annotation) + if origin is Union or origin is types.UnionType: + args = get_args(annotation) + # All args must be primitive types or None + return all(arg is types.NoneType or arg in _ELICITATION_PRIMITIVE_TYPES for arg in args) + + return False + + +async def elicit_with_validation( + session: ServerSession, + message: str, + schema: type[ElicitSchemaModelT], + related_request_id: RequestId | None = None, +) -> ElicitationResult[ElicitSchemaModelT]: + """Elicit information from the client/user with schema validation. + + This method can be used to interactively ask for additional information from the + client within a tool's execution. The client might display the message to the + user and collect a response according to the provided schema. Or in case a + client is an agent, it might decide how to handle the elicitation -- either by asking + the user or automatically generating a response. + """ + # Validate that schema only contains primitive types and fail loudly if not + _validate_elicitation_schema(schema) + + json_schema = schema.model_json_schema() + + result = await session.elicit( + message=message, + requestedSchema=json_schema, + related_request_id=related_request_id, + ) + + if result.action == "accept" and result.content: + # Validate and parse the content using the schema + validated_data = schema.model_validate(result.content) + return AcceptedElicitation(data=validated_data) + elif result.action == "decline": + return DeclinedElicitation() + elif result.action == "cancel": + return CancelledElicitation() + else: + # This should never happen, but handle it just in case + raise ValueError(f"Unexpected elicitation action: {result.action}") diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index a62974bc9..a85c7117c 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -34,6 +34,7 @@ from mcp.server.auth.settings import ( AuthSettings, ) +from mcp.server.elicitation import ElicitationResult, ElicitSchemaModelT, elicit_with_validation from mcp.server.fastmcp.exceptions import ResourceError from mcp.server.fastmcp.prompts import Prompt, PromptManager from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager @@ -972,6 +973,37 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent assert self._fastmcp is not None, "Context is not available outside of a request" return await self._fastmcp.read_resource(uri) + async def elicit( + self, + message: str, + schema: type[ElicitSchemaModelT], + ) -> ElicitationResult[ElicitSchemaModelT]: + """Elicit information from the client/user. + + This method can be used to interactively ask for additional information from the + client within a tool's execution. The client might display the message to the + user and collect a response according to the provided schema. Or in case a + client is an agent, it might decide how to handle the elicitation -- either by asking + the user or automatically generating a response. + + Args: + schema: A Pydantic model class defining the expected response structure, according to the specification, + only primive types are allowed. + message: Optional message to present to the user. If not provided, will use + a default message based on the schema + + Returns: + An ElicitationResult containing the action taken and the data if accepted + + Note: + Check the result.action to determine if the user accepted, declined, or cancelled. + The result.data will only be populated if action is "accept" and validation succeeded. + """ + + return await elicit_with_validation( + session=self.request_context.session, message=message, schema=schema, related_request_id=self.request_id + ) + async def log( self, level: Literal["debug", "info", "warning", "error"], diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index e6611b0d4..5c696b136 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -121,6 +121,10 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool: if client_caps.sampling is None: return False + if capability.elicitation is not None: + if client_caps.elicitation is None: + return False + if capability.experimental is not None: if client_caps.experimental is None: return False @@ -251,6 +255,35 @@ async def list_roots(self) -> types.ListRootsResult: types.ListRootsResult, ) + async def elicit( + self, + message: str, + requestedSchema: types.ElicitRequestedSchema, + related_request_id: types.RequestId | None = None, + ) -> types.ElicitResult: + """Send an elicitation/create request. + + Args: + message: The message to present to the user + requestedSchema: Schema defining the expected response structure + + Returns: + The client's response + """ + return await self.send_request( + types.ServerRequest( + types.ElicitRequest( + method="elicitation/create", + params=types.ElicitRequestParams( + message=message, + requestedSchema=requestedSchema, + ), + ) + ), + types.ElicitResult, + metadata=ServerMessageMetadata(related_request_id=related_request_id), + ) + async def send_ping(self) -> types.EmptyResult: """Send a ping request.""" return await self.send_request( diff --git a/src/mcp/shared/memory.py b/src/mcp/shared/memory.py index f088d3f8b..c94e5e6ac 100644 --- a/src/mcp/shared/memory.py +++ b/src/mcp/shared/memory.py @@ -11,7 +11,14 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream import mcp.types as types -from mcp.client.session import ClientSession, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT +from mcp.client.session import ( + ClientSession, + ElicitationFnT, + ListRootsFnT, + LoggingFnT, + MessageHandlerFnT, + SamplingFnT, +) from mcp.server import Server from mcp.shared.message import SessionMessage @@ -53,6 +60,7 @@ async def create_connected_server_and_client_session( message_handler: MessageHandlerFnT | None = None, client_info: types.Implementation | None = None, raise_exceptions: bool = False, + elicitation_callback: ElicitationFnT | None = None, ) -> AsyncGenerator[ClientSession, None]: """Creates a ClientSession that is connected to a running MCP server.""" async with create_client_server_memory_streams() as ( @@ -83,6 +91,7 @@ async def create_connected_server_and_client_session( logging_callback=logging_callback, message_handler=message_handler, client_info=client_info, + elicitation_callback=elicitation_callback, ) as client_session: await client_session.initialize() yield client_session diff --git a/src/mcp/types.py b/src/mcp/types.py index be8d7326f..a678f101d 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -216,6 +216,12 @@ class SamplingCapability(BaseModel): model_config = ConfigDict(extra="allow") +class ElicitationCapability(BaseModel): + """Capability for elicitation operations.""" + + model_config = ConfigDict(extra="allow") + + class ClientCapabilities(BaseModel): """Capabilities a client may support.""" @@ -223,6 +229,8 @@ class ClientCapabilities(BaseModel): """Experimental, non-standard capabilities that the client supports.""" sampling: SamplingCapability | None = None """Present if the client supports sampling from an LLM.""" + elicitation: ElicitationCapability | None = None + """Present if the client supports elicitation from the user.""" roots: RootsCapability | None = None """Present if the client supports listing roots.""" model_config = ConfigDict(extra="allow") @@ -1186,11 +1194,49 @@ class ClientNotification( pass -class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult]): +# Type for elicitation schema - a JSON Schema dict +ElicitRequestedSchema: TypeAlias = dict[str, Any] +"""Schema for elicitation requests.""" + + +class ElicitRequestParams(RequestParams): + """Parameters for elicitation requests.""" + + message: str + requestedSchema: ElicitRequestedSchema + model_config = ConfigDict(extra="allow") + + +class ElicitRequest(Request[ElicitRequestParams, Literal["elicitation/create"]]): + """A request from the server to elicit information from the client.""" + + method: Literal["elicitation/create"] + params: ElicitRequestParams + + +class ElicitResult(Result): + """The client's response to an elicitation request.""" + + action: Literal["accept", "decline", "cancel"] + """ + The user action in response to the elicitation. + - "accept": User submitted the form/confirmed the action + - "decline": User explicitly declined the action + - "cancel": User dismissed without making an explicit choice + """ + + content: dict[str, str | int | float | bool | None] | None = None + """ + The submitted form data, only present when action is "accept". + Contains values matching the requested schema. + """ + + +class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult | ElicitResult]): pass -class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest]): +class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest | ElicitRequest]): pass diff --git a/tests/server/fastmcp/test_elicitation.py b/tests/server/fastmcp/test_elicitation.py new file mode 100644 index 000000000..20937d91d --- /dev/null +++ b/tests/server/fastmcp/test_elicitation.py @@ -0,0 +1,210 @@ +""" +Test the elicitation feature using stdio transport. +""" + +import pytest +from pydantic import BaseModel, Field + +from mcp.server.fastmcp import Context, FastMCP +from mcp.shared.memory import create_connected_server_and_client_session +from mcp.types import ElicitResult, TextContent + + +# Shared schema for basic tests +class AnswerSchema(BaseModel): + answer: str = Field(description="The user's answer to the question") + + +def create_ask_user_tool(mcp: FastMCP): + """Create a standard ask_user tool that handles all elicitation responses.""" + + @mcp.tool(description="A tool that uses elicitation") + async def ask_user(prompt: str, ctx: Context) -> str: + result = await ctx.elicit( + message=f"Tool wants to ask: {prompt}", + schema=AnswerSchema, + ) + + if result.action == "accept" and result.data: + return f"User answered: {result.data.answer}" + elif result.action == "decline": + return "User declined to answer" + else: + return "User cancelled" + + return ask_user + + +async def call_tool_and_assert( + mcp: FastMCP, + elicitation_callback, + tool_name: str, + args: dict, + expected_text: str | None = None, + text_contains: list[str] | None = None, +): + """Helper to create session, call tool, and assert result.""" + async with create_connected_server_and_client_session( + mcp._mcp_server, elicitation_callback=elicitation_callback + ) as client_session: + await client_session.initialize() + + result = await client_session.call_tool(tool_name, args) + assert len(result.content) == 1 + assert isinstance(result.content[0], TextContent) + + if expected_text is not None: + assert result.content[0].text == expected_text + elif text_contains is not None: + for substring in text_contains: + assert substring in result.content[0].text + + return result + + +@pytest.mark.anyio +async def test_stdio_elicitation(): + """Test the elicitation feature using stdio transport.""" + mcp = FastMCP(name="StdioElicitationServer") + create_ask_user_tool(mcp) + + # Create a custom handler for elicitation requests + async def elicitation_callback(context, params): + if params.message == "Tool wants to ask: What is your name?": + return ElicitResult(action="accept", content={"answer": "Test User"}) + else: + raise ValueError(f"Unexpected elicitation message: {params.message}") + + await call_tool_and_assert( + mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User answered: Test User" + ) + + +@pytest.mark.anyio +async def test_stdio_elicitation_decline(): + """Test elicitation with user declining.""" + mcp = FastMCP(name="StdioElicitationDeclineServer") + create_ask_user_tool(mcp) + + async def elicitation_callback(context, params): + return ElicitResult(action="decline") + + await call_tool_and_assert( + mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User declined to answer" + ) + + +@pytest.mark.anyio +async def test_elicitation_schema_validation(): + """Test that elicitation schemas must only contain primitive types.""" + mcp = FastMCP(name="ValidationTestServer") + + def create_validation_tool(name: str, schema_class: type[BaseModel]): + @mcp.tool(name=name, description=f"Tool testing {name}") + async def tool(ctx: Context) -> str: + try: + await ctx.elicit(message="This should fail validation", schema=schema_class) + return "Should not reach here" + except TypeError as e: + return f"Validation failed as expected: {str(e)}" + + return tool + + # Test cases for invalid schemas + class InvalidListSchema(BaseModel): + names: list[str] = Field(description="List of names") + + class NestedModel(BaseModel): + value: str + + class InvalidNestedSchema(BaseModel): + nested: NestedModel = Field(description="Nested model") + + create_validation_tool("invalid_list", InvalidListSchema) + create_validation_tool("nested_model", InvalidNestedSchema) + + # Dummy callback (won't be called due to validation failure) + async def elicitation_callback(context, params): + return ElicitResult(action="accept", content={}) + + async with create_connected_server_and_client_session( + mcp._mcp_server, elicitation_callback=elicitation_callback + ) as client_session: + await client_session.initialize() + + # Test both invalid schemas + for tool_name, field_name in [("invalid_list", "names"), ("nested_model", "nested")]: + result = await client_session.call_tool(tool_name, {}) + assert len(result.content) == 1 + assert isinstance(result.content[0], TextContent) + assert "Validation failed as expected" in result.content[0].text + assert field_name in result.content[0].text + + +@pytest.mark.anyio +async def test_elicitation_with_optional_fields(): + """Test that Optional fields work correctly in elicitation schemas.""" + mcp = FastMCP(name="OptionalFieldServer") + + class OptionalSchema(BaseModel): + required_name: str = Field(description="Your name (required)") + optional_age: int | None = Field(default=None, description="Your age (optional)") + optional_email: str | None = Field(default=None, description="Your email (optional)") + subscribe: bool | None = Field(default=False, description="Subscribe to newsletter?") + + @mcp.tool(description="Tool with optional fields") + async def optional_tool(ctx: Context) -> str: + result = await ctx.elicit(message="Please provide your information", schema=OptionalSchema) + + if result.action == "accept" and result.data: + info = [f"Name: {result.data.required_name}"] + if result.data.optional_age is not None: + info.append(f"Age: {result.data.optional_age}") + if result.data.optional_email is not None: + info.append(f"Email: {result.data.optional_email}") + info.append(f"Subscribe: {result.data.subscribe}") + return ", ".join(info) + else: + return f"User {result.action}" + + # Test cases with different field combinations + test_cases = [ + ( + # All fields provided + {"required_name": "John Doe", "optional_age": 30, "optional_email": "john@example.com", "subscribe": True}, + "Name: John Doe, Age: 30, Email: john@example.com, Subscribe: True", + ), + ( + # Only required fields + {"required_name": "Jane Smith"}, + "Name: Jane Smith, Subscribe: False", + ), + ] + + for content, expected in test_cases: + + async def callback(context, params): + return ElicitResult(action="accept", content=content) + + await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected) + + # Test invalid optional field + class InvalidOptionalSchema(BaseModel): + name: str = Field(description="Name") + optional_list: list[str] | None = Field(default=None, description="Invalid optional list") + + @mcp.tool(description="Tool with invalid optional field") + async def invalid_optional_tool(ctx: Context) -> str: + try: + await ctx.elicit(message="This should fail", schema=InvalidOptionalSchema) + return "Should not reach here" + except TypeError as e: + return f"Validation failed: {str(e)}" + + await call_tool_and_assert( + mcp, + lambda c, p: ElicitResult(action="accept", content={}), + "invalid_optional_tool", + {}, + text_contains=["Validation failed:", "optional_list"], + ) diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index 3caf994e0..3eb013946 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -14,14 +14,14 @@ import pytest import uvicorn -from pydantic import AnyUrl +from pydantic import AnyUrl, BaseModel, Field from starlette.applications import Starlette from starlette.requests import Request from mcp.client.session import ClientSession from mcp.client.sse import sse_client from mcp.client.streamable_http import streamablehttp_client -from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp import Context, FastMCP from mcp.server.fastmcp.resources import FunctionResource from mcp.shared.context import RequestContext from mcp.types import ( @@ -30,6 +30,7 @@ CompletionContext, CreateMessageRequestParams, CreateMessageResult, + ElicitResult, GetPromptResult, InitializeResult, LoggingMessageNotification, @@ -98,6 +99,20 @@ def make_fastmcp_app(): def echo(message: str) -> str: return f"Echo: {message}" + # Add a tool that uses elicitation + @mcp.tool(description="A tool that uses elicitation") + async def ask_user(prompt: str, ctx: Context) -> str: + class AnswerSchema(BaseModel): + answer: str = Field(description="The user's answer to the question") + + result = await ctx.elicit(message=f"Tool wants to ask: {prompt}", schema=AnswerSchema) + + if result.action == "accept" and result.data: + return f"User answered: {result.data.answer}" + else: + # Handle cancellation or decline + return f"User cancelled or declined: {result.action}" + # Create the SSE app app = mcp.sse_app() @@ -261,6 +276,49 @@ def echo_context(custom_request_id: str, ctx: Context[Any, Any, Request]) -> str context_data["path"] = request.url.path return json.dumps(context_data) + # Restaurant booking tool with elicitation + @mcp.tool(description="Book a table at a restaurant with elicitation") + async def book_restaurant( + date: str, + time: str, + party_size: int, + ctx: Context, + ) -> str: + """Book a table - uses elicitation if requested date is unavailable.""" + + class AlternativeDateSchema(BaseModel): + checkAlternative: bool = Field(description="Would you like to try another date?") + alternativeDate: str = Field( + default="2024-12-26", + description="What date would you prefer? (YYYY-MM-DD)", + ) + + # For testing: assume dates starting with "2024-12-25" are unavailable + if date.startswith("2024-12-25"): + # Use elicitation to ask about alternatives + result = await ctx.elicit( + message=( + f"No tables available for {party_size} people on {date} " + f"at {time}. Would you like to check another date?" + ), + schema=AlternativeDateSchema, + ) + + if result.action == "accept" and result.data: + if result.data.checkAlternative: + alt_date = result.data.alternativeDate + return f"✅ Booked table for {party_size} on {alt_date} at {time}" + else: + return "❌ No booking made" + elif result.action in ("decline", "cancel"): + return "❌ Booking cancelled" + else: + # Handle case where action is "accept" but data is None + return "❌ No booking data received" + else: + # Available - book directly + return f"✅ Booked table for {party_size} on {date} at {time}" + return mcp @@ -652,6 +710,22 @@ async def handle_generic_notification(self, message) -> None: await self.handle_tool_list_changed(message.root.params) +async def create_test_elicitation_callback(context, params): + """Shared elicitation callback for tests. + + Handles elicitation requests for restaurant booking tests. + """ + # For restaurant booking test + if "No tables available" in params.message: + return ElicitResult( + action="accept", + content={"checkAlternative": True, "alternativeDate": "2024-12-26"}, + ) + else: + # Default response + return ElicitResult(action="decline") + + async def call_all_mcp_features(session: ClientSession, collector: NotificationCollector) -> None: """ Test all MCP features using the provided session. @@ -747,6 +821,21 @@ async def progress_callback(progress: float, total: float | None, message: str | assert "info" in log_levels assert "warning" in log_levels + # 5. Test elicitation tool + # Test restaurant booking with unavailable date (triggers elicitation) + booking_result = await session.call_tool( + "book_restaurant", + { + "date": "2024-12-25", # Unavailable date to trigger elicitation + "time": "19:00", + "party_size": 4, + }, + ) + assert len(booking_result.content) == 1 + assert isinstance(booking_result.content[0], TextContent) + # Should have booked the alternative date from elicitation callback + assert "✅ Booked table for 4 on 2024-12-26" in booking_result.content[0].text + # Test resources # 1. Static resource resources = await session.list_resources() @@ -887,8 +976,6 @@ async def test_fastmcp_all_features_sse(everything_server: None, everything_serv # Create notification collector collector = NotificationCollector() - # Create a sampling callback that simulates an LLM - # Connect to the server with callbacks async with sse_client(everything_server_url + "/sse") as streams: # Set up message handler to capture notifications @@ -901,6 +988,7 @@ async def message_handler(message): async with ClientSession( *streams, sampling_callback=sampling_callback, + elicitation_callback=create_test_elicitation_callback, message_handler=message_handler, ) as session: # Run the common test suite @@ -933,7 +1021,37 @@ async def message_handler(message): read_stream, write_stream, sampling_callback=sampling_callback, + elicitation_callback=create_test_elicitation_callback, message_handler=message_handler, ) as session: # Run the common test suite with HTTP-specific test suffix await call_all_mcp_features(session, collector) + + +@pytest.mark.anyio +async def test_elicitation_feature(server: None, server_url: str) -> None: + """Test the elicitation feature.""" + + # Create a custom handler for elicitation requests + async def elicitation_callback(context, params): + # Verify the elicitation parameters + if params.message == "Tool wants to ask: What is your name?": + return ElicitResult(content={"answer": "Test User"}, action="accept") + else: + raise ValueError("Unexpected elicitation message") + + # Connect to the server with our custom elicitation handler + async with sse_client(server_url + "/sse") as streams: + async with ClientSession(*streams, elicitation_callback=elicitation_callback) as session: + # First initialize the session + result = await session.initialize() + assert isinstance(result, InitializeResult) + assert result.serverInfo.name == "NoAuthServer" + + # Call the tool that uses elicitation + tool_result = await session.call_tool("ask_user", {"prompt": "What is your name?"}) + # Verify the result + assert len(tool_result.content) == 1 + assert isinstance(tool_result.content[0], TextContent) + # # The test should only succeed with the successful elicitation response + assert tool_result.content[0].text == "User answered: Test User" From 5429ac6153de20cc198e5edf3d937197d055548b Mon Sep 17 00:00:00 2001 From: Inna Harper Date: Tue, 17 Jun 2025 11:54:59 +0100 Subject: [PATCH 10/14] Update _meta usage guidance in types (#971) --- src/mcp/types.py | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/src/mcp/types.py b/src/mcp/types.py index a678f101d..d56a9bce8 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -69,7 +69,8 @@ class Meta(BaseModel): meta: Meta | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ @@ -106,7 +107,8 @@ class Result(BaseModel): meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -401,7 +403,8 @@ class Resource(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -426,7 +429,8 @@ class ResourceTemplate(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -476,7 +480,8 @@ class ResourceContents(BaseModel): """The MIME type of this resource, if known.""" meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -607,7 +612,8 @@ class Prompt(BaseModel): """A list of arguments to use for templating the prompt.""" meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -644,7 +650,8 @@ class TextContent(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -663,7 +670,8 @@ class ImageContent(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -682,7 +690,8 @@ class AudioContent(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -708,7 +717,8 @@ class EmbeddedResource(BaseModel): annotations: Annotations | None = None meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -811,7 +821,8 @@ class Tool(BaseModel): """Optional additional tool information.""" meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") @@ -1117,7 +1128,8 @@ class Root(BaseModel): """ meta: dict[str, Any] | None = Field(alias="_meta", default=None) """ - Reserved by MCP for protocol-level metadata; implementations must not make assumptions about its contents. + See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields) + for notes on _meta usage. """ model_config = ConfigDict(extra="allow") From 029f434109d6a28e8fc97f16dfa1d5452be7b23c Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Tue, 17 Jun 2025 03:55:29 -0700 Subject: [PATCH 11/14] Allow longer duration in test_188_concurrency (#969) --- tests/issues/test_188_concurrency.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/issues/test_188_concurrency.py b/tests/issues/test_188_concurrency.py index d0a86885f..9ccffefa9 100644 --- a/tests/issues/test_188_concurrency.py +++ b/tests/issues/test_188_concurrency.py @@ -35,7 +35,7 @@ async def slow_resource(): end_time = anyio.current_time() duration = end_time - start_time - assert duration < 6 * _sleep_time_seconds + assert duration < 10 * _sleep_time_seconds print(duration) From a2f8766a8a6bf8d6e6745b6c6f48d799dffe9583 Mon Sep 17 00:00:00 2001 From: Felix Weinberger <3823880+felixweinberger@users.noreply.github.com> Date: Tue, 17 Jun 2025 10:53:14 -0400 Subject: [PATCH 12/14] Add title to tools, resources, prompts (#972) Co-authored-by: Claude --- README.md | 48 +++- .../simple-chatbot/mcp_simple_chatbot/main.py | 22 +- .../simple-prompt/mcp_simple_prompt/server.py | 1 + .../mcp_simple_resource/server.py | 18 +- .../simple-tool/mcp_simple_tool/server.py | 1 + src/mcp/server/fastmcp/prompts/base.py | 3 + src/mcp/server/fastmcp/resources/base.py | 1 + .../fastmcp/resources/resource_manager.py | 2 + src/mcp/server/fastmcp/resources/templates.py | 4 + src/mcp/server/fastmcp/resources/types.py | 2 + src/mcp/server/fastmcp/server.py | 23 +- src/mcp/server/fastmcp/tools/base.py | 3 + src/mcp/server/fastmcp/tools/tool_manager.py | 3 +- src/mcp/shared/metadata_utils.py | 45 ++++ src/mcp/types.py | 36 +-- tests/server/fastmcp/test_integration.py | 99 +++++++- tests/server/fastmcp/test_title.py | 215 ++++++++++++++++++ tests/server/fastmcp/test_tool_manager.py | 1 + 18 files changed, 481 insertions(+), 46 deletions(-) create mode 100644 src/mcp/shared/metadata_utils.py create mode 100644 tests/server/fastmcp/test_title.py diff --git a/README.md b/README.md index 1bdae1d20..d8a2db2b6 100644 --- a/README.md +++ b/README.md @@ -212,13 +212,13 @@ from mcp.server.fastmcp import FastMCP mcp = FastMCP("My App") -@mcp.resource("config://app") +@mcp.resource("config://app", title="Application Configuration") def get_config() -> str: """Static configuration data""" return "App configuration here" -@mcp.resource("users://{user_id}/profile") +@mcp.resource("users://{user_id}/profile", title="User Profile") def get_user_profile(user_id: str) -> str: """Dynamic user data""" return f"Profile data for user {user_id}" @@ -235,13 +235,13 @@ from mcp.server.fastmcp import FastMCP mcp = FastMCP("My App") -@mcp.tool() +@mcp.tool(title="BMI Calculator") def calculate_bmi(weight_kg: float, height_m: float) -> float: """Calculate BMI given weight in kg and height in meters""" return weight_kg / (height_m**2) -@mcp.tool() +@mcp.tool(title="Weather Fetcher") async def fetch_weather(city: str) -> str: """Fetch current weather for a city""" async with httpx.AsyncClient() as client: @@ -260,12 +260,12 @@ from mcp.server.fastmcp.prompts import base mcp = FastMCP("My App") -@mcp.prompt() +@mcp.prompt(title="Code Review") def review_code(code: str) -> str: return f"Please review this code:\n\n{code}" -@mcp.prompt() +@mcp.prompt(title="Debug Assistant") def debug_error(error: str) -> list[base.Message]: return [ base.UserMessage("I'm seeing this error:"), @@ -918,6 +918,42 @@ async def main(): tool_result = await session.call_tool("echo", {"message": "hello"}) ``` +### Client Display Utilities + +When building MCP clients, the SDK provides utilities to help display human-readable names for tools, resources, and prompts: + +```python +from mcp.shared.metadata_utils import get_display_name +from mcp.client.session import ClientSession + + +async def display_tools(session: ClientSession): + """Display available tools with human-readable names""" + tools_response = await session.list_tools() + + for tool in tools_response.tools: + # get_display_name() returns the title if available, otherwise the name + display_name = get_display_name(tool) + print(f"Tool: {display_name}") + if tool.description: + print(f" {tool.description}") + + +async def display_resources(session: ClientSession): + """Display available resources with human-readable names""" + resources_response = await session.list_resources() + + for resource in resources_response.resources: + display_name = get_display_name(resource) + print(f"Resource: {display_name} ({resource.uri})") +``` + +The `get_display_name()` function implements the proper precedence rules for displaying names: +- For tools: `title` > `annotations.title` > `name` +- For other objects: `title` > `name` + +This ensures your client UI shows the most user-friendly names that servers provide. + ### OAuth Authentication for Clients The SDK includes [authorization support](https://modelcontextprotocol.io/specification/2025-03-26/basic/authorization) for connecting to protected MCP servers: diff --git a/examples/clients/simple-chatbot/mcp_simple_chatbot/main.py b/examples/clients/simple-chatbot/mcp_simple_chatbot/main.py index 180d709aa..b97b85080 100644 --- a/examples/clients/simple-chatbot/mcp_simple_chatbot/main.py +++ b/examples/clients/simple-chatbot/mcp_simple_chatbot/main.py @@ -123,7 +123,7 @@ async def list_tools(self) -> list[Any]: for item in tools_response: if isinstance(item, tuple) and item[0] == "tools": tools.extend( - Tool(tool.name, tool.description, tool.inputSchema) + Tool(tool.name, tool.description, tool.inputSchema, tool.title) for tool in item[1] ) @@ -189,9 +189,14 @@ class Tool: """Represents a tool with its properties and formatting.""" def __init__( - self, name: str, description: str, input_schema: dict[str, Any] + self, + name: str, + description: str, + input_schema: dict[str, Any], + title: str | None = None, ) -> None: self.name: str = name + self.title: str | None = title self.description: str = description self.input_schema: dict[str, Any] = input_schema @@ -211,13 +216,20 @@ def format_for_llm(self) -> str: arg_desc += " (required)" args_desc.append(arg_desc) - return f""" -Tool: {self.name} -Description: {self.description} + # Build the formatted output with title as a separate field + output = f"Tool: {self.name}\n" + + # Add human-readable title if available + if self.title: + output += f"User-readable title: {self.title}\n" + + output += f"""Description: {self.description} Arguments: {chr(10).join(args_desc)} """ + return output + class LLMClient: """Manages communication with the LLM provider.""" diff --git a/examples/servers/simple-prompt/mcp_simple_prompt/server.py b/examples/servers/simple-prompt/mcp_simple_prompt/server.py index eca0bbcf3..b562cc932 100644 --- a/examples/servers/simple-prompt/mcp_simple_prompt/server.py +++ b/examples/servers/simple-prompt/mcp_simple_prompt/server.py @@ -53,6 +53,7 @@ async def list_prompts() -> list[types.Prompt]: return [ types.Prompt( name="simple", + title="Simple Assistant Prompt", description="A simple prompt that can take optional context and topic " "arguments", arguments=[ diff --git a/examples/servers/simple-resource/mcp_simple_resource/server.py b/examples/servers/simple-resource/mcp_simple_resource/server.py index 545b415db..cef29b851 100644 --- a/examples/servers/simple-resource/mcp_simple_resource/server.py +++ b/examples/servers/simple-resource/mcp_simple_resource/server.py @@ -5,9 +5,18 @@ from pydantic import AnyUrl, FileUrl SAMPLE_RESOURCES = { - "greeting": "Hello! This is a sample text resource.", - "help": "This server provides a few sample text resources for testing.", - "about": "This is the simple-resource MCP server implementation.", + "greeting": { + "content": "Hello! This is a sample text resource.", + "title": "Welcome Message", + }, + "help": { + "content": "This server provides a few sample text resources for testing.", + "title": "Help Documentation", + }, + "about": { + "content": "This is the simple-resource MCP server implementation.", + "title": "About This Server", + }, } @@ -28,6 +37,7 @@ async def list_resources() -> list[types.Resource]: types.Resource( uri=FileUrl(f"file:///{name}.txt"), name=name, + title=SAMPLE_RESOURCES[name]["title"], description=f"A sample text resource named {name}", mimeType="text/plain", ) @@ -43,7 +53,7 @@ async def read_resource(uri: AnyUrl) -> str | bytes: if name not in SAMPLE_RESOURCES: raise ValueError(f"Unknown resource: {uri}") - return SAMPLE_RESOURCES[name] + return SAMPLE_RESOURCES[name]["content"] if transport == "sse": from mcp.server.sse import SseServerTransport diff --git a/examples/servers/simple-tool/mcp_simple_tool/server.py b/examples/servers/simple-tool/mcp_simple_tool/server.py index 29c741c2c..a6a3f0bd2 100644 --- a/examples/servers/simple-tool/mcp_simple_tool/server.py +++ b/examples/servers/simple-tool/mcp_simple_tool/server.py @@ -41,6 +41,7 @@ async def list_tools() -> list[types.Tool]: return [ types.Tool( name="fetch", + title="Website Fetcher", description="Fetches a website and returns its content", inputSchema={ "type": "object", diff --git a/src/mcp/server/fastmcp/prompts/base.py b/src/mcp/server/fastmcp/prompts/base.py index a269d6ab4..e4751feb1 100644 --- a/src/mcp/server/fastmcp/prompts/base.py +++ b/src/mcp/server/fastmcp/prompts/base.py @@ -58,6 +58,7 @@ class Prompt(BaseModel): """A prompt template that can be rendered with parameters.""" name: str = Field(description="Name of the prompt") + title: str | None = Field(None, description="Human-readable title of the prompt") description: str | None = Field(None, description="Description of what the prompt does") arguments: list[PromptArgument] | None = Field(None, description="Arguments that can be passed to the prompt") fn: Callable[..., PromptResult | Awaitable[PromptResult]] = Field(exclude=True) @@ -67,6 +68,7 @@ def from_function( cls, fn: Callable[..., PromptResult | Awaitable[PromptResult]], name: str | None = None, + title: str | None = None, description: str | None = None, ) -> "Prompt": """Create a Prompt from a function. @@ -103,6 +105,7 @@ def from_function( return cls( name=func_name, + title=title, description=description or fn.__doc__ or "", arguments=arguments, fn=fn, diff --git a/src/mcp/server/fastmcp/resources/base.py b/src/mcp/server/fastmcp/resources/base.py index 24bf7d531..f57631cc1 100644 --- a/src/mcp/server/fastmcp/resources/base.py +++ b/src/mcp/server/fastmcp/resources/base.py @@ -21,6 +21,7 @@ class Resource(BaseModel, abc.ABC): uri: Annotated[AnyUrl, UrlConstraints(host_required=False)] = Field(default=..., description="URI of the resource") name: str | None = Field(description="Name of the resource", default=None) + title: str | None = Field(description="Human-readable title of the resource", default=None) description: str | None = Field(description="Description of the resource", default=None) mime_type: str = Field( default="text/plain", diff --git a/src/mcp/server/fastmcp/resources/resource_manager.py b/src/mcp/server/fastmcp/resources/resource_manager.py index d27e6ac12..35e4ec04d 100644 --- a/src/mcp/server/fastmcp/resources/resource_manager.py +++ b/src/mcp/server/fastmcp/resources/resource_manager.py @@ -51,6 +51,7 @@ def add_template( fn: Callable[..., Any], uri_template: str, name: str | None = None, + title: str | None = None, description: str | None = None, mime_type: str | None = None, ) -> ResourceTemplate: @@ -59,6 +60,7 @@ def add_template( fn, uri_template=uri_template, name=name, + title=title, description=description, mime_type=mime_type, ) diff --git a/src/mcp/server/fastmcp/resources/templates.py b/src/mcp/server/fastmcp/resources/templates.py index 19d7dedd1..b1c7b2711 100644 --- a/src/mcp/server/fastmcp/resources/templates.py +++ b/src/mcp/server/fastmcp/resources/templates.py @@ -17,6 +17,7 @@ class ResourceTemplate(BaseModel): uri_template: str = Field(description="URI template with parameters (e.g. weather://{city}/current)") name: str = Field(description="Name of the resource") + title: str | None = Field(description="Human-readable title of the resource", default=None) description: str | None = Field(description="Description of what the resource does") mime_type: str = Field(default="text/plain", description="MIME type of the resource content") fn: Callable[..., Any] = Field(exclude=True) @@ -28,6 +29,7 @@ def from_function( fn: Callable[..., Any], uri_template: str, name: str | None = None, + title: str | None = None, description: str | None = None, mime_type: str | None = None, ) -> ResourceTemplate: @@ -45,6 +47,7 @@ def from_function( return cls( uri_template=uri_template, name=func_name, + title=title, description=description or fn.__doc__ or "", mime_type=mime_type or "text/plain", fn=fn, @@ -71,6 +74,7 @@ async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: return FunctionResource( uri=uri, # type: ignore name=self.name, + title=self.title, description=self.description, mime_type=self.mime_type, fn=lambda: result, # Capture result in closure diff --git a/src/mcp/server/fastmcp/resources/types.py b/src/mcp/server/fastmcp/resources/types.py index 63377f423..9c980dff1 100644 --- a/src/mcp/server/fastmcp/resources/types.py +++ b/src/mcp/server/fastmcp/resources/types.py @@ -72,6 +72,7 @@ def from_function( fn: Callable[..., Any], uri: str, name: str | None = None, + title: str | None = None, description: str | None = None, mime_type: str | None = None, ) -> "FunctionResource": @@ -86,6 +87,7 @@ def from_function( return cls( uri=AnyUrl(uri), name=func_name, + title=title, description=description or fn.__doc__ or "", mime_type=mime_type or "text/plain", fn=fn, diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index a85c7117c..3adc465b8 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -237,6 +237,7 @@ async def list_tools(self) -> list[MCPTool]: return [ MCPTool( name=info.name, + title=info.title, description=info.description, inputSchema=info.parameters, annotations=info.annotations, @@ -270,6 +271,7 @@ async def list_resources(self) -> list[MCPResource]: MCPResource( uri=resource.uri, name=resource.name or "", + title=resource.title, description=resource.description, mimeType=resource.mime_type, ) @@ -282,6 +284,7 @@ async def list_resource_templates(self) -> list[MCPResourceTemplate]: MCPResourceTemplate( uriTemplate=template.uri_template, name=template.name, + title=template.title, description=template.description, ) for template in templates @@ -305,6 +308,7 @@ def add_tool( self, fn: AnyFunction, name: str | None = None, + title: str | None = None, description: str | None = None, annotations: ToolAnnotations | None = None, ) -> None: @@ -316,14 +320,16 @@ def add_tool( Args: fn: The function to register as a tool name: Optional name for the tool (defaults to function name) + title: Optional human-readable title for the tool description: Optional description of what the tool does annotations: Optional ToolAnnotations providing additional tool information """ - self._tool_manager.add_tool(fn, name=name, description=description, annotations=annotations) + self._tool_manager.add_tool(fn, name=name, title=title, description=description, annotations=annotations) def tool( self, name: str | None = None, + title: str | None = None, description: str | None = None, annotations: ToolAnnotations | None = None, ) -> Callable[[AnyFunction], AnyFunction]: @@ -335,6 +341,7 @@ def tool( Args: name: Optional name for the tool (defaults to function name) + title: Optional human-readable title for the tool description: Optional description of what the tool does annotations: Optional ToolAnnotations providing additional tool information @@ -360,7 +367,7 @@ async def async_tool(x: int, context: Context) -> str: ) def decorator(fn: AnyFunction) -> AnyFunction: - self.add_tool(fn, name=name, description=description, annotations=annotations) + self.add_tool(fn, name=name, title=title, description=description, annotations=annotations) return fn return decorator @@ -396,6 +403,7 @@ def resource( uri: str, *, name: str | None = None, + title: str | None = None, description: str | None = None, mime_type: str | None = None, ) -> Callable[[AnyFunction], AnyFunction]: @@ -413,6 +421,7 @@ def resource( Args: uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}") name: Optional name for the resource + title: Optional human-readable title for the resource description: Optional description of the resource mime_type: Optional MIME type for the resource @@ -462,6 +471,7 @@ def decorator(fn: AnyFunction) -> AnyFunction: fn=fn, uri_template=uri, name=name, + title=title, description=description, mime_type=mime_type, ) @@ -471,6 +481,7 @@ def decorator(fn: AnyFunction) -> AnyFunction: fn=fn, uri=uri, name=name, + title=title, description=description, mime_type=mime_type, ) @@ -487,11 +498,14 @@ def add_prompt(self, prompt: Prompt) -> None: """ self._prompt_manager.add_prompt(prompt) - def prompt(self, name: str | None = None, description: str | None = None) -> Callable[[AnyFunction], AnyFunction]: + def prompt( + self, name: str | None = None, title: str | None = None, description: str | None = None + ) -> Callable[[AnyFunction], AnyFunction]: """Decorator to register a prompt. Args: name: Optional name for the prompt (defaults to function name) + title: Optional human-readable title for the prompt description: Optional description of what the prompt does Example: @@ -529,7 +543,7 @@ async def analyze_file(path: str) -> list[Message]: ) def decorator(func: AnyFunction) -> AnyFunction: - prompt = Prompt.from_function(func, name=name, description=description) + prompt = Prompt.from_function(func, name=name, title=title, description=description) self.add_prompt(prompt) return func @@ -831,6 +845,7 @@ async def list_prompts(self) -> list[MCPPrompt]: return [ MCPPrompt( name=prompt.name, + title=prompt.title, description=prompt.description, arguments=[ MCPPromptArgument( diff --git a/src/mcp/server/fastmcp/tools/base.py b/src/mcp/server/fastmcp/tools/base.py index b25597215..2f7c48e8b 100644 --- a/src/mcp/server/fastmcp/tools/base.py +++ b/src/mcp/server/fastmcp/tools/base.py @@ -22,6 +22,7 @@ class Tool(BaseModel): fn: Callable[..., Any] = Field(exclude=True) name: str = Field(description="Name of the tool") + title: str | None = Field(None, description="Human-readable title of the tool") description: str = Field(description="Description of what the tool does") parameters: dict[str, Any] = Field(description="JSON schema for tool parameters") fn_metadata: FuncMetadata = Field( @@ -36,6 +37,7 @@ def from_function( cls, fn: Callable[..., Any], name: str | None = None, + title: str | None = None, description: str | None = None, context_kwarg: str | None = None, annotations: ToolAnnotations | None = None, @@ -69,6 +71,7 @@ def from_function( return cls( fn=fn, name=func_name, + title=title, description=func_doc, parameters=parameters, fn_metadata=func_arg_metadata, diff --git a/src/mcp/server/fastmcp/tools/tool_manager.py b/src/mcp/server/fastmcp/tools/tool_manager.py index 1cd301299..b9ca1655d 100644 --- a/src/mcp/server/fastmcp/tools/tool_manager.py +++ b/src/mcp/server/fastmcp/tools/tool_manager.py @@ -46,11 +46,12 @@ def add_tool( self, fn: Callable[..., Any], name: str | None = None, + title: str | None = None, description: str | None = None, annotations: ToolAnnotations | None = None, ) -> Tool: """Add a tool to the server.""" - tool = Tool.from_function(fn, name=name, description=description, annotations=annotations) + tool = Tool.from_function(fn, name=name, title=title, description=description, annotations=annotations) existing = self._tools.get(tool.name) if existing: if self.warn_on_duplicate_tools: diff --git a/src/mcp/shared/metadata_utils.py b/src/mcp/shared/metadata_utils.py new file mode 100644 index 000000000..e3f49daf4 --- /dev/null +++ b/src/mcp/shared/metadata_utils.py @@ -0,0 +1,45 @@ +"""Utility functions for working with metadata in MCP types. + +These utilities are primarily intended for client-side usage to properly display +human-readable names in user interfaces in a spec compliant way. +""" + +from mcp.types import Implementation, Prompt, Resource, ResourceTemplate, Tool + + +def get_display_name(obj: Tool | Resource | Prompt | ResourceTemplate | Implementation) -> str: + """ + Get the display name for an MCP object with proper precedence. + + This is a client-side utility function designed to help MCP clients display + human-readable names in their user interfaces. When servers provide a 'title' + field, it should be preferred over the programmatic 'name' field for display. + + For tools: title > annotations.title > name + For other objects: title > name + + Example: + # In a client displaying available tools + tools = await session.list_tools() + for tool in tools.tools: + display_name = get_display_name(tool) + print(f"Available tool: {display_name}") + + Args: + obj: An MCP object with name and optional title fields + + Returns: + The display name to use for UI presentation + """ + if isinstance(obj, Tool): + # Tools have special precedence: title > annotations.title > name + if hasattr(obj, "title") and obj.title is not None: + return obj.title + if obj.annotations and hasattr(obj.annotations, "title") and obj.annotations.title is not None: + return obj.annotations.title + return obj.name + else: + # All other objects: title > name + if hasattr(obj, "title") and obj.title is not None: + return obj.title + return obj.name diff --git a/src/mcp/types.py b/src/mcp/types.py index d56a9bce8..f53813412 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -196,10 +196,26 @@ class EmptyResult(Result): """A response that indicates success but carries no data.""" -class Implementation(BaseModel): - """Describes the name and version of an MCP implementation.""" +class BaseMetadata(BaseModel): + """Base class for entities with name and optional title fields.""" name: str + """The programmatic name of the entity.""" + + title: str | None = None + """ + Intended for UI and end-user contexts — optimized to be human-readable and easily understood, + even by those unfamiliar with domain-specific terminology. + + If not provided, the name should be used for display (except for Tool, + where `annotations.title` should be given precedence over using `name`, + if present). + """ + + +class Implementation(BaseMetadata): + """Describes the name and version of an MCP implementation.""" + version: str model_config = ConfigDict(extra="allow") @@ -382,13 +398,11 @@ class Annotations(BaseModel): model_config = ConfigDict(extra="allow") -class Resource(BaseModel): +class Resource(BaseMetadata): """A known resource that the server is capable of reading.""" uri: Annotated[AnyUrl, UrlConstraints(host_required=False)] """The URI of this resource.""" - name: str - """A human-readable name for this resource.""" description: str | None = None """A description of what this resource represents.""" mimeType: str | None = None @@ -409,7 +423,7 @@ class Resource(BaseModel): model_config = ConfigDict(extra="allow") -class ResourceTemplate(BaseModel): +class ResourceTemplate(BaseMetadata): """A template description for resources available on the server.""" uriTemplate: str @@ -417,8 +431,6 @@ class ResourceTemplate(BaseModel): A URI template (according to RFC 6570) that can be used to construct resource URIs. """ - name: str - """A human-readable name for the type of resource this template refers to.""" description: str | None = None """A human-readable description of what this template is for.""" mimeType: str | None = None @@ -601,11 +613,9 @@ class PromptArgument(BaseModel): model_config = ConfigDict(extra="allow") -class Prompt(BaseModel): +class Prompt(BaseMetadata): """A prompt or prompt template that the server offers.""" - name: str - """The name of the prompt or prompt template.""" description: str | None = None """An optional description of what this prompt provides.""" arguments: list[PromptArgument] | None = None @@ -808,11 +818,9 @@ class ToolAnnotations(BaseModel): model_config = ConfigDict(extra="allow") -class Tool(BaseModel): +class Tool(BaseMetadata): """Definition for a tool the client can call.""" - name: str - """The name of the tool.""" description: str | None = None """A human-readable description of the tool.""" inputSchema: dict[str, Any] diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index 3eb013946..7224a8e48 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -126,7 +126,7 @@ def make_everything_fastmcp() -> FastMCP: mcp = FastMCP(name="EverythingServer") # Tool with context for logging and progress - @mcp.tool(description="A tool that demonstrates logging and progress") + @mcp.tool(description="A tool that demonstrates logging and progress", title="Progress Tool") async def tool_with_progress(message: str, ctx: Context, steps: int = 3) -> str: await ctx.info(f"Starting processing of '{message}' with {steps} steps") @@ -143,12 +143,12 @@ async def tool_with_progress(message: str, ctx: Context, steps: int = 3) -> str: return f"Processed '{message}' in {steps} steps" # Simple tool for basic functionality - @mcp.tool(description="A simple echo tool") + @mcp.tool(description="A simple echo tool", title="Echo Tool") def echo(message: str) -> str: return f"Echo: {message}" # Tool with sampling capability - @mcp.tool(description="A tool that uses sampling to generate content") + @mcp.tool(description="A tool that uses sampling to generate content", title="Sampling Tool") async def sampling_tool(prompt: str, ctx: Context) -> str: await ctx.info(f"Requesting sampling for prompt: {prompt}") @@ -167,7 +167,7 @@ async def sampling_tool(prompt: str, ctx: Context) -> str: return f"Sampling result: {str(result.content)[:100]}..." # Tool that sends notifications and logging - @mcp.tool(description="A tool that demonstrates notifications and logging") + @mcp.tool(description="A tool that demonstrates notifications and logging", title="Notification Tool") async def notification_tool(message: str, ctx: Context) -> str: # Send different log levels await ctx.debug("Debug: Starting notification tool") @@ -188,35 +188,36 @@ def get_static_info() -> str: static_resource = FunctionResource( uri=AnyUrl("resource://static/info"), name="Static Info", + title="Static Information", description="Static information resource", fn=get_static_info, ) mcp.add_resource(static_resource) # Resource - dynamic function - @mcp.resource("resource://dynamic/{category}") + @mcp.resource("resource://dynamic/{category}", title="Dynamic Resource") def dynamic_resource(category: str) -> str: return f"Dynamic resource content for category: {category}" # Resource template - @mcp.resource("resource://template/{id}/data") + @mcp.resource("resource://template/{id}/data", title="Template Resource") def template_resource(id: str) -> str: return f"Template resource data for ID: {id}" # Prompt - simple - @mcp.prompt(description="A simple prompt") + @mcp.prompt(description="A simple prompt", title="Simple Prompt") def simple_prompt(topic: str) -> str: return f"Tell me about {topic}" # Prompt - complex with multiple messages - @mcp.prompt(description="Complex prompt with context") + @mcp.prompt(description="Complex prompt with context", title="Complex Prompt") def complex_prompt(user_query: str, context: str = "general") -> str: # For simplicity, return a single string that incorporates the context # Since FastMCP doesn't support system messages in the same way return f"Context: {context}. Query: {user_query}" # Resource template with completion support - @mcp.resource("github://repos/{owner}/{repo}") + @mcp.resource("github://repos/{owner}/{repo}", title="GitHub Repository") def github_repo_resource(owner: str, repo: str) -> str: return f"Repository: {owner}/{repo}" @@ -250,7 +251,7 @@ async def handle_completion( return Completion(values=[], total=0, hasMore=False) # Tool that echoes request headers from context - @mcp.tool(description="Echo request headers from context") + @mcp.tool(description="Echo request headers from context", title="Echo Headers") def echo_headers(ctx: Context[Any, Any, Request]) -> str: """Returns the request headers as JSON.""" headers_info = {} @@ -260,7 +261,7 @@ def echo_headers(ctx: Context[Any, Any, Request]) -> str: return json.dumps(headers_info) # Tool that returns full request context - @mcp.tool(description="Echo request context with custom data") + @mcp.tool(description="Echo request context with custom data", title="Echo Context") def echo_context(custom_request_id: str, ctx: Context[Any, Any, Request]) -> str: """Returns request context including headers and custom data.""" context_data = { @@ -277,7 +278,7 @@ def echo_context(custom_request_id: str, ctx: Context[Any, Any, Request]) -> str return json.dumps(context_data) # Restaurant booking tool with elicitation - @mcp.tool(description="Book a table at a restaurant with elicitation") + @mcp.tool(description="Book a table at a restaurant with elicitation", title="Restaurant Booking") async def book_restaurant( date: str, time: str, @@ -1055,3 +1056,77 @@ async def elicitation_callback(context, params): assert isinstance(tool_result.content[0], TextContent) # # The test should only succeed with the successful elicitation response assert tool_result.content[0].text == "User answered: Test User" + + +@pytest.mark.anyio +async def test_title_precedence(everything_server: None, everything_server_url: str) -> None: + """Test that titles are properly returned for tools, resources, and prompts.""" + from mcp.shared.metadata_utils import get_display_name + + async with sse_client(everything_server_url + "/sse") as streams: + async with ClientSession(*streams) as session: + # Initialize the session + result = await session.initialize() + assert isinstance(result, InitializeResult) + + # Test tools have titles + tools_result = await session.list_tools() + assert tools_result.tools + + # Check specific tools have titles + tool_names_to_titles = { + "tool_with_progress": "Progress Tool", + "echo": "Echo Tool", + "sampling_tool": "Sampling Tool", + "notification_tool": "Notification Tool", + "echo_headers": "Echo Headers", + "echo_context": "Echo Context", + "book_restaurant": "Restaurant Booking", + } + + for tool in tools_result.tools: + if tool.name in tool_names_to_titles: + assert tool.title == tool_names_to_titles[tool.name] + # Test get_display_name utility + assert get_display_name(tool) == tool_names_to_titles[tool.name] + + # Test resources have titles + resources_result = await session.list_resources() + assert resources_result.resources + + # Check specific resources have titles + static_resource = next((r for r in resources_result.resources if r.name == "Static Info"), None) + assert static_resource is not None + assert static_resource.title == "Static Information" + assert get_display_name(static_resource) == "Static Information" + + # Test resource templates have titles + resource_templates = await session.list_resource_templates() + assert resource_templates.resourceTemplates + + # Check specific resource templates have titles + template_uris_to_titles = { + "resource://dynamic/{category}": "Dynamic Resource", + "resource://template/{id}/data": "Template Resource", + "github://repos/{owner}/{repo}": "GitHub Repository", + } + + for template in resource_templates.resourceTemplates: + if template.uriTemplate in template_uris_to_titles: + assert template.title == template_uris_to_titles[template.uriTemplate] + assert get_display_name(template) == template_uris_to_titles[template.uriTemplate] + + # Test prompts have titles + prompts_result = await session.list_prompts() + assert prompts_result.prompts + + # Check specific prompts have titles + prompt_names_to_titles = { + "simple_prompt": "Simple Prompt", + "complex_prompt": "Complex Prompt", + } + + for prompt in prompts_result.prompts: + if prompt.name in prompt_names_to_titles: + assert prompt.title == prompt_names_to_titles[prompt.name] + assert get_display_name(prompt) == prompt_names_to_titles[prompt.name] diff --git a/tests/server/fastmcp/test_title.py b/tests/server/fastmcp/test_title.py new file mode 100644 index 000000000..a94f6671d --- /dev/null +++ b/tests/server/fastmcp/test_title.py @@ -0,0 +1,215 @@ +"""Integration tests for title field functionality.""" + +import pytest +from pydantic import AnyUrl + +from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp.resources import FunctionResource +from mcp.shared.memory import create_connected_server_and_client_session +from mcp.shared.metadata_utils import get_display_name +from mcp.types import Prompt, Resource, ResourceTemplate, Tool, ToolAnnotations + + +@pytest.mark.anyio +async def test_tool_title_precedence(): + """Test that tool title precedence works correctly: title > annotations.title > name.""" + # Create server with various tool configurations + mcp = FastMCP(name="TitleTestServer") + + # Tool with only name + @mcp.tool(description="Basic tool") + def basic_tool(message: str) -> str: + return message + + # Tool with title + @mcp.tool(description="Tool with title", title="User-Friendly Tool") + def tool_with_title(message: str) -> str: + return message + + # Tool with annotations.title (when title is not supported on decorator) + # We'll need to add this manually after registration + @mcp.tool(description="Tool with annotations") + def tool_with_annotations(message: str) -> str: + return message + + # Tool with both title and annotations.title + @mcp.tool(description="Tool with both", title="Primary Title") + def tool_with_both(message: str) -> str: + return message + + # Start server and connect client + async with create_connected_server_and_client_session(mcp._mcp_server) as client: + await client.initialize() + + # List tools + tools_result = await client.list_tools() + tools = {tool.name: tool for tool in tools_result.tools} + + # Verify basic tool uses name + assert "basic_tool" in tools + basic = tools["basic_tool"] + # Since we haven't implemented get_display_name yet, we'll check the raw fields + assert basic.title is None + assert basic.name == "basic_tool" + + # Verify tool with title + assert "tool_with_title" in tools + titled = tools["tool_with_title"] + assert titled.title == "User-Friendly Tool" + + # For now, we'll skip the annotations.title test as it requires modifying + # the tool after registration, which we'll implement later + + # Verify tool with both uses title over annotations.title + assert "tool_with_both" in tools + both = tools["tool_with_both"] + assert both.title == "Primary Title" + + +@pytest.mark.anyio +async def test_prompt_title(): + """Test that prompt titles work correctly.""" + mcp = FastMCP(name="PromptTitleServer") + + # Prompt with only name + @mcp.prompt(description="Basic prompt") + def basic_prompt(topic: str) -> str: + return f"Tell me about {topic}" + + # Prompt with title + @mcp.prompt(description="Titled prompt", title="Ask About Topic") + def titled_prompt(topic: str) -> str: + return f"Tell me about {topic}" + + # Start server and connect client + async with create_connected_server_and_client_session(mcp._mcp_server) as client: + await client.initialize() + + # List prompts + prompts_result = await client.list_prompts() + prompts = {prompt.name: prompt for prompt in prompts_result.prompts} + + # Verify basic prompt uses name + assert "basic_prompt" in prompts + basic = prompts["basic_prompt"] + assert basic.title is None + assert basic.name == "basic_prompt" + + # Verify prompt with title + assert "titled_prompt" in prompts + titled = prompts["titled_prompt"] + assert titled.title == "Ask About Topic" + + +@pytest.mark.anyio +async def test_resource_title(): + """Test that resource titles work correctly.""" + mcp = FastMCP(name="ResourceTitleServer") + + # Static resource without title + def get_basic_data() -> str: + return "Basic data" + + basic_resource = FunctionResource( + uri=AnyUrl("resource://basic"), + name="basic_resource", + description="Basic resource", + fn=get_basic_data, + ) + mcp.add_resource(basic_resource) + + # Static resource with title + def get_titled_data() -> str: + return "Titled data" + + titled_resource = FunctionResource( + uri=AnyUrl("resource://titled"), + name="titled_resource", + title="User-Friendly Resource", + description="Resource with title", + fn=get_titled_data, + ) + mcp.add_resource(titled_resource) + + # Dynamic resource without title + @mcp.resource("resource://dynamic/{id}") + def dynamic_resource(id: str) -> str: + return f"Data for {id}" + + # Dynamic resource with title (when supported) + @mcp.resource("resource://titled-dynamic/{id}", title="Dynamic Data") + def titled_dynamic_resource(id: str) -> str: + return f"Data for {id}" + + # Start server and connect client + async with create_connected_server_and_client_session(mcp._mcp_server) as client: + await client.initialize() + + # List resources + resources_result = await client.list_resources() + resources = {str(res.uri): res for res in resources_result.resources} + + # Verify basic resource uses name + assert "resource://basic" in resources + basic = resources["resource://basic"] + assert basic.title is None + assert basic.name == "basic_resource" + + # Verify resource with title + assert "resource://titled" in resources + titled = resources["resource://titled"] + assert titled.title == "User-Friendly Resource" + + # List resource templates + templates_result = await client.list_resource_templates() + templates = {tpl.uriTemplate: tpl for tpl in templates_result.resourceTemplates} + + # Verify dynamic resource template + assert "resource://dynamic/{id}" in templates + dynamic = templates["resource://dynamic/{id}"] + assert dynamic.title is None + assert dynamic.name == "dynamic_resource" + + # Verify titled dynamic resource template (when supported) + if "resource://titled-dynamic/{id}" in templates: + titled_dynamic = templates["resource://titled-dynamic/{id}"] + assert titled_dynamic.title == "Dynamic Data" + + +@pytest.mark.anyio +async def test_get_display_name_utility(): + """Test the get_display_name utility function.""" + + # Test tool precedence: title > annotations.title > name + tool_name_only = Tool(name="test_tool", inputSchema={}) + assert get_display_name(tool_name_only) == "test_tool" + + tool_with_title = Tool(name="test_tool", title="Test Tool", inputSchema={}) + assert get_display_name(tool_with_title) == "Test Tool" + + tool_with_annotations = Tool(name="test_tool", inputSchema={}, annotations=ToolAnnotations(title="Annotated Tool")) + assert get_display_name(tool_with_annotations) == "Annotated Tool" + + tool_with_both = Tool( + name="test_tool", title="Primary Title", inputSchema={}, annotations=ToolAnnotations(title="Secondary Title") + ) + assert get_display_name(tool_with_both) == "Primary Title" + + # Test other types: title > name + resource = Resource(uri=AnyUrl("file://test"), name="test_res") + assert get_display_name(resource) == "test_res" + + resource_with_title = Resource(uri=AnyUrl("file://test"), name="test_res", title="Test Resource") + assert get_display_name(resource_with_title) == "Test Resource" + + prompt = Prompt(name="test_prompt") + assert get_display_name(prompt) == "test_prompt" + + prompt_with_title = Prompt(name="test_prompt", title="Test Prompt") + assert get_display_name(prompt_with_title) == "Test Prompt" + + template = ResourceTemplate(uriTemplate="file://{id}", name="test_template") + assert get_display_name(template) == "test_template" + + template_with_title = ResourceTemplate(uriTemplate="file://{id}", name="test_template", title="Test Template") + assert get_display_name(template_with_title) == "Test Template" diff --git a/tests/server/fastmcp/test_tool_manager.py b/tests/server/fastmcp/test_tool_manager.py index 7954e1729..206df42d7 100644 --- a/tests/server/fastmcp/test_tool_manager.py +++ b/tests/server/fastmcp/test_tool_manager.py @@ -44,6 +44,7 @@ class AddArguments(ArgModelBase): original_tool = Tool( name="add", + title="Add Tool", description="Add two numbers.", fn=add, fn_metadata=fn_metadata, From 86bb54cb284c6db293b280671f1a8d0915d93ccb Mon Sep 17 00:00:00 2001 From: David Dworken Date: Tue, 17 Jun 2025 12:13:00 -0700 Subject: [PATCH 13/14] Add support for DNS rebinding protections (#861) --- src/mcp/server/fastmcp/server.py | 6 + src/mcp/server/sse.py | 24 +- src/mcp/server/streamable_http.py | 16 + src/mcp/server/streamable_http_manager.py | 5 + src/mcp/server/transport_security.py | 127 ++++++++ tests/server/fastmcp/test_integration.py | 25 +- tests/server/test_sse_security.py | 293 ++++++++++++++++++ tests/server/test_streamable_http_security.py | 293 ++++++++++++++++++ tests/shared/test_sse.py | 13 +- tests/shared/test_streamable_http.py | 10 +- 10 files changed, 799 insertions(+), 13 deletions(-) create mode 100644 src/mcp/server/transport_security.py create mode 100644 tests/server/test_sse_security.py create mode 100644 tests/server/test_streamable_http_security.py diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 3adc465b8..e555a5239 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -50,6 +50,7 @@ from mcp.server.stdio import stdio_server from mcp.server.streamable_http import EventStore from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from mcp.server.transport_security import TransportSecuritySettings from mcp.shared.context import LifespanContextT, RequestContext, RequestT from mcp.types import ( AnyFunction, @@ -118,6 +119,9 @@ class Settings(BaseSettings, Generic[LifespanResultT]): auth: AuthSettings | None = None + # Transport security settings (DNS rebinding protection) + transport_security: TransportSecuritySettings | None = None + def lifespan_wrapper( app: FastMCP, @@ -674,6 +678,7 @@ def sse_app(self, mount_path: str | None = None) -> Starlette: sse = SseServerTransport( normalized_message_endpoint, + security_settings=self.settings.transport_security, ) async def handle_sse(scope: Scope, receive: Receive, send: Send): @@ -779,6 +784,7 @@ def streamable_http_app(self) -> Starlette: event_store=self._event_store, json_response=self.settings.json_response, stateless=self.settings.stateless_http, # Use the stateless setting + security_settings=self.settings.transport_security, ) # Create the ASGI handler diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py index 52f273968..41145e49f 100644 --- a/src/mcp/server/sse.py +++ b/src/mcp/server/sse.py @@ -52,6 +52,10 @@ async def handle_sse(request): from starlette.types import Receive, Scope, Send import mcp.types as types +from mcp.server.transport_security import ( + TransportSecurityMiddleware, + TransportSecuritySettings, +) from mcp.shared.message import ServerMessageMetadata, SessionMessage logger = logging.getLogger(__name__) @@ -71,16 +75,22 @@ class SseServerTransport: _endpoint: str _read_stream_writers: dict[UUID, MemoryObjectSendStream[SessionMessage | Exception]] + _security: TransportSecurityMiddleware - def __init__(self, endpoint: str) -> None: + def __init__(self, endpoint: str, security_settings: TransportSecuritySettings | None = None) -> None: """ Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL given. + + Args: + endpoint: The relative or absolute URL for POST messages. + security_settings: Optional security settings for DNS rebinding protection. """ super().__init__() self._endpoint = endpoint self._read_stream_writers = {} + self._security = TransportSecurityMiddleware(security_settings) logger.debug(f"SseServerTransport initialized with endpoint: {endpoint}") @asynccontextmanager @@ -89,6 +99,13 @@ async def connect_sse(self, scope: Scope, receive: Receive, send: Send): logger.error("connect_sse received non-HTTP request") raise ValueError("connect_sse can only handle HTTP requests") + # Validate request headers for DNS rebinding protection + request = Request(scope, receive) + error_response = await self._security.validate_request(request, is_post=False) + if error_response: + await error_response(scope, receive, send) + raise ValueError("Request validation failed") + logger.debug("Setting up SSE connection") read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception] @@ -160,6 +177,11 @@ async def handle_post_message(self, scope: Scope, receive: Receive, send: Send) logger.debug("Handling POST message") request = Request(scope, receive) + # Validate request headers for DNS rebinding protection + error_response = await self._security.validate_request(request, is_post=True) + if error_response: + return await error_response(scope, receive, send) + session_id_param = request.query_params.get("session_id") if session_id_param is None: logger.warning("Received request without session_id") diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py index 13ee27b32..d46549929 100644 --- a/src/mcp/server/streamable_http.py +++ b/src/mcp/server/streamable_http.py @@ -24,6 +24,10 @@ from starlette.responses import Response from starlette.types import Receive, Scope, Send +from mcp.server.transport_security import ( + TransportSecurityMiddleware, + TransportSecuritySettings, +) from mcp.shared.message import ServerMessageMetadata, SessionMessage from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS from mcp.types import ( @@ -130,12 +134,14 @@ class StreamableHTTPServerTransport: _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception] | None = None _write_stream: MemoryObjectSendStream[SessionMessage] | None = None _write_stream_reader: MemoryObjectReceiveStream[SessionMessage] | None = None + _security: TransportSecurityMiddleware def __init__( self, mcp_session_id: str | None, is_json_response_enabled: bool = False, event_store: EventStore | None = None, + security_settings: TransportSecuritySettings | None = None, ) -> None: """ Initialize a new StreamableHTTP server transport. @@ -148,6 +154,7 @@ def __init__( event_store: Event store for resumability support. If provided, resumability will be enabled, allowing clients to reconnect and resume messages. + security_settings: Optional security settings for DNS rebinding protection. Raises: ValueError: If the session ID contains invalid characters. @@ -158,6 +165,7 @@ def __init__( self.mcp_session_id = mcp_session_id self.is_json_response_enabled = is_json_response_enabled self._event_store = event_store + self._security = TransportSecurityMiddleware(security_settings) self._request_streams: dict[ RequestId, tuple[ @@ -251,6 +259,14 @@ async def _clean_up_memory_streams(self, request_id: RequestId) -> None: async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None: """Application entry point that handles all HTTP requests""" request = Request(scope, receive) + + # Validate request headers for DNS rebinding protection + is_post = request.method == "POST" + error_response = await self._security.validate_request(request, is_post=is_post) + if error_response: + await error_response(scope, receive, send) + return + if self._terminated: # If the session has been terminated, return 404 Not Found response = self._create_error_response( diff --git a/src/mcp/server/streamable_http_manager.py b/src/mcp/server/streamable_http_manager.py index 3af9829d1..41b807388 100644 --- a/src/mcp/server/streamable_http_manager.py +++ b/src/mcp/server/streamable_http_manager.py @@ -22,6 +22,7 @@ EventStore, StreamableHTTPServerTransport, ) +from mcp.server.transport_security import TransportSecuritySettings logger = logging.getLogger(__name__) @@ -60,11 +61,13 @@ def __init__( event_store: EventStore | None = None, json_response: bool = False, stateless: bool = False, + security_settings: TransportSecuritySettings | None = None, ): self.app = app self.event_store = event_store self.json_response = json_response self.stateless = stateless + self.security_settings = security_settings # Session tracking (only used if not stateless) self._session_creation_lock = anyio.Lock() @@ -162,6 +165,7 @@ async def _handle_stateless_request( mcp_session_id=None, # No session tracking in stateless mode is_json_response_enabled=self.json_response, event_store=None, # No event store in stateless mode + security_settings=self.security_settings, ) # Start server in a new task @@ -217,6 +221,7 @@ async def _handle_stateful_request( mcp_session_id=new_session_id, is_json_response_enabled=self.json_response, event_store=self.event_store, # May be None (no resumability) + security_settings=self.security_settings, ) assert http_transport.mcp_session_id is not None diff --git a/src/mcp/server/transport_security.py b/src/mcp/server/transport_security.py new file mode 100644 index 000000000..3a884ee2b --- /dev/null +++ b/src/mcp/server/transport_security.py @@ -0,0 +1,127 @@ +"""DNS rebinding protection for MCP server transports.""" + +import logging + +from pydantic import BaseModel, Field +from starlette.requests import Request +from starlette.responses import Response + +logger = logging.getLogger(__name__) + + +class TransportSecuritySettings(BaseModel): + """Settings for MCP transport security features. + + These settings help protect against DNS rebinding attacks by validating + incoming request headers. + """ + + enable_dns_rebinding_protection: bool = Field( + default=True, + description="Enable DNS rebinding protection (recommended for production)", + ) + + allowed_hosts: list[str] = Field( + default=[], + description="List of allowed Host header values. Only applies when " + + "enable_dns_rebinding_protection is True.", + ) + + allowed_origins: list[str] = Field( + default=[], + description="List of allowed Origin header values. Only applies when " + + "enable_dns_rebinding_protection is True.", + ) + + +class TransportSecurityMiddleware: + """Middleware to enforce DNS rebinding protection for MCP transport endpoints.""" + + def __init__(self, settings: TransportSecuritySettings | None = None): + # If not specified, disable DNS rebinding protection by default + # for backwards compatibility + self.settings = settings or TransportSecuritySettings(enable_dns_rebinding_protection=False) + + def _validate_host(self, host: str | None) -> bool: + """Validate the Host header against allowed values.""" + if not host: + logger.warning("Missing Host header in request") + return False + + # Check exact match first + if host in self.settings.allowed_hosts: + return True + + # Check wildcard port patterns + for allowed in self.settings.allowed_hosts: + if allowed.endswith(":*"): + # Extract base host from pattern + base_host = allowed[:-2] + # Check if the actual host starts with base host and has a port + if host.startswith(base_host + ":"): + return True + + logger.warning(f"Invalid Host header: {host}") + return False + + def _validate_origin(self, origin: str | None) -> bool: + """Validate the Origin header against allowed values.""" + # Origin can be absent for same-origin requests + if not origin: + return True + + # Check exact match first + if origin in self.settings.allowed_origins: + return True + + # Check wildcard port patterns + for allowed in self.settings.allowed_origins: + if allowed.endswith(":*"): + # Extract base origin from pattern + base_origin = allowed[:-2] + # Check if the actual origin starts with base origin and has a port + if origin.startswith(base_origin + ":"): + return True + + logger.warning(f"Invalid Origin header: {origin}") + return False + + def _validate_content_type(self, content_type: str | None) -> bool: + """Validate the Content-Type header for POST requests.""" + if not content_type: + logger.warning("Missing Content-Type header in POST request") + return False + + # Content-Type must start with application/json + if not content_type.lower().startswith("application/json"): + logger.warning(f"Invalid Content-Type header: {content_type}") + return False + + return True + + async def validate_request(self, request: Request, is_post: bool = False) -> Response | None: + """Validate request headers for DNS rebinding protection. + + Returns None if validation passes, or an error Response if validation fails. + """ + # Always validate Content-Type for POST requests + if is_post: + content_type = request.headers.get("content-type") + if not self._validate_content_type(content_type): + return Response("Invalid Content-Type header", status_code=400) + + # Skip remaining validation if DNS rebinding protection is disabled + if not self.settings.enable_dns_rebinding_protection: + return None + + # Validate Host header + host = request.headers.get("host") + if not self._validate_host(host): + return Response("Invalid Host header", status_code=421) + + # Validate Origin header + origin = request.headers.get("origin") + if not self._validate_origin(origin): + return Response("Invalid Origin header", status_code=400) + + return None diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index 7224a8e48..c41e710cd 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -23,6 +23,7 @@ from mcp.client.streamable_http import streamablehttp_client from mcp.server.fastmcp import Context, FastMCP from mcp.server.fastmcp.resources import FunctionResource +from mcp.server.transport_security import TransportSecuritySettings from mcp.shared.context import RequestContext from mcp.types import ( Completion, @@ -92,7 +93,10 @@ def stateless_http_server_url(stateless_http_server_port: int) -> str: # Create a function to make the FastMCP server app def make_fastmcp_app(): """Create a FastMCP server without auth settings.""" - mcp = FastMCP(name="NoAuthServer") + transport_security = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + mcp = FastMCP(name="NoAuthServer", transport_security=transport_security) # Add a simple tool @mcp.tool(description="A simple echo tool") @@ -121,9 +125,10 @@ class AnswerSchema(BaseModel): def make_everything_fastmcp() -> FastMCP: """Create a FastMCP server with all features enabled for testing.""" - from mcp.server.fastmcp import Context - - mcp = FastMCP(name="EverythingServer") + transport_security = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + mcp = FastMCP(name="EverythingServer", transport_security=transport_security) # Tool with context for logging and progress @mcp.tool(description="A tool that demonstrates logging and progress", title="Progress Tool") @@ -333,8 +338,10 @@ def make_everything_fastmcp_app(): def make_fastmcp_streamable_http_app(): """Create a FastMCP server with StreamableHTTP transport.""" - - mcp = FastMCP(name="NoAuthServer") + transport_security = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + mcp = FastMCP(name="NoAuthServer", transport_security=transport_security) # Add a simple tool @mcp.tool(description="A simple echo tool") @@ -359,8 +366,10 @@ def make_everything_fastmcp_streamable_http_app(): def make_fastmcp_stateless_http_app(): """Create a FastMCP server with stateless StreamableHTTP transport.""" - - mcp = FastMCP(name="StatelessServer", stateless_http=True) + transport_security = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + mcp = FastMCP(name="StatelessServer", stateless_http=True, transport_security=transport_security) # Add a simple tool @mcp.tool(description="A simple echo tool") diff --git a/tests/server/test_sse_security.py b/tests/server/test_sse_security.py new file mode 100644 index 000000000..43af35061 --- /dev/null +++ b/tests/server/test_sse_security.py @@ -0,0 +1,293 @@ +"""Tests for SSE server DNS rebinding protection.""" + +import logging +import multiprocessing +import socket +import time + +import httpx +import pytest +import uvicorn +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import Response +from starlette.routing import Mount, Route + +from mcp.server import Server +from mcp.server.sse import SseServerTransport +from mcp.server.transport_security import TransportSecuritySettings +from mcp.types import Tool + +logger = logging.getLogger(__name__) +SERVER_NAME = "test_sse_security_server" + + +@pytest.fixture +def server_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def server_url(server_port: int) -> str: + return f"http://127.0.0.1:{server_port}" + + +class SecurityTestServer(Server): + def __init__(self): + super().__init__(SERVER_NAME) + + async def on_list_tools(self) -> list[Tool]: + return [] + + +def run_server_with_settings(port: int, security_settings: TransportSecuritySettings | None = None): + """Run the SSE server with specified security settings.""" + app = SecurityTestServer() + sse_transport = SseServerTransport("/messages/", security_settings) + + async def handle_sse(request: Request): + try: + async with sse_transport.connect_sse(request.scope, request.receive, request._send) as streams: + if streams: + await app.run(streams[0], streams[1], app.create_initialization_options()) + except ValueError as e: + # Validation error was already handled inside connect_sse + logger.debug(f"SSE connection failed validation: {e}") + return Response() + + routes = [ + Route("/sse", endpoint=handle_sse), + Mount("/messages/", app=sse_transport.handle_post_message), + ] + + starlette_app = Starlette(routes=routes) + uvicorn.run(starlette_app, host="127.0.0.1", port=port, log_level="error") + + +def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None): + """Start server in a separate process.""" + process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings)) + process.start() + # Give server time to start + time.sleep(1) + return process + + +@pytest.mark.anyio +async def test_sse_security_default_settings(server_port: int): + """Test SSE with default security settings (protection disabled).""" + process = start_server_process(server_port) + + try: + headers = {"Host": "evil.com", "Origin": "http://evil.com"} + + async with httpx.AsyncClient(timeout=5.0) as client: + async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response: + assert response.status_code == 200 + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_invalid_host_header(server_port: int): + """Test SSE with invalid Host header.""" + # Enable security by providing settings with an empty allowed_hosts list + security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["example.com"]) + process = start_server_process(server_port, security_settings) + + try: + # Test with invalid host header + headers = {"Host": "evil.com"} + + async with httpx.AsyncClient() as client: + response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers) + assert response.status_code == 421 + assert response.text == "Invalid Host header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_invalid_origin_header(server_port: int): + """Test SSE with invalid Origin header.""" + # Configure security to allow the host but restrict origins + security_settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://localhost:*"] + ) + process = start_server_process(server_port, security_settings) + + try: + # Test with invalid origin header + headers = {"Origin": "http://evil.com"} + + async with httpx.AsyncClient() as client: + response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers) + assert response.status_code == 400 + assert response.text == "Invalid Origin header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_post_invalid_content_type(server_port: int): + """Test POST endpoint with invalid Content-Type header.""" + # Configure security to allow the host + security_settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"] + ) + process = start_server_process(server_port, security_settings) + + try: + async with httpx.AsyncClient(timeout=5.0) as client: + # Test POST with invalid content type + fake_session_id = "12345678123456781234567812345678" + response = await client.post( + f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}", + headers={"Content-Type": "text/plain"}, + content="test", + ) + assert response.status_code == 400 + assert response.text == "Invalid Content-Type header" + + # Test POST with missing content type + response = await client.post( + f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}", content="test" + ) + assert response.status_code == 400 + assert response.text == "Invalid Content-Type header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_disabled(server_port: int): + """Test SSE with security disabled.""" + settings = TransportSecuritySettings(enable_dns_rebinding_protection=False) + process = start_server_process(server_port, settings) + + try: + # Test with invalid host header - should still work + headers = {"Host": "evil.com"} + + async with httpx.AsyncClient(timeout=5.0) as client: + # For SSE endpoints, we need to use stream to avoid timeout + async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response: + # Should connect successfully even with invalid host + assert response.status_code == 200 + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_custom_allowed_hosts(server_port: int): + """Test SSE with custom allowed hosts.""" + settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, + allowed_hosts=["localhost", "127.0.0.1", "custom.host"], + allowed_origins=["http://localhost", "http://127.0.0.1", "http://custom.host"], + ) + process = start_server_process(server_port, settings) + + try: + # Test with custom allowed host + headers = {"Host": "custom.host"} + + async with httpx.AsyncClient(timeout=5.0) as client: + # For SSE endpoints, we need to use stream to avoid timeout + async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response: + # Should connect successfully with custom host + assert response.status_code == 200 + + # Test with non-allowed host + headers = {"Host": "evil.com"} + + async with httpx.AsyncClient() as client: + response = await client.get(f"http://127.0.0.1:{server_port}/sse", headers=headers) + assert response.status_code == 421 + assert response.text == "Invalid Host header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_wildcard_ports(server_port: int): + """Test SSE with wildcard port patterns.""" + settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, + allowed_hosts=["localhost:*", "127.0.0.1:*"], + allowed_origins=["http://localhost:*", "http://127.0.0.1:*"], + ) + process = start_server_process(server_port, settings) + + try: + # Test with various port numbers + for test_port in [8080, 3000, 9999]: + headers = {"Host": f"localhost:{test_port}"} + + async with httpx.AsyncClient(timeout=5.0) as client: + # For SSE endpoints, we need to use stream to avoid timeout + async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response: + # Should connect successfully with any port + assert response.status_code == 200 + + headers = {"Origin": f"http://localhost:{test_port}"} + + async with httpx.AsyncClient(timeout=5.0) as client: + # For SSE endpoints, we need to use stream to avoid timeout + async with client.stream("GET", f"http://127.0.0.1:{server_port}/sse", headers=headers) as response: + # Should connect successfully with any port + assert response.status_code == 200 + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_sse_security_post_valid_content_type(server_port: int): + """Test POST endpoint with valid Content-Type headers.""" + # Configure security to allow the host + security_settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"], allowed_origins=["http://127.0.0.1:*"] + ) + process = start_server_process(server_port, security_settings) + + try: + async with httpx.AsyncClient() as client: + # Test with various valid content types + valid_content_types = [ + "application/json", + "application/json; charset=utf-8", + "application/json;charset=utf-8", + "APPLICATION/JSON", # Case insensitive + ] + + for content_type in valid_content_types: + # Use a valid UUID format (even though session won't exist) + fake_session_id = "12345678123456781234567812345678" + response = await client.post( + f"http://127.0.0.1:{server_port}/messages/?session_id={fake_session_id}", + headers={"Content-Type": content_type}, + json={"test": "data"}, + ) + # Will get 404 because session doesn't exist, but that's OK + # We're testing that it passes the content-type check + assert response.status_code == 404 + assert response.text == "Could not find session" + + finally: + process.terminate() + process.join() diff --git a/tests/server/test_streamable_http_security.py b/tests/server/test_streamable_http_security.py new file mode 100644 index 000000000..eed791924 --- /dev/null +++ b/tests/server/test_streamable_http_security.py @@ -0,0 +1,293 @@ +"""Tests for StreamableHTTP server DNS rebinding protection.""" + +import logging +import multiprocessing +import socket +import time +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager + +import httpx +import pytest +import uvicorn +from starlette.applications import Starlette +from starlette.routing import Mount +from starlette.types import Receive, Scope, Send + +from mcp.server import Server +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from mcp.server.transport_security import TransportSecuritySettings +from mcp.types import Tool + +logger = logging.getLogger(__name__) +SERVER_NAME = "test_streamable_http_security_server" + + +@pytest.fixture +def server_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def server_url(server_port: int) -> str: + return f"http://127.0.0.1:{server_port}" + + +class SecurityTestServer(Server): + def __init__(self): + super().__init__(SERVER_NAME) + + async def on_list_tools(self) -> list[Tool]: + return [] + + +def run_server_with_settings(port: int, security_settings: TransportSecuritySettings | None = None): + """Run the StreamableHTTP server with specified security settings.""" + app = SecurityTestServer() + + # Create session manager with security settings + session_manager = StreamableHTTPSessionManager( + app=app, + json_response=False, + stateless=False, + security_settings=security_settings, + ) + + # Create the ASGI handler + async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> None: + await session_manager.handle_request(scope, receive, send) + + # Create Starlette app with lifespan + @asynccontextmanager + async def lifespan(app: Starlette) -> AsyncGenerator[None, None]: + async with session_manager.run(): + yield + + routes = [ + Mount("/", app=handle_streamable_http), + ] + + starlette_app = Starlette(routes=routes, lifespan=lifespan) + uvicorn.run(starlette_app, host="127.0.0.1", port=port, log_level="error") + + +def start_server_process(port: int, security_settings: TransportSecuritySettings | None = None): + """Start server in a separate process.""" + process = multiprocessing.Process(target=run_server_with_settings, args=(port, security_settings)) + process.start() + # Give server time to start + time.sleep(1) + return process + + +@pytest.mark.anyio +async def test_streamable_http_security_default_settings(server_port: int): + """Test StreamableHTTP with default security settings (protection enabled).""" + process = start_server_process(server_port) + + try: + # Test with valid localhost headers + async with httpx.AsyncClient(timeout=5.0) as client: + # POST request to initialize session + response = await client.post( + f"http://127.0.0.1:{server_port}/", + json={"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {}}, + headers={ + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + }, + ) + assert response.status_code == 200 + assert "mcp-session-id" in response.headers + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_invalid_host_header(server_port: int): + """Test StreamableHTTP with invalid Host header.""" + security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True) + process = start_server_process(server_port, security_settings) + + try: + # Test with invalid host header + headers = { + "Host": "evil.com", + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.post( + f"http://127.0.0.1:{server_port}/", + json={"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {}}, + headers=headers, + ) + assert response.status_code == 421 + assert response.text == "Invalid Host header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_invalid_origin_header(server_port: int): + """Test StreamableHTTP with invalid Origin header.""" + security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1:*"]) + process = start_server_process(server_port, security_settings) + + try: + # Test with invalid origin header + headers = { + "Origin": "http://evil.com", + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.post( + f"http://127.0.0.1:{server_port}/", + json={"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {}}, + headers=headers, + ) + assert response.status_code == 400 + assert response.text == "Invalid Origin header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_invalid_content_type(server_port: int): + """Test StreamableHTTP POST with invalid Content-Type header.""" + process = start_server_process(server_port) + + try: + async with httpx.AsyncClient(timeout=5.0) as client: + # Test POST with invalid content type + response = await client.post( + f"http://127.0.0.1:{server_port}/", + headers={ + "Content-Type": "text/plain", + "Accept": "application/json, text/event-stream", + }, + content="test", + ) + assert response.status_code == 400 + assert response.text == "Invalid Content-Type header" + + # Test POST with missing content type + response = await client.post( + f"http://127.0.0.1:{server_port}/", + headers={"Accept": "application/json, text/event-stream"}, + content="test", + ) + assert response.status_code == 400 + assert response.text == "Invalid Content-Type header" + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_disabled(server_port: int): + """Test StreamableHTTP with security disabled.""" + settings = TransportSecuritySettings(enable_dns_rebinding_protection=False) + process = start_server_process(server_port, settings) + + try: + # Test with invalid host header - should still work + headers = { + "Host": "evil.com", + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.post( + f"http://127.0.0.1:{server_port}/", + json={"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {}}, + headers=headers, + ) + # Should connect successfully even with invalid host + assert response.status_code == 200 + + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_custom_allowed_hosts(server_port: int): + """Test StreamableHTTP with custom allowed hosts.""" + settings = TransportSecuritySettings( + enable_dns_rebinding_protection=True, + allowed_hosts=["localhost", "127.0.0.1", "custom.host"], + allowed_origins=["http://localhost", "http://127.0.0.1", "http://custom.host"], + ) + process = start_server_process(server_port, settings) + + try: + # Test with custom allowed host + headers = { + "Host": "custom.host", + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.post( + f"http://127.0.0.1:{server_port}/", + json={"jsonrpc": "2.0", "method": "initialize", "id": 1, "params": {}}, + headers=headers, + ) + # Should connect successfully with custom host + assert response.status_code == 200 + finally: + process.terminate() + process.join() + + +@pytest.mark.anyio +async def test_streamable_http_security_get_request(server_port: int): + """Test StreamableHTTP GET request with security.""" + security_settings = TransportSecuritySettings(enable_dns_rebinding_protection=True, allowed_hosts=["127.0.0.1"]) + process = start_server_process(server_port, security_settings) + + try: + # Test GET request with invalid host header + headers = { + "Host": "evil.com", + "Accept": "text/event-stream", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.get(f"http://127.0.0.1:{server_port}/", headers=headers) + assert response.status_code == 421 + assert response.text == "Invalid Host header" + + # Test GET request with valid host header + headers = { + "Host": "127.0.0.1", + "Accept": "text/event-stream", + } + + async with httpx.AsyncClient(timeout=5.0) as client: + # GET requests need a session ID in StreamableHTTP + # So it will fail with "Missing session ID" not security error + response = await client.get(f"http://127.0.0.1:{server_port}/", headers=headers) + # This should pass security but fail on session validation + assert response.status_code == 400 + body = response.json() + assert "Missing session ID" in body["error"]["message"] + + finally: + process.terminate() + process.join() diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 4d8f7717e..8e1912e9b 100644 --- a/tests/shared/test_sse.py +++ b/tests/shared/test_sse.py @@ -20,6 +20,7 @@ from mcp.client.sse import sse_client from mcp.server import Server from mcp.server.sse import SseServerTransport +from mcp.server.transport_security import TransportSecuritySettings from mcp.shared.exceptions import McpError from mcp.types import ( EmptyResult, @@ -80,7 +81,11 @@ async def handle_call_tool(name: str, args: dict) -> list[TextContent]: # Test fixtures def make_server_app() -> Starlette: """Create test Starlette app with SSE transport""" - sse = SseServerTransport("/messages/") + # Configure security with allowed hosts/origins for testing + security_settings = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + sse = SseServerTransport("/messages/", security_settings=security_settings) server = ServerTest() async def handle_sse(request: Request) -> Response: @@ -339,7 +344,11 @@ async def handle_list_tools() -> list[Tool]: def run_context_server(server_port: int) -> None: """Run a server that captures request context""" - sse = SseServerTransport("/messages/") + # Configure security with allowed hosts/origins for testing + security_settings = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) + sse = SseServerTransport("/messages/", security_settings=security_settings) context_server = RequestContextServer() async def handle_sse(request: Request) -> Response: diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py index 68d58fc91..1ffcc13b0 100644 --- a/tests/shared/test_streamable_http.py +++ b/tests/shared/test_streamable_http.py @@ -37,6 +37,7 @@ StreamId, ) from mcp.server.streamable_http_manager import StreamableHTTPSessionManager +from mcp.server.transport_security import TransportSecuritySettings from mcp.shared.context import RequestContext from mcp.shared.exceptions import McpError from mcp.shared.message import ( @@ -227,10 +228,14 @@ def create_app(is_json_response_enabled=False, event_store: EventStore | None = server = ServerTest() # Create the session manager + security_settings = TransportSecuritySettings( + allowed_hosts=["127.0.0.1:*", "localhost:*"], allowed_origins=["http://127.0.0.1:*", "http://localhost:*"] + ) session_manager = StreamableHTTPSessionManager( app=server, event_store=event_store, json_response=is_json_response_enabled, + security_settings=security_settings, ) # Create an ASGI application that uses the session manager @@ -436,8 +441,9 @@ def test_content_type_validation(basic_server, basic_server_url): }, data="This is not JSON", ) - assert response.status_code == 415 - assert "Unsupported Media Type" in response.text + + assert response.status_code == 400 + assert "Invalid Content-Type" in response.text def test_json_validation(basic_server, basic_server_url): From d0443a18328a2fc9e87da6feee9130f95c0b37a7 Mon Sep 17 00:00:00 2001 From: Inna Harper Date: Tue, 17 Jun 2025 20:23:07 +0100 Subject: [PATCH 14/14] Add resource Link (#974) --- .../server.py | 2 +- .../mcp_simple_streamablehttp/server.py | 2 +- .../simple-tool/mcp_simple_tool/server.py | 4 +-- src/mcp/server/fastmcp/prompts/base.py | 10 +++--- src/mcp/server/fastmcp/server.py | 8 ++--- src/mcp/server/lowlevel/server.py | 2 +- src/mcp/types.py | 20 ++++++++++-- tests/issues/test_88_random_error.py | 4 +-- tests/server/fastmcp/test_integration.py | 32 ++++++++++++++++++- tests/server/fastmcp/test_server.py | 4 +-- 10 files changed, 66 insertions(+), 22 deletions(-) diff --git a/examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py b/examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py index d76e19550..6a9ff9364 100644 --- a/examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py +++ b/examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py @@ -41,7 +41,7 @@ def main( app = Server("mcp-streamable-http-stateless-demo") @app.call_tool() - async def call_tool(name: str, arguments: dict) -> list[types.Content]: + async def call_tool(name: str, arguments: dict) -> list[types.ContentBlock]: ctx = app.request_context interval = arguments.get("interval", 1.0) count = arguments.get("count", 5) diff --git a/examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py b/examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py index c6b13ab0f..85eb1369f 100644 --- a/examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py +++ b/examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py @@ -45,7 +45,7 @@ def main( app = Server("mcp-streamable-http-demo") @app.call_tool() - async def call_tool(name: str, arguments: dict) -> list[types.Content]: + async def call_tool(name: str, arguments: dict) -> list[types.ContentBlock]: ctx = app.request_context interval = arguments.get("interval", 1.0) count = arguments.get("count", 5) diff --git a/examples/servers/simple-tool/mcp_simple_tool/server.py b/examples/servers/simple-tool/mcp_simple_tool/server.py index a6a3f0bd2..bf3683c9e 100644 --- a/examples/servers/simple-tool/mcp_simple_tool/server.py +++ b/examples/servers/simple-tool/mcp_simple_tool/server.py @@ -7,7 +7,7 @@ async def fetch_website( url: str, -) -> list[types.Content]: +) -> list[types.ContentBlock]: headers = { "User-Agent": "MCP Test Server (github.com/modelcontextprotocol/python-sdk)" } @@ -29,7 +29,7 @@ def main(port: int, transport: str) -> int: app = Server("mcp-website-fetcher") @app.call_tool() - async def fetch_tool(name: str, arguments: dict) -> list[types.Content]: + async def fetch_tool(name: str, arguments: dict) -> list[types.ContentBlock]: if name != "fetch": raise ValueError(f"Unknown tool: {name}") if "url" not in arguments: diff --git a/src/mcp/server/fastmcp/prompts/base.py b/src/mcp/server/fastmcp/prompts/base.py index e4751feb1..b45cfc917 100644 --- a/src/mcp/server/fastmcp/prompts/base.py +++ b/src/mcp/server/fastmcp/prompts/base.py @@ -7,16 +7,16 @@ import pydantic_core from pydantic import BaseModel, Field, TypeAdapter, validate_call -from mcp.types import Content, TextContent +from mcp.types import ContentBlock, TextContent class Message(BaseModel): """Base class for all prompt messages.""" role: Literal["user", "assistant"] - content: Content + content: ContentBlock - def __init__(self, content: str | Content, **kwargs: Any): + def __init__(self, content: str | ContentBlock, **kwargs: Any): if isinstance(content, str): content = TextContent(type="text", text=content) super().__init__(content=content, **kwargs) @@ -27,7 +27,7 @@ class UserMessage(Message): role: Literal["user", "assistant"] = "user" - def __init__(self, content: str | Content, **kwargs: Any): + def __init__(self, content: str | ContentBlock, **kwargs: Any): super().__init__(content=content, **kwargs) @@ -36,7 +36,7 @@ class AssistantMessage(Message): role: Literal["user", "assistant"] = "assistant" - def __init__(self, content: str | Content, **kwargs: Any): + def __init__(self, content: str | ContentBlock, **kwargs: Any): super().__init__(content=content, **kwargs) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index e555a5239..1b761e917 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -54,7 +54,7 @@ from mcp.shared.context import LifespanContextT, RequestContext, RequestT from mcp.types import ( AnyFunction, - Content, + ContentBlock, GetPromptResult, TextContent, ToolAnnotations, @@ -260,7 +260,7 @@ def get_context(self) -> Context[ServerSession, object, Request]: request_context = None return Context(request_context=request_context, fastmcp=self) - async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[Content]: + async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock]: """Call a tool by name with arguments.""" context = self.get_context() result = await self._tool_manager.call_tool(name, arguments, context=context) @@ -878,12 +878,12 @@ async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) - def _convert_to_content( result: Any, -) -> Sequence[Content]: +) -> Sequence[ContentBlock]: """Convert a result to a sequence of content objects.""" if result is None: return [] - if isinstance(result, Content): + if isinstance(result, ContentBlock): return [result] if isinstance(result, Image): diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index 617a4813d..0a8ab7f97 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -384,7 +384,7 @@ def call_tool(self): def decorator( func: Callable[ ..., - Awaitable[Iterable[types.Content]], + Awaitable[Iterable[types.ContentBlock]], ], ): logger.debug("Registering handler for CallToolRequest") diff --git a/src/mcp/types.py b/src/mcp/types.py index f53813412..d5663dad6 100644 --- a/src/mcp/types.py +++ b/src/mcp/types.py @@ -733,14 +733,28 @@ class EmbeddedResource(BaseModel): model_config = ConfigDict(extra="allow") -Content = TextContent | ImageContent | AudioContent | EmbeddedResource +class ResourceLink(Resource): + """ + A resource that the server is capable of reading, included in a prompt or tool call result. + + Note: resource links returned by tools are not guaranteed to appear in the results of `resources/list` requests. + """ + + type: Literal["resource_link"] + + +ContentBlock = TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource +"""A content block that can be used in prompts and tool results.""" + +Content: TypeAlias = ContentBlock +# """DEPRECATED: Content is deprecated, you should use ContentBlock directly.""" class PromptMessage(BaseModel): """Describes a message returned as part of a prompt.""" role: Role - content: Content + content: ContentBlock model_config = ConfigDict(extra="allow") @@ -859,7 +873,7 @@ class CallToolRequest(Request[CallToolRequestParams, Literal["tools/call"]]): class CallToolResult(Result): """The server's response to a tool call.""" - content: list[Content] + content: list[ContentBlock] isError: bool = False diff --git a/tests/issues/test_88_random_error.py b/tests/issues/test_88_random_error.py index 3647871a5..7ba970f0b 100644 --- a/tests/issues/test_88_random_error.py +++ b/tests/issues/test_88_random_error.py @@ -11,7 +11,7 @@ from mcp.client.session import ClientSession from mcp.server.lowlevel import Server from mcp.shared.exceptions import McpError -from mcp.types import Content, TextContent +from mcp.types import ContentBlock, TextContent @pytest.mark.anyio @@ -31,7 +31,7 @@ async def test_notification_validation_error(tmp_path: Path): slow_request_complete = anyio.Event() @server.call_tool() - async def slow_tool(name: str, arg) -> Sequence[Content]: + async def slow_tool(name: str, arg) -> Sequence[ContentBlock]: nonlocal request_count request_count += 1 diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py index c41e710cd..526201f9a 100644 --- a/tests/server/fastmcp/test_integration.py +++ b/tests/server/fastmcp/test_integration.py @@ -38,6 +38,7 @@ ProgressNotification, PromptReference, ReadResourceResult, + ResourceLink, ResourceListChangedNotification, ResourceTemplateReference, SamplingMessage, @@ -152,6 +153,25 @@ async def tool_with_progress(message: str, ctx: Context, steps: int = 3) -> str: def echo(message: str) -> str: return f"Echo: {message}" + # Tool that returns ResourceLinks + @mcp.tool(description="Lists files and returns resource links", title="List Files Tool") + def list_files() -> list[ResourceLink]: + """Returns a list of resource links for files matching the pattern.""" + + # Mock some file resources for testing + file_resources = [ + { + "type": "resource_link", + "uri": "file:///project/README.md", + "name": "README.md", + "mimeType": "text/markdown", + } + ] + + result: list[ResourceLink] = [ResourceLink.model_validate(file_json) for file_json in file_resources] + + return result + # Tool with sampling capability @mcp.tool(description="A tool that uses sampling to generate content", title="Sampling Tool") async def sampling_tool(prompt: str, ctx: Context) -> str: @@ -762,7 +782,17 @@ async def call_all_mcp_features(session: ClientSession, collector: NotificationC assert isinstance(tool_result.content[0], TextContent) assert tool_result.content[0].text == "Echo: hello" - # 2. Tool with context (logging and progress) + # 2. Test tool that returns ResourceLinks + list_files_result = await session.call_tool("list_files") + assert len(list_files_result.content) == 1 + + # Rest should be ResourceLinks + content = list_files_result.content[0] + assert isinstance(content, ResourceLink) + assert str(content.uri).startswith("file:///") + assert content.name is not None + assert content.mimeType is not None + # Test progress callback functionality progress_updates = [] diff --git a/tests/server/fastmcp/test_server.py b/tests/server/fastmcp/test_server.py index e73451290..8719b78d5 100644 --- a/tests/server/fastmcp/test_server.py +++ b/tests/server/fastmcp/test_server.py @@ -18,7 +18,7 @@ from mcp.types import ( AudioContent, BlobResourceContents, - Content, + ContentBlock, EmbeddedResource, ImageContent, TextContent, @@ -194,7 +194,7 @@ def image_tool_fn(path: str) -> Image: return Image(path) -def mixed_content_tool_fn() -> list[Content]: +def mixed_content_tool_fn() -> list[ContentBlock]: return [ TextContent(type="text", text="Hello"), ImageContent(type="image", data="abc", mimeType="image/png"),