From cf1b7ccd4cc9eed00e41e922fe99657b1944d9b1 Mon Sep 17 00:00:00 2001 From: evalstate <1936278+evalstate@users.noreply.github.com> Date: Fri, 21 Mar 2025 23:04:15 +0000 Subject: [PATCH 1/2] basic fix (pre-pyright and ruff) --- src/mcp/server/lowlevel/server.py | 2 +- tests/issues/test_342_base64_encoding.py | 87 ++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 tests/issues/test_342_base64_encoding.py diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index e14f73e19..dbaff3051 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -301,7 +301,7 @@ def create_content(data: str | bytes, mime_type: str | None): return types.BlobResourceContents( uri=req.params.uri, - blob=base64.urlsafe_b64encode(data).decode(), + blob=base64.b64encode(data).decode(), mimeType=mime_type or "application/octet-stream", ) diff --git a/tests/issues/test_342_base64_encoding.py b/tests/issues/test_342_base64_encoding.py new file mode 100644 index 000000000..00350cbcd --- /dev/null +++ b/tests/issues/test_342_base64_encoding.py @@ -0,0 +1,87 @@ +"""Test for base64 encoding issue in MCP server. + +This test demonstrates the issue in server.py where the server uses +urlsafe_b64encode but the BlobResourceContents validator expects standard +base64 encoding. + +The test should FAIL before fixing server.py to use b64encode instead of +urlsafe_b64encode. +After the fix, the test should PASS. +""" + +import base64 +from typing import cast + +import pytest +from pydantic import AnyUrl + +from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.server.lowlevel.server import Server +from mcp.types import ( + BlobResourceContents, + ReadResourceRequest, + ReadResourceRequestParams, + ReadResourceResult, + ServerResult, +) + + +@pytest.mark.anyio +async def test_server_base64_encoding_issue(): + """Tests that server response can be validated by BlobResourceContents. + + This test will: + 1. Set up a server that returns binary data + 2. Extract the base64-encoded blob from the server's response + 3. Verify the encoded data can be properly validated by BlobResourceContents + + BEFORE FIX: The test will fail because server uses urlsafe_b64encode + AFTER FIX: The test will pass because server uses standard b64encode + """ + server = Server("test") + + # Create binary data that will definitely result in + and / characters + # when encoded with standard base64 + binary_data = bytes([x for x in range(255)] * 4) + + # Register a resource handler that returns our test data + @server.read_resource() + async def read_resource(uri: AnyUrl) -> list[ReadResourceContents]: + return [ReadResourceContents(content=binary_data, + mime_type="application/octet-stream")] + + # Get the handler directly from the server + handler = server.request_handlers[ReadResourceRequest] + + # Create a request + request = ReadResourceRequest( + method="resources/read", + params=ReadResourceRequestParams(uri=AnyUrl("test://resource")), + ) + + # Call the handler to get the response + result: ServerResult = await handler(request) + + # Type cast and access the contents + read_result : ReadResourceResult= cast(ReadResourceResult, result) + blob_content = read_result.root.contents[0] + + + # First verify our test data actually produces different encodings + urlsafe_b64 = base64.urlsafe_b64encode(binary_data).decode() + standard_b64 = base64.b64encode(binary_data).decode() + assert urlsafe_b64 != standard_b64, "Test data doesn't demonstrate" + " encoding difference" + + # Now validate the server's output with BlobResourceContents.model_validate + # Before the fix: This should fail with "Invalid base64" because server + # uses urlsafe_b64encode + # After the fix: This should pass because server will use standard b64encode + model_dict = blob_content.model_dump() + + # Direct validation - this will fail before fix, pass after fix + blob_model = BlobResourceContents.model_validate(model_dict) + + # Verify we can decode the data back correctly + decoded = base64.b64decode(blob_model.blob) + assert decoded == binary_data \ No newline at end of file From c403341b76047464199661bd954a9d11ae0d3b36 Mon Sep 17 00:00:00 2001 From: evalstate <1936278+evalstate@users.noreply.github.com> Date: Fri, 21 Mar 2025 23:09:24 +0000 Subject: [PATCH 2/2] ruff and pyright --- src/mcp/server/fastmcp/server.py | 6 +-- tests/issues/test_342_base64_encoding.py | 48 ++++++++++++------------ tests/shared/test_sse.py | 4 +- tests/shared/test_ws.py | 4 +- 4 files changed, 30 insertions(+), 32 deletions(-) diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 275bcb36c..0349506f5 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -652,9 +652,9 @@ async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContent Returns: The resource content as either text or bytes """ - assert self._fastmcp is not None, ( - "Context is not available outside of a request" - ) + assert ( + self._fastmcp is not None + ), "Context is not available outside of a request" return await self._fastmcp.read_resource(uri) async def log( diff --git a/tests/issues/test_342_base64_encoding.py b/tests/issues/test_342_base64_encoding.py index 00350cbcd..f92b037db 100644 --- a/tests/issues/test_342_base64_encoding.py +++ b/tests/issues/test_342_base64_encoding.py @@ -1,10 +1,10 @@ """Test for base64 encoding issue in MCP server. -This test demonstrates the issue in server.py where the server uses -urlsafe_b64encode but the BlobResourceContents validator expects standard +This test demonstrates the issue in server.py where the server uses +urlsafe_b64encode but the BlobResourceContents validator expects standard base64 encoding. -The test should FAIL before fixing server.py to use b64encode instead of +The test should FAIL before fixing server.py to use b64encode instead of urlsafe_b64encode. After the fix, the test should PASS. """ @@ -29,59 +29,61 @@ @pytest.mark.anyio async def test_server_base64_encoding_issue(): """Tests that server response can be validated by BlobResourceContents. - + This test will: 1. Set up a server that returns binary data 2. Extract the base64-encoded blob from the server's response 3. Verify the encoded data can be properly validated by BlobResourceContents - + BEFORE FIX: The test will fail because server uses urlsafe_b64encode AFTER FIX: The test will pass because server uses standard b64encode """ server = Server("test") - - # Create binary data that will definitely result in + and / characters + + # Create binary data that will definitely result in + and / characters # when encoded with standard base64 binary_data = bytes([x for x in range(255)] * 4) - + # Register a resource handler that returns our test data @server.read_resource() async def read_resource(uri: AnyUrl) -> list[ReadResourceContents]: - return [ReadResourceContents(content=binary_data, - mime_type="application/octet-stream")] - + return [ + ReadResourceContents( + content=binary_data, mime_type="application/octet-stream" + ) + ] + # Get the handler directly from the server handler = server.request_handlers[ReadResourceRequest] - + # Create a request request = ReadResourceRequest( method="resources/read", params=ReadResourceRequestParams(uri=AnyUrl("test://resource")), ) - + # Call the handler to get the response result: ServerResult = await handler(request) - - # Type cast and access the contents - read_result : ReadResourceResult= cast(ReadResourceResult, result) - blob_content = read_result.root.contents[0] - - + + # After (fixed code): + read_result: ReadResourceResult = cast(ReadResourceResult, result.root) + blob_content = read_result.contents[0] + # First verify our test data actually produces different encodings urlsafe_b64 = base64.urlsafe_b64encode(binary_data).decode() standard_b64 = base64.b64encode(binary_data).decode() assert urlsafe_b64 != standard_b64, "Test data doesn't demonstrate" " encoding difference" - + # Now validate the server's output with BlobResourceContents.model_validate # Before the fix: This should fail with "Invalid base64" because server # uses urlsafe_b64encode # After the fix: This should pass because server will use standard b64encode model_dict = blob_content.model_dump() - + # Direct validation - this will fail before fix, pass after fix blob_model = BlobResourceContents.model_validate(model_dict) - + # Verify we can decode the data back correctly decoded = base64.b64decode(blob_model.blob) - assert decoded == binary_data \ No newline at end of file + assert decoded == binary_data diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 43107b597..f5158c3c3 100644 --- a/tests/shared/test_sse.py +++ b/tests/shared/test_sse.py @@ -138,9 +138,7 @@ def server(server_port: int) -> Generator[None, None, None]: time.sleep(0.1) attempt += 1 else: - raise RuntimeError( - f"Server failed to start after {max_attempts} attempts" - ) + raise RuntimeError(f"Server failed to start after {max_attempts} attempts") yield diff --git a/tests/shared/test_ws.py b/tests/shared/test_ws.py index 2aca97e15..1381c8153 100644 --- a/tests/shared/test_ws.py +++ b/tests/shared/test_ws.py @@ -134,9 +134,7 @@ def server(server_port: int) -> Generator[None, None, None]: time.sleep(0.1) attempt += 1 else: - raise RuntimeError( - f"Server failed to start after {max_attempts} attempts" - ) + raise RuntimeError(f"Server failed to start after {max_attempts} attempts") yield