mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 06:54:18 +01:00
90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
"""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(list(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)
|
|
|
|
# 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
|