Cache compiled output-schema validators on ClientSession (#3134)
Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
parent
11934c90ae
commit
f599cdfcf9
2 changed files with 134 additions and 9 deletions
|
|
@ -657,12 +657,16 @@ async def test_a_complete_listing_prunes_per_tool_state_for_tools_it_no_longer_c
|
|||
with anyio.fail_after(5):
|
||||
async with Client(server) as client:
|
||||
await client.session.list_tools()
|
||||
# Compile the retired tool's output-schema validator so its eviction is observable.
|
||||
await client.session.validate_tool_result("retired", CallToolResult(content=[], structured_content={}))
|
||||
assert set(client.session._x_mcp_header_maps) == {"retired", "survivor"}
|
||||
assert set(client.session._tool_output_schemas) == {"retired", "survivor"}
|
||||
assert set(client.session._tool_output_validators) == {"retired"}
|
||||
|
||||
await client.session.list_tools()
|
||||
assert set(client.session._x_mcp_header_maps) == {"survivor"}
|
||||
assert set(client.session._tool_output_schemas) == {"survivor"}
|
||||
assert client.session._tool_output_validators == {}
|
||||
|
||||
|
||||
async def test_a_complete_listing_prunes_output_schemas_on_a_legacy_session_too() -> None:
|
||||
|
|
|
|||
|
|
@ -64,3 +64,64 @@ async def test_validate_tool_result_raises_on_schema_mismatch() -> None:
|
|||
# Stable SDK prefix only: the message tail is jsonschema text that shifts with the dependency.
|
||||
with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
|
||||
await client.session.validate_tool_result("t", CallToolResult(content=[], structured_content={"x": "no"}))
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_tool_result_raises_on_an_unusable_output_schema() -> None:
|
||||
"""A schema that isn't valid JSON Schema is reported as such, on every call."""
|
||||
server = _make_server({"type": "not-a-json-schema-type"})
|
||||
async with Client(server) as client:
|
||||
result = CallToolResult(content=[], structured_content={"x": 1})
|
||||
for _ in range(2):
|
||||
# Compiling is never cached on failure, so the second call raises like the first.
|
||||
# Stable SDK prefix only: the message tail is jsonschema text that shifts with the dependency.
|
||||
with pytest.raises(RuntimeError, match="Invalid schema for tool t"):
|
||||
await client.session.validate_tool_result("t", result)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_tool_result_compiles_the_output_schema_once_per_tool() -> None:
|
||||
"""Regression guard: compiling dominates validating, so the validator must outlive one call."""
|
||||
server = _make_server({"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]})
|
||||
async with Client(server) as client:
|
||||
result = CallToolResult(content=[], structured_content={"x": 1})
|
||||
await client.session.validate_tool_result("t", result)
|
||||
compiled = client.session._tool_output_validators["t"]
|
||||
await client.session.validate_tool_result("t", result)
|
||||
assert client.session._tool_output_validators["t"] is compiled
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_tool_result_keeps_the_validator_across_a_relisting_of_the_same_schema() -> None:
|
||||
"""SDK-defined: an unchanged schema on a re-listing (or a re-absorbed cache hit) keeps its
|
||||
compiled validator, so relisting between calls doesn't reintroduce the per-call compile."""
|
||||
server = _make_server({"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]})
|
||||
async with Client(server) as client:
|
||||
result = CallToolResult(content=[], structured_content={"x": 1})
|
||||
await client.session.validate_tool_result("t", result)
|
||||
compiled = client.session._tool_output_validators["t"]
|
||||
|
||||
await client.session.list_tools() # a second listing carrying an equal schema
|
||||
await client.session.validate_tool_result("t", result)
|
||||
assert client.session._tool_output_validators["t"] is compiled
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_validate_tool_result_recompiles_when_the_server_changes_the_schema() -> None:
|
||||
"""A relisted tool must not be validated against the schema it used to declare."""
|
||||
schemas = [
|
||||
{"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]},
|
||||
{"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]},
|
||||
]
|
||||
|
||||
async def on_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams | None) -> ListToolsResult:
|
||||
return ListToolsResult(tools=[Tool(name="t", input_schema={"type": "object"}, output_schema=schemas.pop(0))])
|
||||
|
||||
server = Server("test-server", on_list_tools=on_list_tools)
|
||||
async with Client(server) as client:
|
||||
integer_result = CallToolResult(content=[], structured_content={"x": 1})
|
||||
await client.session.validate_tool_result("t", integer_result)
|
||||
|
||||
await client.session.list_tools()
|
||||
with pytest.raises(RuntimeError, match="Invalid structured content returned by tool t"):
|
||||
await client.session.validate_tool_result("t", integer_result)
|
||||
|
|
|
|||
Loading…
Reference in a new issue