Fix ClearToolResults corrupting typed ToolReturnPart subclasses

`rebuild_with_cleared` matched tool results with `isinstance(part, ToolReturnPart)`,
which also caught typed subclasses like `ToolSearchReturnPart`. `dataclasses.replace`
preserves the concrete class, so blanking the content produced a `ToolSearchReturnPart`
whose structured `TypedDict` content was a plain string. Core's `parse_discovered_tools`
trusts that content shape and re-reads it on the next request, raising TypeError.

Narrow the request-side clearing predicate to an exact type check so typed subclasses
are skipped -- they carry discovery state core re-reads and reclaim little.
This commit is contained in:
David SF
2026-07-17 11:49:05 -05:00
parent e4839c654b
commit 5730fb6874
2 changed files with 39 additions and 1 deletions
+6 -1
View File
@@ -387,8 +387,13 @@ def rebuild_with_cleared(
request_parts: list[ModelRequestPart] = []
changed = False
for part in msg.parts:
# Exact type, not `isinstance`: typed `ToolReturnPart` subclasses such as
# `ToolSearchReturnPart` carry structured `TypedDict` content that core re-parses
# (and trusts to be valid) on every request -- see `parse_discovered_tools`. Blanking
# it to a string both breaks that invariant (crash next request) and discards
# discovery state, so only untyped tool results are cleared.
if (
isinstance(part, ToolReturnPart)
type(part) is ToolReturnPart
and part.tool_call_id in clear_return_ids
and str(part.content) != placeholder
):
+33
View File
@@ -16,10 +16,13 @@ from pydantic_ai.messages import (
TextPart,
ToolCallPart,
ToolReturnPart,
ToolSearchReturnContent,
ToolSearchReturnPart,
UserPromptPart,
)
from pydantic_ai.models import Model, ModelRequestContext, ModelRequestParameters
from pydantic_ai.models.test import TestModel
from pydantic_ai.toolsets._tool_search import parse_discovered_tools
from pydantic_ai.usage import RunUsage
from pydantic_ai_harness.compaction import (
@@ -1581,6 +1584,36 @@ class TestClearToolResults:
assert _return_contents(twice) == ['[tool result cleared]']
assert _call_args(twice) == ['{}']
@pytest.mark.anyio
async def test_preserves_typed_tool_search_return(self):
# `ToolSearchReturnPart` subclasses `ToolReturnPart` but carries structured content that
# core's `parse_discovered_tools` re-reads on the next request. Blanking it to a string
# crashed that reader; clearing must skip typed subclasses and touch only plain results.
cap = ClearToolResults(max_messages=1, keep_pairs=0)
search_content: ToolSearchReturnContent = {'discovered_tools': [{'name': 'alpha'}, {'name': 'beta'}]}
messages: list[ModelMessage] = [
_tool_call('fn', 'u1'),
_tool_return('fn', 'u1', 'plain result'),
ModelResponse(parts=[ToolCallPart(tool_name='search_tools', args='{}', tool_call_id='ts1')]),
ModelRequest(parts=[ToolSearchReturnPart(content=search_content, tool_call_id='ts1')]),
]
result = await cap.before_model_request(_make_ctx(), _make_request_context(messages))
returns = {
p.tool_call_id: p
for m in result.messages
if isinstance(m, ModelRequest)
for p in m.parts
if isinstance(p, ToolReturnPart)
}
# Plain result blanked, typed search return kept intact (concrete type + structured content).
assert type(returns['u1']) is ToolReturnPart
assert returns['u1'].content == cap.placeholder
assert type(returns['ts1']) is ToolSearchReturnPart
assert returns['ts1'].content == search_content
# The real next-request regression: core still recovers the discovered names.
assert parse_discovered_tools(result.messages) == {'alpha', 'beta'}
# ---------------------------------------------------------------------------
# DeduplicateFileReads