pydantic-ai 0.4.8__tar.gz → 0.4.10__tar.gz

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai
-Version: 0.4.8
+Version: 0.4.10
 Summary: Agent Framework / shim to use Pydantic with LLMs
 Project-URL: Homepage, https://ai.pydantic.dev
 Project-URL: Source, https://github.com/pydantic/pydantic-ai
@@ -28,11 +28,11 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,mcp,mistral,openai,retries,vertexai]==0.4.8
+Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,bedrock,cli,cohere,evals,google,groq,huggingface,mcp,mistral,openai,retries,vertexai]==0.4.10
 Provides-Extra: a2a
 Requires-Dist: fasta2a>=0.4.1; extra == 'a2a'
 Provides-Extra: examples
-Requires-Dist: pydantic-ai-examples==0.4.8; extra == 'examples'
+Requires-Dist: pydantic-ai-examples==0.4.10; extra == 'examples'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.11.0; extra == 'logfire'
 Description-Content-Type: text/markdown

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/models/test_gemini.py

@@ -194,7 +194,7 @@ async def test_require_response_tool(allow_model_requests: None):
 
 async def test_json_def_replaced(allow_model_requests: None):
     class Axis(BaseModel):
-        label: str
+        label: str = Field(default='<unlabeled axis>', description='The label of the axis')
 
     class Chart(BaseModel):
         x_axis: Axis
