|
| 1 | +"""Test for base64 encoding issue in MCP server. |
| 2 | +
|
| 3 | +This test demonstrates the issue in server.py where the server uses |
| 4 | +urlsafe_b64encode but the BlobResourceContents validator expects standard |
| 5 | +base64 encoding. |
| 6 | +
|
| 7 | +The test should FAIL before fixing server.py to use b64encode instead of |
| 8 | +urlsafe_b64encode. |
| 9 | +After the fix, the test should PASS. |
| 10 | +""" |
| 11 | + |
| 12 | +import base64 |
| 13 | +from typing import cast |
| 14 | + |
| 15 | +import pytest |
| 16 | +from pydantic import AnyUrl |
| 17 | + |
| 18 | +from mcp.server.lowlevel.helper_types import ReadResourceContents |
| 19 | +from mcp.server.lowlevel.server import Server |
| 20 | +from mcp.types import ( |
| 21 | + BlobResourceContents, |
| 22 | + ReadResourceRequest, |
| 23 | + ReadResourceRequestParams, |
| 24 | + ReadResourceResult, |
| 25 | + ServerResult, |
| 26 | +) |
| 27 | + |
| 28 | + |
| 29 | +@pytest.mark.anyio |
| 30 | +async def test_server_base64_encoding_issue(): |
| 31 | + """Tests that server response can be validated by BlobResourceContents. |
| 32 | +
|
| 33 | + This test will: |
| 34 | + 1. Set up a server that returns binary data |
| 35 | + 2. Extract the base64-encoded blob from the server's response |
| 36 | + 3. Verify the encoded data can be properly validated by BlobResourceContents |
| 37 | +
|
| 38 | + BEFORE FIX: The test will fail because server uses urlsafe_b64encode |
| 39 | + AFTER FIX: The test will pass because server uses standard b64encode |
| 40 | + """ |
| 41 | + server = Server("test") |
| 42 | + |
| 43 | + # Create binary data that will definitely result in + and / characters |
| 44 | + # when encoded with standard base64 |
| 45 | + binary_data = bytes([x for x in range(255)] * 4) |
| 46 | + |
| 47 | + # Register a resource handler that returns our test data |
| 48 | + @server.read_resource() |
| 49 | + async def read_resource(uri: AnyUrl) -> list[ReadResourceContents]: |
| 50 | + return [ |
| 51 | + ReadResourceContents( |
| 52 | + content=binary_data, mime_type="application/octet-stream" |
| 53 | + ) |
| 54 | + ] |
| 55 | + |
| 56 | + # Get the handler directly from the server |
| 57 | + handler = server.request_handlers[ReadResourceRequest] |
| 58 | + |
| 59 | + # Create a request |
| 60 | + request = ReadResourceRequest( |
| 61 | + method="resources/read", |
| 62 | + params=ReadResourceRequestParams(uri=AnyUrl("test://resource")), |
| 63 | + ) |
| 64 | + |
| 65 | + # Call the handler to get the response |
| 66 | + result: ServerResult = await handler(request) |
| 67 | + |
| 68 | + # After (fixed code): |
| 69 | + read_result: ReadResourceResult = cast(ReadResourceResult, result.root) |
| 70 | + blob_content = read_result.contents[0] |
| 71 | + |
| 72 | + # First verify our test data actually produces different encodings |
| 73 | + urlsafe_b64 = base64.urlsafe_b64encode(binary_data).decode() |
| 74 | + standard_b64 = base64.b64encode(binary_data).decode() |
| 75 | + assert urlsafe_b64 != standard_b64, "Test data doesn't demonstrate" |
| 76 | + " encoding difference" |
| 77 | + |
| 78 | + # Now validate the server's output with BlobResourceContents.model_validate |
| 79 | + # Before the fix: This should fail with "Invalid base64" because server |
| 80 | + # uses urlsafe_b64encode |
| 81 | + # After the fix: This should pass because server will use standard b64encode |
| 82 | + model_dict = blob_content.model_dump() |
| 83 | + |
| 84 | + # Direct validation - this will fail before fix, pass after fix |
| 85 | + blob_model = BlobResourceContents.model_validate(model_dict) |
| 86 | + |
| 87 | + # Verify we can decode the data back correctly |
| 88 | + decoded = base64.b64decode(blob_model.blob) |
| 89 | + assert decoded == binary_data |
0 commit comments