Skip to content

Commit 5614482

Browse files
Atishay Tiwariclaude
authored andcommitted
fix: invalidate tool schema cache on ToolListChangedNotification and paginate all pages
Two related bugs in ClientSession._validate_tool_result: 1. ToolListChangedNotification was not handled in _received_notification, leaving _tool_output_schemas stale after the server updates its tool list. Subsequent call_tool invocations would validate structured content against the old schema, causing false negatives (valid responses rejected) or silently accepting responses that no longer match the current schema. Fix: clear _tool_output_schemas when ToolListChangedNotification arrives. 2. When a tool was not found in the schema cache, only the first page of list_tools() was fetched. Tools on page 2+ were never cached, so their output schema was silently skipped — invalid structured_content was accepted without raising a RuntimeError. Fix: paginate through all pages until the tool is found or the list ends. Tests added for both cases in test_output_schema_validation.py. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 5cbd259 commit 5614482

File tree

2 files changed

+125
-2
lines changed

2 files changed

+125
-2
lines changed

src/mcp/client/session.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -324,8 +324,13 @@ async def call_tool(
324324
async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
325325
"""Validate the structured content of a tool result against its output schema."""
326326
if name not in self._tool_output_schemas:
327-
# refresh output schema cache
328-
await self.list_tools()
327+
# refresh output schema cache — paginate through all pages so tools
328+
# beyond the first page are also considered before giving up.
329+
list_result = await self.list_tools()
330+
while list_result.next_cursor is not None and name not in self._tool_output_schemas:
331+
list_result = await self.list_tools(
332+
params=types.PaginatedRequestParams(cursor=list_result.next_cursor)
333+
)
329334

330335
output_schema = None
331336
if name in self._tool_output_schemas:
@@ -476,5 +481,9 @@ async def _received_notification(self, notification: types.ServerNotification) -
476481
# Clients MAY use this to retry requests or update UI
477482
# The notification contains the elicitationId of the completed elicitation
478483
pass
484+
case types.ToolListChangedNotification():
485+
# The server's tool list has changed; invalidate the cached output schemas
486+
# so the next call_tool fetches fresh schemas before validating.
487+
self._tool_output_schemas.clear()
479488
case _:
480489
pass

tests/client/test_output_schema_validation.py

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
PaginatedRequestParams,
1313
TextContent,
1414
Tool,
15+
ToolListChangedNotification,
1516
)
1617

1718

@@ -163,3 +164,116 @@ async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams)
163164
assert result.is_error is False
164165

165166
assert "Tool mystery_tool not listed" in caplog.text
167+
168+
169+
@pytest.mark.anyio
170+
async def test_tool_list_changed_notification_clears_schema_cache():
171+
"""ToolListChangedNotification must invalidate the cached output schemas.
172+
173+
Flow:
174+
Call 1 — schema v1 (integer). Client caches v1. Result validates OK.
175+
Call 2 — server switches to v2 (string), sends ToolListChangedNotification
176+
*before* returning the result, then returns a string value.
177+
178+
Without the fix the client keeps the stale v1 schema and validates the
179+
string result against it → RuntimeError (false negative).
180+
With the fix the notification clears the cache, list_tools() re-fetches v2,
181+
and the string result validates correctly → no error.
182+
"""
183+
schema_v1 = {
184+
"type": "object",
185+
"properties": {"result": {"type": "integer"}},
186+
"required": ["result"],
187+
}
188+
schema_v2 = {
189+
"type": "object",
190+
"properties": {"result": {"type": "string"}},
191+
"required": ["result"],
192+
}
193+
194+
use_v2: list[bool] = [False] # mutable container so nested functions can write to it
195+
196+
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
197+
schema = schema_v2 if use_v2[0] else schema_v1
198+
return ListToolsResult(
199+
tools=[Tool(name="dynamic_tool", description="d", input_schema={"type": "object"}, output_schema=schema)]
200+
)
201+
202+
call_count: list[int] = [0]
203+
204+
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
205+
call_count[0] += 1
206+
if call_count[0] == 1:
207+
# First call: v1 schema, no notification, integer result.
208+
return CallToolResult(
209+
content=[TextContent(type="text", text="r")],
210+
structured_content={"result": 42}, # valid for v1 (integer)
211+
)
212+
# Second call: switch schema to v2, notify BEFORE returning the result,
213+
# then return a string value that is valid only under v2.
214+
use_v2[0] = True
215+
await ctx.session.send_notification(ToolListChangedNotification())
216+
return CallToolResult(
217+
content=[TextContent(type="text", text="r")],
218+
structured_content={"result": "hello"}, # valid for v2 (string), invalid for v1
219+
)
220+
221+
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
222+
223+
async with Client(server) as client:
224+
# Call 1: populates the cache with v1 schema and succeeds.
225+
result1 = await client.call_tool("dynamic_tool", {})
226+
assert result1.structured_content == {"result": 42}
227+
228+
# Call 2: notification arrives first → (with fix) cache cleared → list_tools()
229+
# fetches v2 → string "hello" is valid → no error.
230+
# Without the fix: stale v1 still in cache → "hello" fails integer check → RuntimeError.
231+
result2 = await client.call_tool("dynamic_tool", {})
232+
assert result2.structured_content == {"result": "hello"}
233+
234+
235+
@pytest.mark.anyio
236+
async def test_validate_tool_result_paginates_all_pages():
237+
"""_validate_tool_result must paginate through all tool pages when refreshing.
238+
239+
Without the fix, only the first page of list_tools() is fetched. A tool that
240+
sits on a later page is never found in the cache, so its output schema is
241+
silently skipped — invalid structured_content is accepted without error.
242+
"""
243+
output_schema = {
244+
"type": "object",
245+
"properties": {"result": {"type": "integer"}},
246+
"required": ["result"],
247+
}
248+
249+
page1_tools = [
250+
Tool(name=f"tool_{i}", description="d", input_schema={"type": "object"}) for i in range(3)
251+
]
252+
page2_tools = [
253+
Tool(
254+
name="paginated_tool",
255+
description="d",
256+
input_schema={"type": "object"},
257+
output_schema=output_schema,
258+
)
259+
]
260+
261+
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
262+
if params is not None and params.cursor == "page2":
263+
return ListToolsResult(tools=page2_tools, next_cursor=None)
264+
return ListToolsResult(tools=page1_tools, next_cursor="page2")
265+
266+
async def on_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
267+
# Returns a string for "result" — invalid per the integer schema.
268+
return CallToolResult(
269+
content=[TextContent(type="text", text="r")],
270+
structured_content={"result": "not_an_integer"},
271+
)
272+
273+
server = Server("test-server", on_list_tools=on_list_tools, on_call_tool=on_call_tool)
274+
275+
async with Client(server) as client:
276+
# With the fix: both pages are fetched, schema is found, invalid content raises.
277+
# Without the fix: only page 1 is fetched, tool not found, validation silently skipped.
278+
with pytest.raises(RuntimeError, match="Invalid structured content returned by tool paginated_tool"):
279+
await client.call_tool("paginated_tool", {})

0 commit comments

Comments
 (0)