@@ -213,8 +213,14 @@ async def test_json_def_replaced(allow_model_requests: None):
         {
             '$defs': {
                 'Axis': {
-                    'properties': {'label': {'title': 'Label', 'type': 'string'}},
-                    'required': ['label'],
+                    'properties': {
+                        'label': {
+                            'default': '<unlabeled axis>',
+                            'description': 'The label of the axis',
+                            'title': 'Label',
+                            'type': 'string',
+                        }
+                    },
                     'title': 'Axis',
                     'type': 'object',
                 },
@@ -268,17 +274,27 @@ async def test_json_def_replaced(allow_model_requests: None):
                                 'items': {
                                     'properties': {
                                         'lat': {'type': 'number'},
-                                        'lng': {'type': 'number'},
+                                        'lng': {'default': 1.1, 'type': 'number'},
                                         'chart': {
                                             'properties': {
                                                 'x_axis': {
-                                                    'properties': {'label': {'type': 'string'}},
-                                                    'required': ['label'],
+                                                    'properties': {
+                                                        'label': {
+                                                            'default': '<unlabeled axis>',
+                                                            'description': 'The label of the axis',
+                                                            'type': 'string',
+                                                        }
+                                                    },
                                                     'type': 'object',
                                                 },
                                                 'y_axis': {
-                                                    'properties': {'label': {'type': 'string'}},
-                                                    'required': ['label'],
+                                                    'properties': {
+                                                        'label': {
+                                                            'default': '<unlabeled axis>',
+                                                            'description': 'The label of the axis',
+                                                            'type': 'string',
+                                                        }
+                                                    },
                                                     'type': 'object',
                                                 },
                                             },

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/test_ag_ui.py

@@ -7,6 +7,7 @@ import contextlib
 import json
 import uuid
 from collections.abc import AsyncIterator
+from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any
 
@@ -17,7 +18,9 @@ from dirty_equals import IsStr
 from inline_snapshot import snapshot
 from pydantic import BaseModel
 
+from pydantic_ai._run_context import RunContext
 from pydantic_ai.agent import Agent
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelMessage
 from pydantic_ai.models.function import (
     AgentInfo,
@@ -27,8 +30,9 @@ from pydantic_ai.models.function import (
     DeltaToolCalls,
     FunctionModel,
 )
+from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import OutputDataT
-from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from .conftest import IsSameStr
 
@@ -180,7 +184,7 @@ def create_input(
         thread_id=thread_id,
         run_id=uuid_str(),
         messages=list(messages),
-        state=state,
+        state=dict(state) if state else {},
         context=[],
         tools=tools or [],
         forwarded_props=None,
@@ -1050,9 +1054,19 @@ async def test_tool_local_then_ag_ui() -> None:
 async def test_request_with_state() -> None:
     """Test request with state modification."""
 
+    seen_states: list[int] = []
+
+    async def store_state(
+        ctx: RunContext[StateDeps[StateInt]], tool_defs: list[ToolDefinition]
+    ) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state.value)
+        ctx.deps.state.value += 1
+        return tool_defs
+
     agent: Agent[StateDeps[StateInt], str] = Agent(
         model=FunctionModel(stream_function=simple_stream),
         deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
+        prepare_tools=store_state,
     )
     adapter = _Adapter(agent=agent)
     run_inputs = [
@@ -1074,32 +1088,101 @@ async def test_request_with_state() -> None:
                 id='msg_3',
                 content='Hello, how are you?',
             ),
+        ),
+        create_input(
+            UserMessage(
+                id='msg_4',
+                content='Hello, how are you?',
+            ),
             state=StateInt(value=42),
         ),
     ]
 
-    deps = StateDeps(StateInt())
+    deps = StateDeps(StateInt(value=0))
 
-    last_value = deps.state.value
     for run_input in run_inputs:
         events = list[dict[str, Any]]()
         async for event in adapter.run(run_input, deps=deps):
             events.append(json.loads(event.removeprefix('data: ')))
 
         assert events == simple_result()
-        assert deps.state.value == run_input.state.value if run_input.state is not None else last_value
-        last_value = deps.state.value
+    assert seen_states == snapshot(
+        [
+            41,  # run msg_1, prepare_tools call 1
+            42,  # run msg_1, prepare_tools call 2
+            0,  # run msg_2, prepare_tools call 1
+            1,  # run msg_2, prepare_tools call 2
+            0,  # run msg_3, prepare_tools call 1
+            1,  # run msg_3, prepare_tools call 2
+            42,  # run msg_4, prepare_tools call 1
+            43,  # run msg_4, prepare_tools call 2
+        ]
+    )
+
+
+async def test_request_with_state_without_handler() -> None:
+    agent = Agent(model=FunctionModel(stream_function=simple_stream))
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state=StateInt(value=41),
+    )
+
+    with pytest.raises(
+        UserError,
+        match='AG-UI state is provided but `deps` of type `NoneType` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.',
+    ):
+        async for _ in adapter.run(run_input):
+            pass
+
+
+async def test_request_with_state_with_custom_handler() -> None:
+    @dataclass
+    class CustomStateDeps:
+        state: dict[str, Any]
+
+    seen_states: list[dict[str, Any]] = []
+
+    async def store_state(ctx: RunContext[CustomStateDeps], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
+        seen_states.append(ctx.deps.state)
+        return tool_defs
+
+    agent: Agent[CustomStateDeps, str] = Agent(
+        model=FunctionModel(stream_function=simple_stream),
+        deps_type=CustomStateDeps,
+        prepare_tools=store_state,
+    )
+    adapter = _Adapter(agent=agent)
+    run_input = create_input(
+        UserMessage(
+            id='msg_1',
+            content='Hello, how are you?',
+        ),
+        state={'value': 42},
+    )
+
+    async for _ in adapter.run(run_input, deps=CustomStateDeps(state={'value': 0})):
+        pass
 
-    assert deps.state.value == 42
+    assert seen_states[-1] == {'value': 42}
 
 
 async def test_concurrent_runs() -> None:
     """Test concurrent execution of multiple runs."""
     import asyncio
 
-    agent = Agent(
-        model=FunctionModel(stream_function=simple_stream),
+    agent: Agent[StateDeps[StateInt], str] = Agent(
+        model=TestModel(),
+        deps_type=StateDeps[StateInt],  # type: ignore[reportUnknownArgumentType]
     )
+
+    @agent.tool
+    async def get_state(ctx: RunContext[StateDeps[StateInt]]) -> int:
+        return ctx.deps.state.value
+
     adapter = _Adapter(agent=agent)
     concurrent_tasks: list[asyncio.Task[list[dict[str, Any]]]] = []
 
@@ -1109,10 +1192,11 @@ async def test_concurrent_runs() -> None:
                 id=f'msg_{i}',
                 content=f'Message {i}',
             ),
+            state=StateInt(value=i),
             thread_id=f'test_thread_{i}',
         )
 
-        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input))
+        task = asyncio.create_task(collect_events_from_adapter(adapter, run_input, deps=StateDeps(StateInt())))
         concurrent_tasks.append(task)
 
     results = await asyncio.gather(*concurrent_tasks)
@@ -1121,9 +1205,23 @@ async def test_concurrent_runs() -> None:
     for i, events in enumerate(results):
         assert events == [
             {'type': 'RUN_STARTED', 'threadId': f'test_thread_{i}', 'runId': (run_id := IsSameStr())},
+            {
+                'type': 'TOOL_CALL_START',
+                'toolCallId': (tool_call_id := IsSameStr()),
+                'toolCallName': 'get_state',
+                'parentMessageId': IsStr(),
+            },
+            {'type': 'TOOL_CALL_END', 'toolCallId': tool_call_id},
+            {
+                'type': 'TOOL_CALL_RESULT',
+                'messageId': IsStr(),
+                'toolCallId': tool_call_id,
+                'content': str(i),
+                'role': 'tool',
+            },
             {'type': 'TEXT_MESSAGE_START', 'messageId': (message_id := IsSameStr()), 'role': 'assistant'},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'success '},
-            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '(no tool calls)'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': '{"get_s'},
+            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'tate":' + str(i) + '}'},
             {'type': 'TEXT_MESSAGE_END', 'messageId': message_id},
             {'type': 'RUN_FINISHED', 'threadId': f'test_thread_{i}', 'runId': run_id},
         ]

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/test_agent.py

