From 283f43cf1e6b2e8aa06215a1eb741bac9d72209c Mon Sep 17 00:00:00 2001 From: Amit Arora Date: Mon, 28 Apr 2025 17:35:31 +0000 Subject: [PATCH 1/4] fix for SSE URL handling when a server name is specified --- src/mcp/client/sse.py | 43 ++++++++++++++++++++++++++++++++++++++-- tests/client/test_sse.py | 31 +++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 2 deletions(-) create mode 100644 tests/client/test_sse.py diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index 4f6241a72..b86dea291 100644 --- a/src/mcp/client/sse.py +++ b/src/mcp/client/sse.py @@ -1,7 +1,7 @@ import logging from contextlib import asynccontextmanager from typing import Any -from urllib.parse import urljoin, urlparse +from urllib.parse import urljoin, urlparse, urlunparse import anyio import httpx @@ -18,6 +18,31 @@ def remove_request_params(url: str) -> str: return urljoin(url, urlparse(url).path) +def custom_url_join(base_url: str, endpoint: str) -> str: + """ + Custom URL join function to handle the case where the endpoint is relative + to a base URL. This function ensures that the base URL and endpoint are + combined correctly, even if the endpoint is not a full URL. + """ + # Parse the base URL + parsed_base = urlparse(base_url) + + # Get the path prefix (e.g., '/weather') + path_prefix = "/".join(parsed_base.path.split("/")[:-1]) + + # Remove any leading slash from the endpoint + clean_endpoint = endpoint.lstrip("/") + + # Create the new path by joining prefix and endpoint + new_path = f"{path_prefix}/{clean_endpoint}" + + # Create a new parsed URL with the updated path + parsed_new = parsed_base._replace(path=new_path) + + # Convert back to a string URL + return urlunparse(parsed_new) + + @asynccontextmanager async def sse_client( url: str, @@ -43,6 +68,15 @@ async def sse_client( async with anyio.create_task_group() as tg: try: logger.info(f"Connecting to SSE endpoint: {remove_request_params(url)}") + + # extract MCP server name from URL, for example: https://mcp.example.com/weather/sse + # will extract 'weather' + path_tokens = urlparse(url).path.split("/") + optional_mcp_server_name = ( + path_tokens[1:-1] if len(path_tokens) > 2 else None + ) + logger.debug(f"MCP Server name (optional): {optional_mcp_server_name}") + async with httpx.AsyncClient(headers=headers) as client: async with aconnect_sse( client, @@ -61,7 +95,12 @@ async def sse_reader( logger.debug(f"Received SSE event: {sse.event}") match sse.event: case "endpoint": - endpoint_url = urljoin(url, sse.data) + if optional_mcp_server_name: + endpoint_url = custom_url_join( + url, sse.data + ) + else: + endpoint_url = urljoin(url, sse.data) logger.info( f"Received endpoint URL: {endpoint_url}" ) diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py new file mode 100644 index 000000000..e4f33f090 --- /dev/null +++ b/tests/client/test_sse.py @@ -0,0 +1,31 @@ +import pytest +from urllib.parse import urlparse, urlunparse + +from mcp.client.sse import custom_url_join + + +@pytest.mark.parametrize( + "base_url,endpoint,expected", + [ + # Additional test cases to verify behavior with different URL structures + ( + "https://mcp.example.com/weather/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/weather/messages/?session_id=616df71373444d76bd566df4377c9629" + ), + ( + "https://mcp.example.com/weather/clarksburg/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/weather/clarksburg/messages/?session_id=616df71373444d76bd566df4377c9629" + ), + ( + "https://mcp.example.com/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/messages/?session_id=616df71373444d76bd566df4377c9629" + ), + ], +) +def test_custom_url_join(base_url, endpoint, expected): + """Test the custom_url_join function with messages endpoint and session ID.""" + result = custom_url_join(base_url, endpoint) + assert result == expected \ No newline at end of file From ec9dcc1f35e63f480f642b8df3a0ee45a6547baf Mon Sep 17 00:00:00 2001 From: Amit Arora Date: Mon, 28 Apr 2025 17:43:19 +0000 Subject: [PATCH 2/4] ruff formatting changes --- tests/client/test_sse.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py index e4f33f090..0734dc663 100644 --- a/tests/client/test_sse.py +++ b/tests/client/test_sse.py @@ -6,26 +6,26 @@ @pytest.mark.parametrize( "base_url,endpoint,expected", - [ + [ # Additional test cases to verify behavior with different URL structures ( - "https://mcp.example.com/weather/sse", - "/messages/?session_id=616df71373444d76bd566df4377c9629", - "https://mcp.example.com/weather/messages/?session_id=616df71373444d76bd566df4377c9629" + "https://mcp.example.com/weather/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/weather/messages/?session_id=616df71373444d76bd566df4377c9629", ), ( - "https://mcp.example.com/weather/clarksburg/sse", - "/messages/?session_id=616df71373444d76bd566df4377c9629", - "https://mcp.example.com/weather/clarksburg/messages/?session_id=616df71373444d76bd566df4377c9629" + "https://mcp.example.com/weather/clarksburg/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/weather/clarksburg/messages/?session_id=616df71373444d76bd566df4377c9629", ), ( - "https://mcp.example.com/sse", - "/messages/?session_id=616df71373444d76bd566df4377c9629", - "https://mcp.example.com/messages/?session_id=616df71373444d76bd566df4377c9629" + "https://mcp.example.com/sse", + "/messages/?session_id=616df71373444d76bd566df4377c9629", + "https://mcp.example.com/messages/?session_id=616df71373444d76bd566df4377c9629", ), ], ) def test_custom_url_join(base_url, endpoint, expected): """Test the custom_url_join function with messages endpoint and session ID.""" result = custom_url_join(base_url, endpoint) - assert result == expected \ No newline at end of file + assert result == expected From f7d9e8817bb39c190bb2db732b0290f8e6e7b00e Mon Sep 17 00:00:00 2001 From: Amit Arora Date: Mon, 28 Apr 2025 17:44:40 +0000 Subject: [PATCH 3/4] ruff formatting changes --- tests/client/test_sse.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py index 0734dc663..c59118c5a 100644 --- a/tests/client/test_sse.py +++ b/tests/client/test_sse.py @@ -1,6 +1,4 @@ import pytest -from urllib.parse import urlparse, urlunparse - from mcp.client.sse import custom_url_join From 4eb85218121e43a6e8b0b9c4e39cd96022586f31 Mon Sep 17 00:00:00 2001 From: Amit Arora Date: Mon, 28 Apr 2025 17:45:59 +0000 Subject: [PATCH 4/4] make ruff happy --- tests/client/test_sse.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py index c59118c5a..5b3104fda 100644 --- a/tests/client/test_sse.py +++ b/tests/client/test_sse.py @@ -1,4 +1,5 @@ import pytest + from mcp.client.sse import custom_url_join