"""Tests for PatchToolCallsProcessor.""" from __future__ import annotations from typing import Any from pydantic_ai.messages import ( ModelRequest, ModelResponse, RetryPromptPart, TextPart, ToolCallPart, ToolReturnPart, UserPromptPart, ) from pydantic_ai.models.test import TestModel from pydantic_ai.models.test import TestModel as _TestModel from pydantic_ai.tools import RunContext from pydantic_ai.usage import RunUsage from pydantic_deep import ( CANCELLED_MESSAGE, PatchToolCallsCapability, create_deep_agent, patch_tool_calls_processor, ) TEST_MODEL = TestModel() class TestPatchToolCallsProcessor: def test_empty_messages(self): result = patch_tool_calls_processor([]) assert result == [] def test_no_orphans(self): """Normal history with matched tool calls passes through unchanged.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")]), ModelRequest( parts=[ ToolReturnPart(tool_name="t", content="result", tool_call_id="c1"), ] ), ModelResponse(parts=[TextPart(content="done")]), ] result = patch_tool_calls_processor(messages) assert result is messages # same object, no changes def test_no_tool_calls(self): """Messages with no tool calls pass through unchanged.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="hi")]), ] result = patch_tool_calls_processor(messages) assert result is messages def test_single_orphan_at_end(self): """Tool call at the end with no following ModelRequest.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[ToolCallPart(tool_name="execute", args={}, tool_call_id="c1")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 3 # New ModelRequest with synthetic ToolReturnPart patched = result[2] assert isinstance(patched, ModelRequest) assert len(patched.parts) == 1 assert isinstance(patched.parts[0], ToolReturnPart) assert patched.parts[0].tool_name == "execute" assert patched.parts[0].content == CANCELLED_MESSAGE assert patched.parts[0].tool_call_id == "c1" def test_orphan_with_existing_next_request(self): """Orphan when next message is a ModelRequest with other parts.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse( parts=[ ToolCallPart(tool_name="read_file", args={}, tool_call_id="c1"), ToolCallPart(tool_name="write_file", args={}, tool_call_id="c2"), ] ), ModelRequest( parts=[ ToolReturnPart(tool_name="read_file", content="data", tool_call_id="c1"), # c2 is missing! ] ), ModelResponse(parts=[TextPart(content="done")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 4 # The ModelRequest at index 2 should now have synthetic part prepended patched = result[2] assert isinstance(patched, ModelRequest) # Synthetic c2 + original c1 assert len(patched.parts) == 2 synthetic = patched.parts[0] assert isinstance(synthetic, ToolReturnPart) assert synthetic.tool_name == "write_file" assert synthetic.tool_call_id == "c2" assert synthetic.content == CANCELLED_MESSAGE # Original part preserved original = patched.parts[1] assert isinstance(original, ToolReturnPart) assert original.tool_call_id == "c1" def test_multiple_orphans_across_messages(self): """Multiple orphaned tool calls in different positions.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[ToolCallPart(tool_name="t1", args={}, tool_call_id="c1")]), # No response for c1 — next is another response ModelResponse(parts=[ToolCallPart(tool_name="t2", args={}, tool_call_id="c2")]), ModelRequest( parts=[ ToolReturnPart(tool_name="t2", content="ok", tool_call_id="c2"), ] ), ] result = patch_tool_calls_processor(messages) # c1 orphan: index 1, next message is ModelResponse (not ModelRequest) # So a new ModelRequest should be inserted after index 1 assert len(result) == 5 synthetic_req = result[2] assert isinstance(synthetic_req, ModelRequest) assert len(synthetic_req.parts) == 1 assert isinstance(synthetic_req.parts[0], ToolReturnPart) assert synthetic_req.parts[0].tool_call_id == "c1" def test_all_calls_orphaned(self): """All tool calls in a response are orphaned.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse( parts=[ ToolCallPart(tool_name="a", args={}, tool_call_id="c1"), ToolCallPart(tool_name="b", args={}, tool_call_id="c2"), ] ), ] result = patch_tool_calls_processor(messages) assert len(result) == 3 patched = result[2] assert isinstance(patched, ModelRequest) assert len(patched.parts) == 2 tool_names = {p.tool_name for p in patched.parts if isinstance(p, ToolReturnPart)} assert tool_names == {"a", "b"} def test_deterministic_order(self): """Synthetic parts are in sorted order by tool_call_id.""" messages = [ ModelRequest(parts=[UserPromptPart(content="go")]), ModelResponse( parts=[ ToolCallPart(tool_name="z", args={}, tool_call_id="zzz"), ToolCallPart(tool_name="a", args={}, tool_call_id="aaa"), ToolCallPart(tool_name="m", args={}, tool_call_id="mmm"), ] ), ] result = patch_tool_calls_processor(messages) patched = result[2] assert isinstance(patched, ModelRequest) ids = [p.tool_call_id for p in patched.parts if isinstance(p, ToolReturnPart)] assert ids == ["aaa", "mmm", "zzz"] def test_response_with_only_text_parts(self): """ModelResponse with only TextParts is not affected.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="just text")]), ] result = patch_tool_calls_processor(messages) assert result is messages def test_next_request_has_no_tool_returns(self): """Next ModelRequest exists but has no ToolReturnParts at all.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")]), ModelRequest(parts=[UserPromptPart(content="followup")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 3 # The ModelRequest at index 2 should have synthetic part prepended patched = result[2] assert isinstance(patched, ModelRequest) assert len(patched.parts) == 2 assert isinstance(patched.parts[0], ToolReturnPart) assert patched.parts[0].tool_call_id == "c1" # Original UserPromptPart preserved assert isinstance(patched.parts[1], UserPromptPart) def test_cancelled_message_constant(self): assert CANCELLED_MESSAGE == "Tool call was cancelled." def test_model_retry_not_treated_as_orphan(self): """A `RetryPromptPart` answers the `ToolCallPart` — no synthetic return must be injected. Regression for issue #79: when a tool raises `ModelRetry`, pydantic-ai adds a `RetryPromptPart` (not a `ToolReturnPart`) with the same `tool_call_id`. Previously the processor injected a synthetic `ToolReturnPart` with the same id, producing a `ModelRequest` with duplicate tool_call_ids — rejected by Bedrock with `The toolResult blocks ... contain duplicate Ids`. """ messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[ToolCallPart(tool_name="web_fetch", args={}, tool_call_id="c1")]), ModelRequest( parts=[ RetryPromptPart( tool_name="web_fetch", content="Failed to fetch: 404 Not Found", tool_call_id="c1", ), ] ), ] result = patch_tool_calls_processor(messages) # History must be unchanged — no synthetic ToolReturnPart injected assert result is messages patched = result[2] assert isinstance(patched, ModelRequest) # Exactly one part, and it's the original RetryPromptPart — no duplicate id assert len(patched.parts) == 1 assert isinstance(patched.parts[0], RetryPromptPart) assert patched.parts[0].tool_call_id == "c1" def test_mixed_retry_and_return_no_orphans(self): """Parallel tool calls: one answered by `ToolReturnPart`, one by `RetryPromptPart`.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse( parts=[ ToolCallPart(tool_name="read", args={}, tool_call_id="c1"), ToolCallPart(tool_name="fetch", args={}, tool_call_id="c2"), ] ), ModelRequest( parts=[ ToolReturnPart(tool_name="read", content="data", tool_call_id="c1"), RetryPromptPart( tool_name="fetch", content="try again", tool_call_id="c2", ), ] ), ] result = patch_tool_calls_processor(messages) assert result is messages patched = result[2] assert isinstance(patched, ModelRequest) # No synthetic parts added — exactly the two original answers assert len(patched.parts) == 2 ids = { p.tool_call_id for p in patched.parts if isinstance(p, (ToolReturnPart, RetryPromptPart)) } assert ids == {"c1", "c2"} # --- Phase 2: orphaned tool results (missing calls) --- def test_orphaned_tool_result_stripped(self): """ToolReturnPart without matching ToolCallPart in preceding response is removed.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="summary")]), # This ModelRequest has a ToolReturnPart but the preceding # ModelResponse has no ToolCallParts — orphaned result ModelRequest( parts=[ ToolReturnPart(tool_name="read_file", content="data", tool_call_id="c1"), UserPromptPart(content="next"), ] ), ModelResponse(parts=[TextPart(content="done")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 4 # The orphaned ToolReturnPart should be stripped, keeping UserPromptPart patched = result[2] assert isinstance(patched, ModelRequest) assert len(patched.parts) == 1 assert isinstance(patched.parts[0], UserPromptPart) def test_orphaned_tool_result_all_stripped(self): """ModelRequest with only orphaned ToolReturnParts is removed entirely.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="no tool calls here")]), # Only orphaned tool results — entire message should be removed ModelRequest( parts=[ ToolReturnPart(tool_name="t1", content="r1", tool_call_id="c1"), ToolReturnPart(tool_name="t2", content="r2", tool_call_id="c2"), ] ), ModelResponse(parts=[TextPart(content="done")]), ] result = patch_tool_calls_processor(messages) # The empty ModelRequest should be removed assert len(result) == 3 assert isinstance(result[0], ModelRequest) assert isinstance(result[1], ModelResponse) assert isinstance(result[2], ModelResponse) def test_orphaned_tool_result_all_stripped_at_end_keeps_placeholder(self): """A trailing request of only orphaned results survives as an empty placeholder. Dropping the final message would leave the history ending on a `ModelResponse`, tripping pydantic-ai's "Processed history must end with a `ModelRequest`" validation. The stripped request is kept as an empty `ModelRequest` — the structural placeholder pydantic-ai uses when resuming without a prompt. """ messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="no tool calls here")]), # Trailing request holding only an orphaned tool result. ModelRequest(parts=[ToolReturnPart(tool_name="t1", content="r1", tool_call_id="c1")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 3 last = result[-1] assert isinstance(last, ModelRequest) assert last.parts == [] def test_mixed_valid_and_orphaned_results(self): """Some ToolReturnParts have matching calls, some don't.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse( parts=[ ToolCallPart(tool_name="t1", args={}, tool_call_id="c1"), # No c2 call! ] ), ModelRequest( parts=[ ToolReturnPart(tool_name="t1", content="ok", tool_call_id="c1"), ToolReturnPart(tool_name="t2", content="orphaned", tool_call_id="c2"), ] ), ModelResponse(parts=[TextPart(content="done")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 4 patched = result[2] assert isinstance(patched, ModelRequest) # Only c1 should remain assert len(patched.parts) == 1 assert isinstance(patched.parts[0], ToolReturnPart) assert patched.parts[0].tool_call_id == "c1" def test_orphaned_result_at_start(self): """ToolReturnPart at the very start (no preceding message) is orphaned.""" messages = [ ModelRequest( parts=[ ToolReturnPart(tool_name="t", content="data", tool_call_id="c1"), UserPromptPart(content="hello"), ] ), ModelResponse(parts=[TextPart(content="hi")]), ] result = patch_tool_calls_processor(messages) assert len(result) == 2 patched = result[0] assert isinstance(patched, ModelRequest) assert len(patched.parts) == 1 assert isinstance(patched.parts[0], UserPromptPart) def test_both_orphaned_calls_and_results(self): """Both orphaned calls and orphaned results in the same history.""" messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), # Orphaned call (c1 has no result) ModelResponse(parts=[ToolCallPart(tool_name="t1", args={}, tool_call_id="c1")]), # Next response (no ModelRequest between them) ModelResponse(parts=[TextPart(content="text")]), # Orphaned result (c2 has no matching call in prev response) ModelRequest( parts=[ ToolReturnPart(tool_name="t2", content="orphaned", tool_call_id="c2"), UserPromptPart(content="next"), ] ), ] result = patch_tool_calls_processor(messages) # After phase 1: synthetic result for c1 injected # After phase 2: orphaned c2 stripped # Check c1 was patched has_c1_synthetic = False for msg in result: if isinstance(msg, ModelRequest): for part in msg.parts: if isinstance(part, ToolReturnPart) and part.tool_call_id == "c1": has_c1_synthetic = True assert part.content == CANCELLED_MESSAGE assert has_c1_synthetic # Check c2 was stripped for msg in result: if isinstance(msg, ModelRequest): for part in msg.parts: if isinstance(part, ToolReturnPart) and part.tool_call_id == "c2": raise AssertionError( "Orphaned c2 should have been stripped" ) # pragma: no cover class TestCreateDeepAgentPatchToolCalls: def _has_patch_capability(self, agent: object) -> bool: """Check if agent has PatchToolCallsCapability registered.""" root = getattr(agent, "_root_capability", None) if root is None: return False for cap in getattr(root, "capabilities", []): if isinstance(cap, PatchToolCallsCapability): return True return False def test_patch_tool_calls_true(self): """Agent with patch_tool_calls=True has the capability.""" agent = create_deep_agent(model=TEST_MODEL, patch_tool_calls=True, cost_tracking=False) assert self._has_patch_capability(agent) def test_patch_tool_calls_false(self): """Agent with patch_tool_calls=False doesn't have the capability.""" agent = create_deep_agent(model=TEST_MODEL, patch_tool_calls=False, cost_tracking=False) assert not self._has_patch_capability(agent) def test_patch_tool_calls_with_other_processors(self): """patch_tool_calls works alongside other history processors.""" from pydantic_deep import create_sliding_window_processor window = create_sliding_window_processor( trigger=("messages", 100), keep=("messages", 50), ) agent = create_deep_agent( model=TEST_MODEL, patch_tool_calls=True, history_processors=[window], cost_tracking=False, ) assert self._has_patch_capability(agent) class TestPatchToolCallsCapability: """Tests for PatchToolCallsCapability (before_model_request hook).""" async def test_fixes_orphaned_calls(self): """Capability patches orphaned tool calls via before_model_request.""" cap = PatchToolCallsCapability() ctx = RunContext(deps=None, model=_TestModel(), usage=RunUsage()) messages = [ ModelRequest(parts=[UserPromptPart(content="do something")]), ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")]), # No matching ToolReturnPart — orphan ] class FakeReqCtx: def __init__(self, msgs: list[Any]) -> None: self.messages = msgs rc = FakeReqCtx(messages) result = await cap.before_model_request(ctx, rc) # Should have injected a synthetic ToolReturnPart patched = result.messages assert len(patched) == 3 last = patched[-1] assert isinstance(last, ModelRequest) tool_return = [p for p in last.parts if isinstance(p, ToolReturnPart)] assert len(tool_return) == 1 assert tool_return[0].content == CANCELLED_MESSAGE async def test_passes_through_clean_history(self): """Clean history passes through unchanged.""" cap = PatchToolCallsCapability() ctx = RunContext(deps=None, model=_TestModel(), usage=RunUsage()) messages = [ ModelRequest(parts=[UserPromptPart(content="hello")]), ModelResponse(parts=[TextPart(content="hi")]), ] class FakeReqCtx: def __init__(self, msgs: list[Any]) -> None: self.messages = msgs rc = FakeReqCtx(messages) result = await cap.before_model_request(ctx, rc) assert result.messages == messages class TestPatchExports: def test_importable_from_pydantic_deep(self): from pydantic_deep import ( CANCELLED_MESSAGE, PatchToolCallsCapability, patch_tool_calls_processor, ) assert patch_tool_calls_processor is not None assert PatchToolCallsCapability is not None assert CANCELLED_MESSAGE == "Tool call was cancelled." class TestPatchPreservesRequestFieldsB3: """B3: rebuilt ModelRequests keep instructions (and other fields), not dropped.""" def test_orphaned_call_prepend_preserves_instructions(self) -> None: msgs: list[Any] = [ ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")]), ModelRequest(parts=[UserPromptPart(content="next")], instructions="SYSTEM"), ] out = patch_tool_calls_processor(msgs) reqs = [m for m in out if isinstance(m, ModelRequest)] assert any(getattr(m, "instructions", None) == "SYSTEM" for m in reqs) def test_orphaned_result_strip_preserves_instructions(self) -> None: msgs: list[Any] = [ ModelRequest( parts=[ UserPromptPart(content="hi"), ToolReturnPart(tool_name="t", content="r", tool_call_id="ghost"), ], instructions="SYSTEM2", ), ] out = patch_tool_calls_processor(msgs) assert len(out) == 1 assert isinstance(out[0], ModelRequest) assert out[0].instructions == "SYSTEM2" assert all(not isinstance(p, ToolReturnPart) for p in out[0].parts) class TestPatchBlankToolCallIdB11: """B11: a blank tool_call_id isn't treated as an orphan (consistent with phase 2).""" def test_blank_id_call_gets_no_synthetic_return(self) -> None: msgs: list[Any] = [ ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="")]), ] out = patch_tool_calls_processor(msgs) injected = [ p for m in out if isinstance(m, ModelRequest) for p in m.parts if isinstance(p, ToolReturnPart) ] assert injected == []