@@ -3866,7 +3866,7 @@ def test_prepare_output_tools():
     )
 
 
-async def test_context_manager():
+async def test_explicit_context_manager():
     try:
         from pydantic_ai.mcp import MCPServerStdio
     except ImportError:  # pragma: lax no cover
@@ -3886,6 +3886,47 @@ async def test_context_manager():
             assert server2.is_running
 
 
+async def test_implicit_context_manager():
+    try:
+        from pydantic_ai.mcp import MCPServerStdio
+    except ImportError:  # pragma: lax no cover
+        pytest.skip('mcp is not installed')
+
+    server1 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    server2 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix')])
+    agent = Agent('test', toolsets=[toolset])
+
+    async with agent.iter(
+        user_prompt='Hello',
+    ):
+        assert server1.is_running
+        assert server2.is_running
+
+
+def test_parallel_mcp_calls():
+    try:
+        from pydantic_ai.mcp import MCPServerStdio
+    except ImportError:  # pragma: lax no cover
+        pytest.skip('mcp is not installed')
+
+    async def call_tools_parallel(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        if len(messages) == 1:
+            return ModelResponse(
+                parts=[
+                    ToolCallPart(tool_name='get_none'),
+                    ToolCallPart(tool_name='get_multiple_items'),
+                ]
+            )
+        else:
+            return ModelResponse(parts=[TextPart('finished')])
+
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    agent = Agent(FunctionModel(call_tools_parallel), toolsets=[server])
+    result = agent.run_sync()
+    assert result.output == snapshot('finished')
+
+
 def test_set_mcp_sampling_model():
     try:
         from pydantic_ai.mcp import MCPServerStdio

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/test_streaming.py

@@ -1108,6 +1108,42 @@ async def test_iter_stream_structured_output():
                     )
 
 
+async def test_iter_stream_output_tool_dont_hit_retry_limit():
+    class CityLocation(BaseModel):
+        city: str
+        country: str | None = None
+
+    async def text_stream(_messages: list[ModelMessage], agent_info: AgentInfo) -> AsyncIterator[DeltaToolCalls]:
+        """Stream partial JSON data that will initially fail validation."""
+        assert agent_info.output_tools is not None
+        assert len(agent_info.output_tools) == 1
+        name = agent_info.output_tools[0].name
+
+        yield {0: DeltaToolCall(name=name)}
+        yield {0: DeltaToolCall(json_args='{"c')}
+        yield {0: DeltaToolCall(json_args='ity":')}
+        yield {0: DeltaToolCall(json_args=' "Mex')}
+        yield {0: DeltaToolCall(json_args='ico City",')}
+        yield {0: DeltaToolCall(json_args=' "cou')}
+        yield {0: DeltaToolCall(json_args='ntry": "Mexico"}')}
+
+    agent = Agent(FunctionModel(stream_function=text_stream), output_type=CityLocation)
+
+    async with agent.iter('Generate city info') as run:
+        async for node in run:
+            if agent.is_model_request_node(node):
+                async with node.stream(run.ctx) as stream:
+                    assert [c async for c in stream.stream_output(debounce_by=None)] == snapshot(
+                        [
+                            CityLocation(city='Mex'),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                            CityLocation(city='Mexico City', country='Mexico'),
+                        ]
+                    )
+
+
 def test_function_tool_event_tool_call_id_properties():
     """Ensure that the `tool_call_id` property on function tool events mirrors the underlying part's ID."""
     # Prepare a ToolCallPart with a fixed ID

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/test_tools.py

@@ -13,7 +13,16 @@ from typing_extensions import TypedDict
 
 from pydantic_ai import Agent, RunContext, Tool, UserError
 from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior
