Skip to content

Commit 26c6806

Browse files
committed
validation options
1 parent ef4e167 commit 26c6806

File tree

3 files changed

+281
-8
lines changed

3 files changed

+281
-8
lines changed

src/mcp/client/session.py

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import anyio.lowlevel
66
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
77
from jsonschema import SchemaError, ValidationError, validate
8-
from pydantic import AnyUrl, TypeAdapter
8+
from pydantic import AnyUrl, BaseModel, Field, TypeAdapter
99

1010
import mcp.types as types
1111
from mcp.shared.context import RequestContext
@@ -18,6 +18,17 @@
1818
logger = logging.getLogger("client")
1919

2020

21+
class ValidationOptions(BaseModel):
22+
"""Options for controlling validation behavior in MCP client sessions."""
23+
24+
strict_output_validation: bool = Field(
25+
default=True,
26+
description="Whether to raise exceptions when tools don't return structured "
27+
"content as specified by their output schema. When False, validation "
28+
"errors are logged as warnings and execution continues.",
29+
)
30+
31+
2132
class SamplingFnT(Protocol):
2233
async def __call__(
2334
self,
@@ -118,6 +129,7 @@ def __init__(
118129
logging_callback: LoggingFnT | None = None,
119130
message_handler: MessageHandlerFnT | None = None,
120131
client_info: types.Implementation | None = None,
132+
validation_options: ValidationOptions | None = None,
121133
) -> None:
122134
super().__init__(
123135
read_stream,
@@ -133,6 +145,7 @@ def __init__(
133145
self._logging_callback = logging_callback or _default_logging_callback
134146
self._message_handler = message_handler or _default_message_handler
135147
self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
148+
self._validation_options = validation_options or ValidationOptions()
136149

137150
async def initialize(self) -> types.InitializeResult:
138151
sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -324,13 +337,27 @@ async def _validate_tool_result(self, name: str, result: types.CallToolResult) -
324337

325338
if output_schema is not None:
326339
if result.structuredContent is None:
327-
raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
328-
try:
329-
validate(result.structuredContent, output_schema)
330-
except ValidationError as e:
331-
raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}")
332-
except SchemaError as e:
333-
raise RuntimeError(f"Invalid schema for tool {name}: {e}")
340+
if self._validation_options.strict_output_validation:
341+
raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
342+
else:
343+
logger.warning(
344+
f"Tool {name} has an output schema but did not return structured content. "
345+
f"Continuing without structured content validation due to lenient validation mode."
346+
)
347+
else:
348+
try:
349+
validate(result.structuredContent, output_schema)
350+
except ValidationError as e:
351+
if self._validation_options.strict_output_validation:
352+
raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}")
353+
else:
354+
logger.warning(
355+
f"Invalid structured content returned by tool {name}: {e}. "
356+
f"Continuing due to lenient validation mode."
357+
)
358+
except SchemaError as e:
359+
# Schema errors are always raised - they indicate a problem with the schema itself
360+
raise RuntimeError(f"Invalid schema for tool {name}: {e}")
334361

335362
async def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult:
336363
"""Send a prompts/list request."""

src/mcp/shared/memory.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ async def create_connected_server_and_client_session(
6161
client_info: types.Implementation | None = None,
6262
raise_exceptions: bool = False,
6363
elicitation_callback: ElicitationFnT | None = None,
64+
validation_options: Any | None = None,
6465
) -> AsyncGenerator[ClientSession, None]:
6566
"""Creates a ClientSession that is connected to a running MCP server."""
6667
async with create_client_server_memory_streams() as (
@@ -92,6 +93,7 @@ async def create_connected_server_and_client_session(
9293
message_handler=message_handler,
9394
client_info=client_info,
9495
elicitation_callback=elicitation_callback,
96+
validation_options=validation_options,
9597
) as client_session:
9698
await client_session.initialize()
9799
yield client_session
Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
"""Tests for client-side validation options."""
2+
3+
import logging
4+
from unittest.mock import AsyncMock, MagicMock
5+
6+
import pytest
7+
8+
from mcp.client.session import ClientSession, ValidationOptions
9+
from mcp.types import CallToolResult, TextContent
10+
11+
12+
class TestValidationOptions:
13+
"""Test validation options for MCP client sessions."""
14+
15+
@pytest.mark.anyio
16+
async def test_strict_validation_default(self):
17+
"""Test that strict validation is enabled by default."""
18+
# Create a mock client session
19+
read_stream = MagicMock()
20+
write_stream = MagicMock()
21+
22+
client = ClientSession(read_stream, write_stream)
23+
24+
# Set up tool with output schema
25+
client._tool_output_schemas = {
26+
"test_tool": {
27+
"type": "object",
28+
"properties": {"result": {"type": "integer"}},
29+
"required": ["result"],
30+
}
31+
}
32+
33+
# Mock send_request to return a result without structured content
34+
mock_result = CallToolResult(
35+
content=[TextContent(type="text", text="This is unstructured text content")],
36+
structuredContent=None,
37+
isError=False,
38+
)
39+
40+
client.send_request = AsyncMock(return_value=mock_result)
41+
42+
# Should raise by default when structured content is missing
43+
with pytest.raises(RuntimeError) as exc_info:
44+
await client.call_tool("test_tool", {})
45+
assert "has an output schema but did not return structured content" in str(exc_info.value)
46+
47+
@pytest.mark.anyio
48+
async def test_lenient_validation_missing_content(self, caplog):
49+
"""Test lenient validation when structured content is missing."""
50+
# Set logging level to capture warnings
51+
caplog.set_level(logging.WARNING)
52+
53+
# Create client with lenient validation
54+
validation_options = ValidationOptions(strict_output_validation=False)
55+
56+
read_stream = MagicMock()
57+
write_stream = MagicMock()
58+
59+
client = ClientSession(read_stream, write_stream, validation_options=validation_options)
60+
61+
# Set up tool with output schema
62+
client._tool_output_schemas = {
63+
"test_tool": {
64+
"type": "object",
65+
"properties": {"result": {"type": "integer"}},
66+
"required": ["result"],
67+
}
68+
}
69+
70+
# Mock send_request to return a result without structured content
71+
mock_result = CallToolResult(
72+
content=[TextContent(type="text", text="This is unstructured text content")],
73+
structuredContent=None,
74+
isError=False,
75+
)
76+
77+
client.send_request = AsyncMock(return_value=mock_result)
78+
79+
# Should not raise with lenient validation
80+
result = await client.call_tool("test_tool", {})
81+
82+
# Should have logged a warning
83+
assert "has an output schema but did not return structured content" in caplog.text
84+
assert "Continuing without structured content validation" in caplog.text
85+
86+
# Result should still be returned
87+
assert result.isError is False
88+
assert result.structuredContent is None
89+
90+
@pytest.mark.anyio
91+
async def test_lenient_validation_invalid_content(self, caplog):
92+
"""Test lenient validation when structured content is invalid."""
93+
# Set logging level to capture warnings
94+
caplog.set_level(logging.WARNING)
95+
96+
# Create client with lenient validation
97+
validation_options = ValidationOptions(strict_output_validation=False)
98+
99+
read_stream = MagicMock()
100+
write_stream = MagicMock()
101+
102+
client = ClientSession(read_stream, write_stream, validation_options=validation_options)
103+
104+
# Set up tool with output schema
105+
client._tool_output_schemas = {
106+
"test_tool": {
107+
"type": "object",
108+
"properties": {"result": {"type": "integer"}},
109+
"required": ["result"],
110+
}
111+
}
112+
113+
# Mock send_request to return a result with invalid structured content
114+
mock_result = CallToolResult(
115+
content=[TextContent(type="text", text='{"result": "not_an_integer"}')],
116+
structuredContent={"result": "not_an_integer"}, # Invalid: string instead of integer
117+
isError=False,
118+
)
119+
120+
client.send_request = AsyncMock(return_value=mock_result)
121+
122+
# Should not raise with lenient validation
123+
result = await client.call_tool("test_tool", {})
124+
125+
# Should have logged a warning
126+
assert "Invalid structured content returned by tool test_tool" in caplog.text
127+
assert "Continuing due to lenient validation mode" in caplog.text
128+
129+
# Result should still be returned with the invalid content
130+
assert result.isError is False
131+
assert result.structuredContent == {"result": "not_an_integer"}
132+
133+
@pytest.mark.anyio
134+
async def test_strict_validation_with_valid_content(self):
135+
"""Test that valid structured content passes validation."""
136+
read_stream = MagicMock()
137+
write_stream = MagicMock()
138+
139+
client = ClientSession(read_stream, write_stream)
140+
141+
# Set up tool with output schema
142+
client._tool_output_schemas = {
143+
"test_tool": {
144+
"type": "object",
145+
"properties": {"result": {"type": "integer"}},
146+
"required": ["result"],
147+
}
148+
}
149+
150+
# Mock send_request to return a result with valid structured content
151+
mock_result = CallToolResult(
152+
content=[TextContent(type="text", text='{"result": 42}')], structuredContent={"result": 42}, isError=False
153+
)
154+
155+
client.send_request = AsyncMock(return_value=mock_result)
156+
157+
# Should not raise with valid content
158+
result = await client.call_tool("test_tool", {})
159+
assert result.isError is False
160+
assert result.structuredContent == {"result": 42}
161+
162+
@pytest.mark.anyio
163+
async def test_schema_errors_always_raised(self):
164+
"""Test that schema errors are always raised regardless of validation mode."""
165+
# Create client with lenient validation
166+
validation_options = ValidationOptions(strict_output_validation=False)
167+
168+
read_stream = MagicMock()
169+
write_stream = MagicMock()
170+
171+
client = ClientSession(read_stream, write_stream, validation_options=validation_options)
172+
173+
# Set up tool with invalid output schema
174+
client._tool_output_schemas = {
175+
"test_tool": "not a valid schema" # Invalid schema
176+
}
177+
178+
# Mock send_request to return a result with structured content
179+
mock_result = CallToolResult(
180+
content=[TextContent(type="text", text='{"result": 42}')], structuredContent={"result": 42}, isError=False
181+
)
182+
183+
client.send_request = AsyncMock(return_value=mock_result)
184+
185+
# Should still raise for schema errors even in lenient mode
186+
with pytest.raises(RuntimeError) as exc_info:
187+
await client.call_tool("test_tool", {})
188+
assert "Invalid schema for tool test_tool" in str(exc_info.value)
189+
190+
@pytest.mark.anyio
191+
async def test_error_results_not_validated(self):
192+
"""Test that error results are not validated."""
193+
read_stream = MagicMock()
194+
write_stream = MagicMock()
195+
196+
client = ClientSession(read_stream, write_stream)
197+
198+
# Set up tool with output schema
199+
client._tool_output_schemas = {
200+
"test_tool": {
201+
"type": "object",
202+
"properties": {"result": {"type": "integer"}},
203+
"required": ["result"],
204+
}
205+
}
206+
207+
# Mock send_request to return an error result
208+
mock_result = CallToolResult(
209+
content=[TextContent(type="text", text="Tool execution failed")],
210+
structuredContent=None,
211+
isError=True, # Error result
212+
)
213+
214+
client.send_request = AsyncMock(return_value=mock_result)
215+
216+
# Should not validate error results
217+
result = await client.call_tool("test_tool", {})
218+
assert result.isError is True
219+
# No exception should be raised
220+
221+
@pytest.mark.anyio
222+
async def test_tool_without_output_schema(self):
223+
"""Test that tools without output schema don't trigger validation."""
224+
read_stream = MagicMock()
225+
write_stream = MagicMock()
226+
227+
client = ClientSession(read_stream, write_stream)
228+
229+
# Tool has no output schema
230+
client._tool_output_schemas = {"test_tool": None}
231+
232+
# Mock send_request to return a result without structured content
233+
mock_result = CallToolResult(
234+
content=[TextContent(type="text", text="This is unstructured text content")],
235+
structuredContent=None,
236+
isError=False,
237+
)
238+
239+
client.send_request = AsyncMock(return_value=mock_result)
240+
241+
# Should not raise when there's no output schema
242+
result = await client.call_tool("test_tool", {})
243+
assert result.isError is False
244+
assert result.structuredContent is None

0 commit comments

Comments
 (0)