-from pydantic_ai.messages import ModelMessage, ModelRequest, ModelResponse, TextPart, ToolCallPart, ToolReturnPart
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    ToolCallPart,
+    ToolReturn,
+    ToolReturnPart,
+    UserPromptPart,
+)
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import DeferredToolCalls, ToolOutput
@@ -21,8 +30,9 @@ from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.deferred import DeferredToolset
 from pydantic_ai.toolsets.function import FunctionToolset
 from pydantic_ai.toolsets.prefixed import PrefixedToolset
+from pydantic_ai.usage import Usage
 
-from .conftest import IsStr
+from .conftest import IsDatetime, IsStr
 
 
 def test_tool_no_ctx():
@@ -1321,3 +1331,91 @@ def test_output_type_deferred_tool_calls_by_itself():
 def test_output_type_empty():
     with pytest.raises(UserError, match='At least one output type must be provided.'):
         Agent(TestModel(), output_type=[])
+
+
+def test_parallel_tool_return():
+    def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        if len(messages) == 1:
+            return ModelResponse(
+                parts=[ToolCallPart('get_price', {'fruit': 'apple'}), ToolCallPart('get_price', {'fruit': 'banana'})]
+            )
+        else:
+            return ModelResponse(
+                parts=[
+                    TextPart('Done!'),
+                ]
+            )
+
+    agent = Agent(FunctionModel(llm))
+
+    @agent.tool_plain
+    def get_price(fruit: str) -> ToolReturn:
+        return ToolReturn(
+            return_value=10.0,
+            content=f'The price of {fruit} is 10.0',
+            metadata={'foo': 'bar'},
+        )
+
+    result = agent.run_sync('What do an apple and a banana cost?')
+
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='What do an apple and a banana cost?',
+                        timestamp=IsDatetime(),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[
+                    ToolCallPart(
+                        tool_name='get_price',
+                        args={'fruit': 'apple'},
+                        tool_call_id=IsStr(),
+                    ),
+                    ToolCallPart(
+                        tool_name='get_price',
+                        args={'fruit': 'banana'},
+                        tool_call_id=IsStr(),
+                    ),
+                ],
+                usage=Usage(requests=1, request_tokens=58, response_tokens=10, total_tokens=68),
+                model_name='function:llm:',
+                timestamp=IsDatetime(),
+            ),
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name='get_price',
+                        content=10.0,
+                        tool_call_id=IsStr(),
+                        metadata={'foo': 'bar'},
+                        timestamp=IsDatetime(),
+                    ),
+                    ToolReturnPart(
+                        tool_name='get_price',
+                        content=10.0,
+                        tool_call_id=IsStr(),
+                        metadata={'foo': 'bar'},
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='The price of apple is 10.0',
+                        timestamp=IsDatetime(),
+                    ),
+                    UserPromptPart(
+                        content='The price of banana is 10.0',
+                        timestamp=IsDatetime(),
+                    ),
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='Done!')],
+                usage=Usage(requests=1, request_tokens=76, response_tokens=11, total_tokens=87),
+                model_name='function:llm:',
+                timestamp=IsDatetime(),
+            ),
+        ]
+    )

{pydantic_ai-0.4.8 → pydantic_ai-0.4.10}/tests/test_toolsets.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from collections import defaultdict
 from dataclasses import dataclass, replace
 from typing import TypeVar
 from unittest.mock import AsyncMock
@@ -10,7 +11,7 @@ from inline_snapshot import snapshot
 
 from pydantic_ai._run_context import RunContext
 from pydantic_ai._tool_manager import ToolManager
-from pydantic_ai.exceptions import UserError
+from pydantic_ai.exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior, UserError
 from pydantic_ai.messages import ToolCallPart
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.tools import ToolDefinition
@@ -494,3 +495,134 @@ async def test_context_manager_failed_initialization():
             pass
 
     assert server1.is_running is False
+
+
+async def test_tool_manager_retry_logic():
+    """Test the retry logic with failed_tools and for_run_step method."""
+
+    @dataclass
+    class TestDeps:
+        pass
+
+    # Create a toolset with tools that can fail
+    toolset = FunctionToolset[TestDeps](max_retries=2)
+    call_count: defaultdict[str, int] = defaultdict(int)
+
+    @toolset.tool
+    def failing_tool(x: int) -> int:
+        """A tool that always fails"""
+        call_count['failing_tool'] += 1
+        raise ModelRetry('This tool always fails')
+
+    @toolset.tool
+    def other_tool(x: int) -> int:
+        """A tool that works"""
+        call_count['other_tool'] += 1
+        return x * 2
+
+    # Create initial context and tool manager
+    initial_context = build_run_context(TestDeps())
+    tool_manager = await ToolManager[TestDeps].build(toolset, initial_context)
+
+    # Initially no failed tools
+    assert tool_manager.failed_tools == set()
+    assert initial_context.retries == {}
+
+    # Call the failing tool - should add to failed_tools
+    with pytest.raises(ToolRetryError):
+        await tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))
+
+    assert tool_manager.failed_tools == {'failing_tool'}
+    assert call_count['failing_tool'] == 1
+
+    # Call the working tool - should not add to failed_tools
+    result = await tool_manager.handle_call(ToolCallPart(tool_name='other_tool', args={'x': 3}))
+    assert result == 6
+    assert tool_manager.failed_tools == {'failing_tool'}  # unchanged
+    assert call_count['other_tool'] == 1
+
+    # Test for_run_step - should create new tool manager with updated retry counts
+    new_context = build_run_context(TestDeps())
+    new_tool_manager = await tool_manager.for_run_step(new_context)
+
+    # The new tool manager should have retry count for the failed tool
+    assert new_tool_manager.ctx.retries == {'failing_tool': 1}
+    assert new_tool_manager.failed_tools == set()  # reset for new run step
+
+    # Call the failing tool again in the new manager - should have retry=1
+    with pytest.raises(ToolRetryError):
+        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))
+
+    # Call the failing tool another time in the new manager
+    with pytest.raises(ToolRetryError):
+        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))
+
+    # Call the failing tool a third time in the new manager
+    with pytest.raises(ToolRetryError):
+        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))
+
+    assert new_tool_manager.failed_tools == {'failing_tool'}
+    assert call_count['failing_tool'] == 4
+
+    # Create another run step
+    another_context = build_run_context(TestDeps())
+    another_tool_manager = await new_tool_manager.for_run_step(another_context)
+
+    # Should now have retry count of 2 for failing_tool
+    assert another_tool_manager.ctx.retries == {'failing_tool': 2}
+    assert another_tool_manager.failed_tools == set()
+
+    # Call the failing tool _again_, now we should finally hit the limit
+    with pytest.raises(UnexpectedModelBehavior, match="Tool 'failing_tool' exceeded max retries count of 2"):
+        await another_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))
+
+
+async def test_tool_manager_multiple_failed_tools():
+    """Test retry logic when multiple tools fail in the same run step."""
+
+    @dataclass
+    class TestDeps:
+        pass
+
+    toolset = FunctionToolset[TestDeps]()
+
+    @toolset.tool
+    def tool_a(x: int) -> int:
+        """Tool A that fails"""
+        raise ModelRetry('Tool A fails')
+
+    @toolset.tool
+    def tool_b(x: int) -> int:
+        """Tool B that fails"""
+        raise ModelRetry('Tool B fails')
+
+    @toolset.tool
+    def tool_c(x: int) -> int:
+        """Tool C that works"""
+        return x * 3
+
+    # Create tool manager
+    context = build_run_context(TestDeps())
+    tool_manager = await ToolManager[TestDeps].build(toolset, context)
+
+    # Call tool_a - should fail and be added to failed_tools
+    with pytest.raises(ToolRetryError):
+        await tool_manager.handle_call(ToolCallPart(tool_name='tool_a', args={'x': 1}))
+    assert tool_manager.failed_tools == {'tool_a'}
+
+    # Call tool_b - should also fail and be added to failed_tools
+    with pytest.raises(ToolRetryError):
+        await tool_manager.handle_call(ToolCallPart(tool_name='tool_b', args={'x': 1}))
+    assert tool_manager.failed_tools == {'tool_a', 'tool_b'}
+
+    # Call tool_c - should succeed and not be added to failed_tools
+    result = await tool_manager.handle_call(ToolCallPart(tool_name='tool_c', args={'x': 2}))
+    assert result == 6
+    assert tool_manager.failed_tools == {'tool_a', 'tool_b'}  # unchanged
+
+    # Create next run step - should have retry counts for both failed tools
+    new_context = build_run_context(TestDeps())
+    new_tool_manager = await tool_manager.for_run_step(new_context)
+
+    assert new_tool_manager.ctx.retries == {'tool_a': 1, 'tool_b': 1}
+    assert new_tool_manager.failed_tools == set()  # reset for new run step