Repository: pydantic/pydantic-ai (language: python)

Issue:
Toolsets, OpenAPI and Model Context Protocol
TL;DR I spoke to @dsp on Friday, and the plan is to add support for MCP to PydanticAI via something I'm currently calling "toolsets".

One of the things I've realised while building PydanticAI is that building tools to give agents access to resources is going to be quite resource intensive and repetative — say you want to build an agent to interact with slack, 90% of your time will be spend figuring out the slack API and dealing with auth and reties, not building the agent.

I therefore propose the following process to make this kind of thing easy with PydanticAI:
1. Add the concept of reusable "Toolsets" which contain multiple tools, their own state (e.g. http connection) and their own relation to agent dependencies
2. Add a Toolset to support OpenAPI scheme APIs - e.g. so you could register an API which has an OpenAPI schema in a few lines of code - this should cover many common cases where you don't need full bidirectional communications
3. Add support for registering MCP toolsets
4. Create our own toolsets for interacting with common tools like slack, gmail, search etc. - this should be simple since I think many of these are already build for MCP, see [here](https://github.com/modelcontextprotocol/servers)


Discussion / hints:
Was instantly thinking of this as I saw the release. I would love to help with this if you're looking for some extra hands!

Fix patch:
diff --git a/Makefile b/Makefile
index ff20f8e9..c92c3cc3 100644
--- a/Makefile
+++ b/Makefile
@@ -29,7 +29,8 @@ lint:
 
 .PHONY: typecheck-pyright
 typecheck-pyright:
-	uv run pyright
+	@# PYRIGHT_PYTHON_IGNORE_WARNINGS avoids the overhead of making a request to github on every invocation
+	PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright
 
 .PHONY: typecheck-mypy
 typecheck-mypy:
diff --git a/docs/agents.md b/docs/agents.md
index 678ea59c..5f418e3a 100644
--- a/docs/agents.md
+++ b/docs/agents.md
@@ -420,7 +420,7 @@ print(dice_result.data)
 ```
 
 1. The simplest way to register tools via the `Agent` constructor is to pass a list of functions, the function signature is inspected to determine if the tool takes [`RunContext`][pydantic_ai.tools.RunContext].
-2. `agent_a` and `agent_b` are identical — but we can use [`Tool`][pydantic_ai.tools.Tool] to give more fine-grained control over how tools are defined, e.g. setting their name or description.
+2. `agent_a` and `agent_b` are identical — but we can use [`Tool`][pydantic_ai.tools.Tool] to reuse tool definitions and give more fine-grained control over how tools are defined, e.g. setting their name or description, or using a custom [`prepare`](#tool-prepare) method.
 
 _(This example is complete, it can be run "as is")_
 
@@ -459,10 +459,10 @@ def foobar(a: int, b: str, c: dict[str, list[float]]) -> str:
 
 
 def print_schema(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
-    tool = info.function_tools['foobar']
+    tool = info.function_tools[0]
     print(tool.description)
     #> Get me foobar.
-    print(tool.json_schema)
+    print(tool.parameters_json_schema)
     """
     {
         'description': 'Get me foobar.',
@@ -491,7 +491,167 @@ _(This example is complete, it can be run "as is")_
 
 The return type of tool can be anything which Pydantic can serialize to JSON as some models (e.g. Gemini) support semi-structured return values, some expect text (OpenAI) but seem to be just as good at extracting meaning from the data. If a Python object is returned and the model expects a string, the value will be serialized to JSON.
 
-If a tool has a single parameter that can be represented as an object in JSON schema (e.g. dataclass, TypedDict, pydantic model), the schema for the tool is simplified to be just that object. (TODO example)
+If a tool has a single parameter that can be represented as an object in JSON schema (e.g. dataclass, TypedDict, pydantic model), the schema for the tool is simplified to be just that object.
+
+Here's an example, we use [`TestModel.agent_model_function_tools`][pydantic_ai.models.test.TestModel.agent_model_function_tools] to inspect the tool schema that would be passed to the model.
+
+```py title="single_parameter_tool.py"
+from pydantic import BaseModel
+
+from pydantic_ai import Agent
+from pydantic_ai.models.test import TestModel
+
+agent = Agent()
+
+
+class Foobar(BaseModel):
+    """This is a Foobar"""
+
+    x: int
+    y: str
+    z: float = 3.14
+
+
+@agent.tool_plain
+def foobar(f: Foobar) -> str:
+    return str(f)
+
+
+test_model = TestModel()
+result = agent.run_sync('hello', model=test_model)
+print(result.data)
+#> {"foobar":"x=0 y='a' z=3.14"}
+print(test_model.agent_model_function_tools)
+"""
+[
+    ToolDefinition(
+        name='foobar',
+        description='',
+        parameters_json_schema={
+            'description': 'This is a Foobar',
+            'properties': {
+                'x': {'title': 'X', 'type': 'integer'},
+                'y': {'title': 'Y', 'type': 'string'},
+                'z': {'default': 3.14, 'title': 'Z', 'type': 'number'},
+            },
+            'required': ['x', 'y'],
+            'title': 'Foobar',
+            'type': 'object',
+        },
+        outer_typed_dict_key=None,
+    )
+]
+"""
+```
+
+_(This example is complete, it can be run "as is")_
+
+### Dynamic Function tools {#tool-prepare}
+
+Tools can optionally be defined with another function: `prepare`, which is called at each step of a run to
+customize the definition of the tool passed to the model, or omit the tool completely from that step.
+
+A `prepare` method can be registered via the `prepare` kwarg to any of the tool registration mechanisms:
+
+* [`@agent.tool`][pydantic_ai.Agent.tool] decorator
+* [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain] decorator
+* [`Tool`][pydantic_ai.tools.Tool] dataclass
+
+The `prepare` method, should be of type [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc], a function which takes [`RunContext`][pydantic_ai.tools.RunContext] and a pre-built [`ToolDefinition`][pydantic_ai.tools.ToolDefinition], and should either return that `ToolDefinition` with or without modifying it, return a new `ToolDefinition`, or return `None` to indicate this tools should not be registered for that step.
+
+Here's a simple `prepare` method that only includes the tool if the value of the dependency is `42`.
+
+As with the previous example, we use [`TestModel`][pydantic_ai.models.test.TestModel] to demonstrate the behavior without calling a real model.
+
+```py title="tool_only_if_42.py"
+from typing import Union
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.tools import ToolDefinition
+
+agent = Agent('test')
+
+
+async def only_if_42(
+    ctx: RunContext[int], tool_def: ToolDefinition
+) -> Union[ToolDefinition, None]:
+    if ctx.deps == 42:
+        return tool_def
+
+
+@agent.tool(prepare=only_if_42)
+def hitchhiker(ctx: RunContext[int], answer: str) -> str:
+    return f'{ctx.deps} {answer}'
+
+
+result = agent.run_sync('testing...', deps=41)
+print(result.data)
+#> success (no tool calls)
+result = agent.run_sync('testing...', deps=42)
+print(result.data)
+#> {"hitchhiker":"42 a"}
+```
+
+_(This example is complete, it can be run "as is")_
+
+Here's a more complex example where we change the description of the `name` parameter to based on the value of `deps`
+
+For the sake of variation, we create this tool using the [`Tool`][pydantic_ai.tools.Tool] dataclass.
+
+```py title="customize_name.py"
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.test import TestModel
+from pydantic_ai.tools import Tool, ToolDefinition
+
+
+def greet(name: str) -> str:
+    return f'hello {name}'
+
+
+async def prepare_greet(
+    ctx: RunContext[Literal['human', 'machine']], tool_def: ToolDefinition
+) -> ToolDefinition | None:
+    d = f'Name of the {ctx.deps} to greet.'
+    tool_def.parameters_json_schema['properties']['name']['description'] = d
+    return tool_def
+
+
+greet_tool = Tool(greet, prepare=prepare_greet)
+test_model = TestModel()
+agent = Agent(test_model, tools=[greet_tool], deps_type=Literal['human', 'machine'])
+
+result = agent.run_sync('testing...', deps='human')
+print(result.data)
+#> {"greet":"hello a"}
+print(test_model.agent_model_function_tools)
+"""
+[
+    ToolDefinition(
+        name='greet',
+        description='',
+        parameters_json_schema={
+            'properties': {
+                'name': {
+                    'title': 'Name',
+                    'type': 'string',
+                    'description': 'Name of the human to greet.',
+                }
+            },
+            'required': ['name'],
+            'type': 'object',
+            'additionalProperties': False,
+        },
+        outer_typed_dict_key=None,
+    )
+]
+"""
+```
+
+_(This example is complete, it can be run "as is")_
 
 ## Reflection and self-correction
 
diff --git a/pydantic_ai_slim/pydantic_ai/_pydantic.py b/pydantic_ai_slim/pydantic_ai/_pydantic.py
index 4aecd5fc..ce853636 100644
--- a/pydantic_ai_slim/pydantic_ai/_pydantic.py
+++ b/pydantic_ai_slim/pydantic_ai/_pydantic.py
@@ -17,10 +17,10 @@
 from pydantic_core import SchemaValidator, core_schema
 
 from ._griffe import doc_descriptions
-from ._utils import ObjectJsonSchema, check_object_json_schema, is_model_like
+from ._utils import check_object_json_schema, is_model_like
 
 if TYPE_CHECKING:
-    pass
+    from .tools import ObjectJsonSchema
 
 
 __all__ = 'function_schema', 'LazyTypeAdapter'
@@ -168,11 +168,13 @@ def takes_ctx(function: Callable[..., Any]) -> bool:
     """
     sig = signature(function)
     try:
-        _, first_param = next(iter(sig.parameters.items()))
+        first_param_name = next(iter(sig.parameters.keys()))
     except StopIteration:
         return False
     else:
-        return first_param.annotation is not sig.empty and _is_call_ctx(first_param.annotation)
+        type_hints = _typing_extra.get_function_type_hints(function)
+        annotation = type_hints[first_param_name]
+        return annotation is not sig.empty and _is_call_ctx(annotation)
 
 
 def _build_schema(
diff --git a/pydantic_ai_slim/pydantic_ai/_result.py b/pydantic_ai_slim/pydantic_ai/_result.py
index 3622f866..3618f335 100644
--- a/pydantic_ai_slim/pydantic_ai/_result.py
+++ b/pydantic_ai_slim/pydantic_ai/_result.py
@@ -14,7 +14,7 @@
 from .exceptions import ModelRetry
 from .messages import ModelStructuredResponse, ToolCall
 from .result import ResultData
-from .tools import AgentDeps, ResultValidatorFunc, RunContext
+from .tools import AgentDeps, ResultValidatorFunc, RunContext, ToolDefinition
 
 
 @dataclass
@@ -94,10 +94,7 @@ def build(cls, response_type: type[ResultData], name: str, description: str | No
             allow_text_result = False
 
         def _build_tool(a: Any, tool_name_: str, multiple: bool) -> ResultTool[ResultData]:
-            return cast(
-                ResultTool[ResultData],
-                ResultTool.build(a, tool_name_, description, multiple),  # pyright: ignore[reportUnknownMemberType]
-            )
+            return cast(ResultTool[ResultData], ResultTool(a, tool_name_, description, multiple))
 
         tools: dict[str, ResultTool[ResultData]] = {}
         if args := get_union_args(response_type):
@@ -121,38 +118,38 @@ def tool_names(self) -> list[str]:
         """Return the names of the tools."""
         return list(self.tools.keys())
 
+    def tool_defs(self) -> list[ToolDefinition]:
+        """Get tool definitions to register with the model."""
+        return [t.tool_def for t in self.tools.values()]
+
 
 DEFAULT_DESCRIPTION = 'The final response which ends this conversation'
 
 
-@dataclass
+@dataclass(init=False)
 class ResultTool(Generic[ResultData]):
-    name: str
-    description: str
+    tool_def: ToolDefinition
     type_adapter: TypeAdapter[Any]
-    json_schema: _utils.ObjectJsonSchema
-    outer_typed_dict_key: str | None
 
-    @classmethod
-    def build(cls, response_type: type[ResultData], name: str, description: str | None, multiple: bool) -> Self | None:
+    def __init__(self, response_type: type[ResultData], name: str, description: str | None, multiple: bool):
         """Build a ResultTool dataclass from a response type."""
         assert response_type is not str, 'ResultTool does not support str as a response type'
 
         if _utils.is_model_like(response_type):
-            type_adapter = TypeAdapter(response_type)
+            self.type_adapter = TypeAdapter(response_type)
             outer_typed_dict_key: str | None = None
             # noinspection PyArgumentList
-            json_schema = _utils.check_object_json_schema(type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
         else:
             response_data_typed_dict = TypedDict('response_data_typed_dict', {'response': response_type})  # noqa
-            type_adapter = TypeAdapter(response_data_typed_dict)
+            self.type_adapter = TypeAdapter(response_data_typed_dict)
             outer_typed_dict_key = 'response'
             # noinspection PyArgumentList
-            json_schema = _utils.check_object_json_schema(type_adapter.json_schema())
+            parameters_json_schema = _utils.check_object_json_schema(self.type_adapter.json_schema())
             # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
-            json_schema.pop('title')
+            parameters_json_schema.pop('title')
 
-        if json_schema_description := json_schema.pop('description', None):
+        if json_schema_description := parameters_json_schema.pop('description', None):
             if description is None:
                 tool_description = json_schema_description
             else:
@@ -162,11 +159,10 @@ def build(cls, response_type: type[ResultData], name: str, description: str | No
             if multiple:
                 tool_description = f'{union_arg_name(response_type)}: {tool_description}'
 
-        return cls(
+        self.tool_def = ToolDefinition(
             name=name,
             description=tool_description,
-            type_adapter=type_adapter,
-            json_schema=json_schema,
+            parameters_json_schema=parameters_json_schema,
             outer_typed_dict_key=outer_typed_dict_key,
         )
 
@@ -204,7 +200,7 @@ def validate(
             else:
                 raise
         else:
-            if k := self.outer_typed_dict_key:
+            if k := self.tool_def.outer_typed_dict_key:
                 result = result[k]
             return result
 
diff --git a/pydantic_ai_slim/pydantic_ai/_system_prompt.py b/pydantic_ai_slim/pydantic_ai/_system_prompt.py
index 7f05079a..ba21b2bc 100644
--- a/pydantic_ai_slim/pydantic_ai/_system_prompt.py
+++ b/pydantic_ai_slim/pydantic_ai/_system_prompt.py
@@ -21,7 +21,7 @@ def __post_init__(self):
 
     async def run(self, deps: AgentDeps) -> str:
         if self._takes_ctx:
-            args = (RunContext(deps, 0, None),)
+            args = (RunContext(deps, 0),)
         else:
             args = ()
 
diff --git a/pydantic_ai_slim/pydantic_ai/_utils.py b/pydantic_ai_slim/pydantic_ai/_utils.py
index b8879609..e2ef6940 100644
--- a/pydantic_ai_slim/pydantic_ai/_utils.py
+++ b/pydantic_ai_slim/pydantic_ai/_utils.py
@@ -8,12 +8,15 @@
 from datetime import datetime, timezone
 from functools import partial
 from types import GenericAlias
-from typing import Any, Callable, Generic, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast, overload
 
 from pydantic import BaseModel
 from pydantic.json_schema import JsonSchemaValue
 from typing_extensions import ParamSpec, TypeAlias, TypeGuard, is_typeddict
 
+if TYPE_CHECKING:
+    from .tools import ObjectJsonSchema
+
 _P = ParamSpec('_P')
 _R = TypeVar('_R')
 
@@ -39,10 +42,6 @@ def is_model_like(type_: Any) -> bool:
     )
 
 
-# With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_items=Any`
-ObjectJsonSchema: TypeAlias = dict[str, Any]
-
-
 def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
     from .exceptions import UserError
 
@@ -127,6 +126,12 @@ def is_left(self) -> bool:
     def whichever(self) -> Left | Right:
         return self._left.value if self._left is not None else self.right
 
+    def __repr__(self):
+        if left := self._left:
+            return f'Either(left={left.value!r})'
+        else:
+            return f'Either(right={self.right!r})'
+
 
 @asynccontextmanager
 async def group_by_temporal(
@@ -218,7 +223,7 @@ async def async_iter_groups() -> AsyncIterator[list[T]]:
 
     try:
         yield async_iter_groups()
-    finally:
+    finally:  # pragma: no cover
         # after iteration if a tasks still exists, cancel it, this will only happen if an error occurred
         if task:
             task.cancel('Cancelling due to error in iterator')
diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py
index b087c4ab..284ba7e6 100644
--- a/pydantic_ai_slim/pydantic_ai/agent.py
+++ b/pydantic_ai_slim/pydantic_ai/agent.py
@@ -22,7 +22,17 @@
     result,
 )
 from .result import ResultData
-from .tools import AgentDeps, RunContext, Tool, ToolFuncContext, ToolFuncEither, ToolFuncPlain, ToolParams
+from .tools import (
+    AgentDeps,
+    RunContext,
+    Tool,
+    ToolDefinition,
+    ToolFuncContext,
+    ToolFuncEither,
+    ToolFuncPlain,
+    ToolParams,
+    ToolPrepareFunc,
+)
 
 __all__ = ('Agent',)
 
@@ -136,7 +146,10 @@ def __init__(
         self._function_tools = {}
         self._default_retries = retries
         for tool in tools:
-            self._register_tool(Tool.infer(tool))
+            if isinstance(tool, Tool):
+                self._register_tool(tool)
+            else:
+                self._register_tool(Tool(tool))
         self._deps_type = deps_type
         self._system_prompt_functions = []
         self._max_result_retries = result_retries if result_retries is not None else retries
@@ -166,7 +179,7 @@ async def run(
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
-        model_used, custom_model, agent_model = await self._get_agent_model(model)
+        model_used, mode_selection = await self._get_model(model)
 
         deps = self._get_deps(deps)
 
@@ -174,7 +187,7 @@ async def run(
             '{agent_name} run {prompt=}',
             prompt=user_prompt,
             agent=self,
-            custom_model=custom_model,
+            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
@@ -182,14 +195,17 @@ async def run(
             self.last_run_messages = messages
 
             for tool in self._function_tools.values():
-                tool.reset()
+                tool.current_retry = 0
 
             cost = result.Cost()
 
             run_step = 0
             while True:
                 run_step += 1
-                with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                    agent_model = await self._prepare_model(model_used, deps)
+
+                with _logfire.span('model request', run_step=run_step) as model_req_span:
                     model_response, request_cost = await agent_model.request(messages)
                     model_req_span.set_attribute('response', model_response)
                     model_req_span.set_attribute('cost', request_cost)
@@ -198,7 +214,7 @@ async def run(
                 messages.append(model_response)
                 cost += request_cost
 
-                with _logfire.span('handle model response') as handle_span:
+                with _logfire.span('handle model response', run_step=run_step) as handle_span:
                     either = await self._handle_model_response(model_response, deps)
 
                     if isinstance(either, _MarkFinalResult):
@@ -273,7 +289,7 @@ async def run_stream(
             # f_back because `asynccontextmanager` adds one frame
             if frame := inspect.currentframe():  # pragma: no branch
                 self._infer_name(frame.f_back)
-        model_used, custom_model, agent_model = await self._get_agent_model(model)
+        model_used, mode_selection = await self._get_model(model)
 
         deps = self._get_deps(deps)
 
@@ -281,7 +297,7 @@ async def run_stream(
             '{agent_name} run stream {prompt=}',
             prompt=user_prompt,
             agent=self,
-            custom_model=custom_model,
+            mode_selection=mode_selection,
             model_name=model_used.name(),
             agent_name=self.name or 'agent',
         ) as run_span:
@@ -289,13 +305,17 @@ async def run_stream(
             self.last_run_messages = messages
 
             for tool in self._function_tools.values():
-                tool.reset()
+                tool.current_retry = 0
 
             cost = result.Cost()
 
             run_step = 0
             while True:
                 run_step += 1
+
+                with _logfire.span('preparing model and tools {run_step=}', run_step=run_step):
+                    agent_model = await self._prepare_model(model_used, deps)
+
                 with _logfire.span('model request {run_step=}', run_step=run_step) as model_req_span:
                     async with agent_model.request_stream(messages) as model_response:
                         model_req_span.set_attribute('response_type', model_response.__class__.__name__)
@@ -477,7 +497,11 @@ def tool(self, func: ToolFuncContext[AgentDeps, ToolParams], /) -> ToolFuncConte
 
     @overload
     def tool(
-        self, /, *, retries: int | None = None
+        self,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Callable[[ToolFuncContext[AgentDeps, ToolParams]], ToolFuncContext[AgentDeps, ToolParams]]: ...
 
     def tool(
@@ -486,9 +510,9 @@ def tool(
         /,
         *,
         retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Any:
-        """Decorator to register a tool function which takes
-        [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
+        """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
         Can decorate a sync or async functions.
 
@@ -521,20 +545,23 @@ async def spam(ctx: RunContext[str], y: float) -> float:
             func: The tool function to register.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
-        """  # noqa: D205
+            prepare: custom method to prepare the tool definition for each step, return `None` to omit this
+                tool from a given step. This is useful if you want to customise a tool at call time,
+                or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
+        """
         if func is None:
 
             def tool_decorator(
                 func_: ToolFuncContext[AgentDeps, ToolParams],
             ) -> ToolFuncContext[AgentDeps, ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, True, retries)
+                self._register_function(func_, True, retries, prepare)
                 return func_
 
             return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_function(func, True, retries)
+            self._register_function(func, True, retries, prepare)
             return func
 
     @overload
@@ -542,10 +569,21 @@ def tool_plain(self, func: ToolFuncPlain[ToolParams], /) -> ToolFuncPlain[ToolPa
 
     @overload
     def tool_plain(
-        self, /, *, retries: int | None = None
+        self,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
-    def tool_plain(self, func: ToolFuncPlain[ToolParams] | None = None, /, *, retries: int | None = None) -> Any:
+    def tool_plain(
+        self,
+        func: ToolFuncPlain[ToolParams] | None = None,
+        /,
+        *,
+        retries: int | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
+    ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
         Can decorate a sync or async functions.
@@ -579,30 +617,38 @@ async def spam(ctx: RunContext[str]) -> float:
             func: The tool function to register.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
+            prepare: custom method to prepare the tool definition for each step, return `None` to omit this
+                tool from a given step. This is useful if you want to customise a tool at call time,
+                or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
         """
         if func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_function(func_, False, retries)
+                self._register_function(func_, False, retries, prepare)
                 return func_
 
             return tool_decorator
         else:
-            self._register_function(func, False, retries)
+            self._register_function(func, False, retries, prepare)
             return func
 
     def _register_function(
-        self, func: ToolFuncEither[AgentDeps, ToolParams], takes_ctx: bool, retries: int | None
+        self,
+        func: ToolFuncEither[AgentDeps, ToolParams],
+        takes_ctx: bool,
+        retries: int | None,
+        prepare: ToolPrepareFunc[AgentDeps] | None,
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
-        tool = Tool(func, takes_ctx, max_retries=retries_)
+        tool = Tool(func, takes_ctx=takes_ctx, max_retries=retries_, prepare=prepare)
         self._register_tool(tool)
 
     def _register_tool(self, tool: Tool[AgentDeps]) -> None:
         """Private utility to register a tool instance."""
         if tool.max_retries is None:
+            # noinspection PyTypeChecker
             tool = dataclasses.replace(tool, max_retries=self._default_retries)
 
         if tool.name in self._function_tools:
@@ -613,16 +659,14 @@ def _register_tool(self, tool: Tool[AgentDeps]) -> None:
 
         self._function_tools[tool.name] = tool
 
-    async def _get_agent_model(
-        self, model: models.Model | models.KnownModelName | None
-    ) -> tuple[models.Model, models.Model | None, models.AgentModel]:
+    async def _get_model(self, model: models.Model | models.KnownModelName | None) -> tuple[models.Model, str]:
         """Create a model configured for this agent.
 
         Args:
             model: model to use for this run, required if `model` was not set when creating the agent.
 
         Returns:
-            a tuple of `(model used, custom_model if any, agent_model)`
+            a tuple of `(model used, how the model was selected)`
         """
         model_: models.Model
         if some_model := self._override_model:
@@ -633,19 +677,35 @@ async def _get_agent_model(
                     '(Even when `override(model=...)` is customizing the model that will actually be called)'
                 )
             model_ = some_model.value
-            custom_model = None
+            mode_selection = 'override-model'
         elif model is not None:
-            custom_model = model_ = models.infer_model(model)
+            model_ = models.infer_model(model)
+            mode_selection = 'custom'
         elif self.model is not None:
             # noinspection PyTypeChecker
             model_ = self.model = models.infer_model(self.model)
-            custom_model = None
+            mode_selection = 'from-agent'
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
-        result_tools = list(self._result_schema.tools.values()) if self._result_schema else None
-        agent_model = await model_.agent_model(self._function_tools, self._allow_text_result, result_tools)
-        return model_, custom_model, agent_model
+        return model_, mode_selection
+
+    async def _prepare_model(self, model: models.Model, deps: AgentDeps) -> models.AgentModel:
+        """Create building tools and create an agent model."""
+        function_tools: list[ToolDefinition] = []
+
+        async def add_tool(tool: Tool[AgentDeps]) -> None:
+            ctx = RunContext(deps, tool.current_retry, tool.name)
+            if tool_def := await tool.prepare_tool_def(ctx):
+                function_tools.append(tool_def)
+
+        await asyncio.gather(*map(add_tool, self._function_tools.values()))
+
+        return await model.agent_model(
+            function_tools=function_tools,
+            allow_text_result=self._allow_text_result,
+            result_tools=self._result_schema.tool_defs() if self._result_schema is not None else [],
+        )
 
     async def _prepare_messages(
         self, deps: AgentDeps, user_prompt: str, message_history: list[_messages.Message] | None
diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py
index 9ab01165..283c1b25 100644
--- a/pydantic_ai_slim/pydantic_ai/models/__init__.py
+++ b/pydantic_ai_slim/pydantic_ai/models/__init__.py
@@ -7,11 +7,11 @@
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable, Iterator
 from contextlib import asynccontextmanager, contextmanager
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING, Literal, Protocol, Union
+from typing import TYPE_CHECKING, Literal, Union
 
 import httpx
 
@@ -19,8 +19,8 @@
 from ..messages import Message, ModelAnyResponse, ModelStructuredResponse
 
 if TYPE_CHECKING:
-    from .._utils import ObjectJsonSchema
     from ..result import Cost
+    from ..tools import ToolDefinition
 
 
 KnownModelName = Literal[
@@ -63,11 +63,12 @@ class Model(ABC):
     @abstractmethod
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
-        """Create an agent model.
+        """Create an agent model, this is called for each step of an agent run.
 
         This is async in case slow/async config checks need to be performed that can't be done in `__init__`.
 
@@ -87,7 +88,7 @@ def name(self) -> str:
 
 
 class AgentModel(ABC):
-    """Model configured for a specific agent."""
+    """Model configured for each step of an Agent run."""
 
     @abstractmethod
     async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
@@ -256,35 +257,6 @@ def infer_model(model: Model | KnownModelName) -> Model:
         raise UserError(f'Unknown model: {model}')
 
 
-class AbstractToolDefinition(Protocol):
-    """Abstract definition of a function/tool.
-
-    This is used for both function tools and result tools.
-    """
-
-    @property
-    def name(self) -> str:
-        """The name of the tool."""
-        ...
-
-    @property
-    def description(self) -> str:
-        """The description of the tool."""
-        ...
-
-    @property
-    def json_schema(self) -> ObjectJsonSchema:
-        """The JSON schema for the tool's arguments."""
-        ...
-
-    @property
-    def outer_typed_dict_key(self) -> str | None:
-        """The key in the outer [TypedDict] that wraps a result tool.
-
-        This will only be set for result tools which don't have an `object` JSON schema.
-        """
-
-
 @cache
 def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     """Cached HTTPX async client so multiple agents and calls can share the same client.
diff --git a/pydantic_ai_slim/pydantic_ai/models/function.py b/pydantic_ai_slim/pydantic_ai/models/function.py
index 9ca06d00..043fb1c8 100644
--- a/pydantic_ai_slim/pydantic_ai/models/function.py
+++ b/pydantic_ai_slim/pydantic_ai/models/function.py
@@ -2,7 +2,7 @@
 
 import inspect
 import re
-from collections.abc import AsyncIterator, Awaitable, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Awaitable, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -14,14 +14,8 @@
 
 from .. import _utils, result
 from ..messages import ArgsJson, Message, ModelAnyResponse, ModelStructuredResponse, ToolCall
-from . import (
-    AbstractToolDefinition,
-    AgentModel,
-    EitherStreamedResponse,
-    Model,
-    StreamStructuredResponse,
-    StreamTextResponse,
-)
+from ..tools import ToolDefinition
+from . import AgentModel, EitherStreamedResponse, Model, StreamStructuredResponse, StreamTextResponse
 
 
 @dataclass(init=False)
@@ -59,11 +53,11 @@ def __init__(self, function: FunctionDef | None = None, *, stream_function: Stre
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
-        result_tools = list(result_tools) if result_tools is not None else None
         return FunctionAgentModel(
             self.function, self.stream_function, AgentInfo(function_tools, allow_text_result, result_tools)
         )
@@ -84,7 +78,7 @@ class AgentInfo:
     This is passed as the second to functions used within [`FunctionModel`][pydantic_ai.models.function.FunctionModel].
     """
 
-    function_tools: Mapping[str, AbstractToolDefinition]
+    function_tools: list[ToolDefinition]
     """The function tools available on this agent.
 
     These are the tools registered via the [`tool`][pydantic_ai.Agent.tool] and
@@ -92,7 +86,7 @@ class AgentInfo:
     """
     allow_text_result: bool
     """Whether a plain text result is allowed."""
-    result_tools: list[AbstractToolDefinition] | None
+    result_tools: list[ToolDefinition]
     """The tools that can called as the final result of the run."""
 
 
diff --git a/pydantic_ai_slim/pydantic_ai/models/gemini.py b/pydantic_ai_slim/pydantic_ai/models/gemini.py
index 7f1b3ca3..382fbe85 100644
--- a/pydantic_ai_slim/pydantic_ai/models/gemini.py
+++ b/pydantic_ai_slim/pydantic_ai/models/gemini.py
@@ -2,7 +2,7 @@
 
 import os
 import re
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
@@ -25,8 +25,8 @@
     ToolCall,
     ToolReturn,
 )
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -90,9 +90,10 @@ def __init__(
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
         return GeminiAgentModel(
             http_client=self.http_client,
@@ -142,13 +143,13 @@ def __init__(
         model_name: GeminiModelName,
         auth: AuthProtocol,
         url: str,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ):
         check_allow_model_requests()
-        tools = [_function_from_abstract_tool(t) for t in function_tools.values()]
-        if result_tools is not None:
+        tools = [_function_from_abstract_tool(t) for t in function_tools]
+        if result_tools:
             tools += [_function_from_abstract_tool(t) for t in result_tools]
 
         if allow_text_result:
@@ -504,8 +505,8 @@ class _GeminiFunction(TypedDict):
     """
 
 
-def _function_from_abstract_tool(tool: AbstractToolDefinition) -> _GeminiFunction:
-    json_schema = _GeminiJsonSchema(tool.json_schema).simplify()
+def _function_from_abstract_tool(tool: ToolDefinition) -> _GeminiFunction:
+    json_schema = _GeminiJsonSchema(tool.parameters_json_schema).simplify()
     f = _GeminiFunction(
         name=tool.name,
         description=tool.description,
diff --git a/pydantic_ai_slim/pydantic_ai/models/groq.py b/pydantic_ai_slim/pydantic_ai/models/groq.py
index 46420b48..7e938089 100644
--- a/pydantic_ai_slim/pydantic_ai/models/groq.py
+++ b/pydantic_ai_slim/pydantic_ai/models/groq.py
@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
@@ -21,8 +21,8 @@
     ToolReturn,
 )
 from ..result import Cost
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -109,13 +109,14 @@ def __init__(
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
         check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools.values()]
-        if result_tools is not None:
+        tools = [self._map_tool_definition(r) for r in function_tools]
+        if result_tools:
             tools += [self._map_tool_definition(r) for r in result_tools]
         return GroqAgentModel(
             self.client,
@@ -128,13 +129,13 @@ def name(self) -> str:
         return f'groq:{self.model_name}'
 
     @staticmethod
-    def _map_tool_definition(f: AbstractToolDefinition) -> chat.ChatCompletionToolParam:
+    def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam:
         return {
             'type': 'function',
             'function': {
                 'name': f.name,
                 'description': f.description,
-                'parameters': f.json_schema,
+                'parameters': f.parameters_json_schema,
             },
         }
 
diff --git a/pydantic_ai_slim/pydantic_ai/models/openai.py b/pydantic_ai_slim/pydantic_ai/models/openai.py
index 4b409f72..68f96b54 100644
--- a/pydantic_ai_slim/pydantic_ai/models/openai.py
+++ b/pydantic_ai_slim/pydantic_ai/models/openai.py
@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
@@ -21,8 +21,8 @@
     ToolReturn,
 )
 from ..result import Cost
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -89,13 +89,14 @@ def __init__(
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
         check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools.values()]
-        if result_tools is not None:
+        tools = [self._map_tool_definition(r) for r in function_tools]
+        if result_tools:
             tools += [self._map_tool_definition(r) for r in result_tools]
         return OpenAIAgentModel(
             self.client,
@@ -108,13 +109,13 @@ def name(self) -> str:
         return f'openai:{self.model_name}'
 
     @staticmethod
-    def _map_tool_definition(f: AbstractToolDefinition) -> chat.ChatCompletionToolParam:
+    def _map_tool_definition(f: ToolDefinition) -> chat.ChatCompletionToolParam:
         return {
             'type': 'function',
             'function': {
                 'name': f.name,
                 'description': f.description,
-                'parameters': f.json_schema,
+                'parameters': f.parameters_json_schema,
             },
         }
 
diff --git a/pydantic_ai_slim/pydantic_ai/models/vertexai.py b/pydantic_ai_slim/pydantic_ai/models/vertexai.py
index cdab5812..95060189 100644
--- a/pydantic_ai_slim/pydantic_ai/models/vertexai.py
+++ b/pydantic_ai_slim/pydantic_ai/models/vertexai.py
@@ -1,6 +1,5 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -10,7 +9,8 @@
 
 from .._utils import run_in_executor
 from ..exceptions import UserError
-from . import AbstractToolDefinition, Model, cached_async_http_client
+from ..tools import ToolDefinition
+from . import Model, cached_async_http_client
 from .gemini import GeminiAgentModel, GeminiModelName
 
 try:
@@ -109,9 +109,10 @@ def __init__(
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
         url, auth = await self._ainit()
         return GeminiAgentModel(
diff --git a/pydantic_ai_slim/pydantic_ai/tools.py b/pydantic_ai_slim/pydantic_ai/tools.py
index 8021003e..85f32814 100644
--- a/pydantic_ai_slim/pydantic_ai/tools.py
+++ b/pydantic_ai_slim/pydantic_ai/tools.py
@@ -7,7 +7,7 @@
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
-from typing_extensions import Concatenate, ParamSpec, final
+from typing_extensions import Concatenate, ParamSpec, TypeAlias
 
 from . import _pydantic, _utils, messages
 from .exceptions import ModelRetry, UnexpectedModelBehavior
@@ -27,7 +27,10 @@
     'ToolFuncPlain',
     'ToolFuncEither',
     'ToolParams',
+    'ToolPrepareFunc',
     'Tool',
+    'ObjectJsonSchema',
+    'ToolDefinition',
 )
 
 AgentDeps = TypeVar('AgentDeps')
@@ -42,7 +45,7 @@ class RunContext(Generic[AgentDeps]):
     """Dependencies for the agent."""
     retry: int
     """Number of retries so far."""
-    tool_name: str | None
+    tool_name: str | None = None
     """Name of the tool being called."""
 
 
@@ -91,11 +94,37 @@ class RunContext(Generic[AgentDeps]):
 
 Usage `ToolFuncEither[AgentDeps, ToolParams]`.
 """
+ToolPrepareFunc: TypeAlias = 'Callable[[RunContext[AgentDeps], ToolDefinition], Awaitable[ToolDefinition | None]]'
+"""Definition of a function that can prepare a tool definition at call time.
+
+See [tool docs](../agents.md#tool-prepare) for more information.
+
+Example — here `only_if_42` is valid as a `ToolPrepareFunc`:
+
+```py
+from typing import Union
+
+from pydantic_ai import RunContext, Tool
+from pydantic_ai.tools import ToolDefinition
+
+async def only_if_42(
+    ctx: RunContext[int], tool_def: ToolDefinition
+) -> Union[ToolDefinition, None]:
+    if ctx.deps == 42:
+        return tool_def
+
+def hitchhiker(ctx: RunContext[int], answer: str) -> str:
+    return f'{ctx.deps} {answer}'
+
+hitchhiker = Tool(hitchhiker, prepare=only_if_42)
+```
+
+Usage `ToolPrepareFunc[AgentDeps]`.
+"""
 
 A = TypeVar('A')
 
 
-@final
 @dataclass(init=False)
 class Tool(Generic[AgentDeps]):
     """A tool function for an agent."""
@@ -105,22 +134,24 @@ class Tool(Generic[AgentDeps]):
     max_retries: int | None
     name: str
     description: str
+    prepare: ToolPrepareFunc[AgentDeps] | None
     _is_async: bool = field(init=False)
     _single_arg_name: str | None = field(init=False)
     _positional_fields: list[str] = field(init=False)
     _var_positional_field: str | None = field(init=False)
     _validator: SchemaValidator = field(init=False, repr=False)
-    _json_schema: _utils.ObjectJsonSchema = field(init=False)
-    _current_retry: int = field(default=0, init=False)
+    _parameters_json_schema: ObjectJsonSchema = field(init=False)
+    current_retry: int = field(default=0, init=False)
 
     def __init__(
         self,
         function: ToolFuncEither[AgentDeps, ...],
-        takes_ctx: bool,
         *,
+        takes_ctx: bool | None = None,
         max_retries: int | None = None,
         name: str | None = None,
         description: str | None = None,
+        prepare: ToolPrepareFunc[AgentDeps] | None = None,
     ):
         """Create a new tool instance.
 
@@ -132,47 +163,77 @@ def __init__(
         async def my_tool(ctx: RunContext[int], x: int, y: int) -> str:
             return f'{ctx.deps} {x} {y}'
 
-        agent = Agent('test', tools=[Tool(my_tool, True)])
+        agent = Agent('test', tools=[Tool(my_tool)])
+        ```
+
+        or with a custom prepare method:
+
+        ```py
+        from typing import Union
+
+        from pydantic_ai import Agent, RunContext, Tool
+        from pydantic_ai.tools import ToolDefinition
+
+        async def my_tool(ctx: RunContext[int], x: int, y: int) -> str:
+            return f'{ctx.deps} {x} {y}'
+
+        async def prep_my_tool(
+            ctx: RunContext[int], tool_def: ToolDefinition
+        ) -> Union[ToolDefinition, None]:
+            # only register the tool if `deps == 42`
+            if ctx.deps == 42:
+                return tool_def
+
+        agent = Agent('test', tools=[Tool(my_tool, prepare=prep_my_tool)])
         ```
 
+
         Args:
             function: The Python function to call as the tool.
-            takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] first argument.
+            takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] first argument,
+                this is inferred if unset.
             max_retries: Maximum number of retries allowed for this tool, set to the agent default if `None`.
             name: Name of the tool, inferred from the function if `None`.
             description: Description of the tool, inferred from the function if `None`.
+            prepare: custom method to prepare the tool definition for each step, return `None` to omit this
+                tool from a given step. This is useful if you want to customise a tool at call time,
+                or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc].
         """
+        if takes_ctx is None:
+            takes_ctx = _pydantic.takes_ctx(function)
+
         f = _pydantic.function_schema(function, takes_ctx)
         self.function = function
         self.takes_ctx = takes_ctx
         self.max_retries = max_retries
         self.name = name or function.__name__
         self.description = description or f['description']
+        self.prepare = prepare
         self._is_async = inspect.iscoroutinefunction(self.function)
         self._single_arg_name = f['single_arg_name']
         self._positional_fields = f['positional_fields']
         self._var_positional_field = f['var_positional_field']
         self._validator = f['validator']
-        self._json_schema = f['json_schema']
+        self._parameters_json_schema = f['json_schema']
 
-    @staticmethod
-    def infer(function: ToolFuncEither[A, ...] | Tool[A]) -> Tool[A]:
-        """Create a tool from a pure function, inferring whether it takes `RunContext` as its first argument.
+    async def prepare_tool_def(self, ctx: RunContext[AgentDeps]) -> ToolDefinition | None:
+        """Get the tool definition.
 
-        Args:
-            function: The tool function to wrap; or for convenience, a `Tool` instance.
+        By default, this method creates a tool definition, then either returns it, or calls `self.prepare`
+        if it's set.
 
         Returns:
-            A new `Tool` instance.
+            return a `ToolDefinition` or `None` if the tools should not be registered for this run.
         """
-        if isinstance(function, Tool):
-            return function
+        tool_def = ToolDefinition(
+            name=self.name,
+            description=self.description,
+            parameters_json_schema=self._parameters_json_schema,
+        )
+        if self.prepare is not None:
+            return await self.prepare(ctx, tool_def)
         else:
-            return Tool(function, takes_ctx=_pydantic.takes_ctx(function))
-
-    def reset(self) -> None:
-        """Reset the current retry count."""
-        self._current_retry = 0
+            return tool_def
 
     async def run(self, deps: AgentDeps, message: messages.ToolCall) -> messages.Message:
         """Run the tool function asynchronously."""
@@ -195,28 +256,20 @@ async def run(self, deps: AgentDeps, message: messages.ToolCall) -> messages.Mes
         except ModelRetry as e:
             return self._on_error(e, message)
 
-        self._current_retry = 0
+        self.current_retry = 0
         return messages.ToolReturn(
             tool_name=message.tool_name,
             content=response_content,
             tool_id=message.tool_id,
         )
 
-    @property
-    def json_schema(self) -> _utils.ObjectJsonSchema:
-        return self._json_schema
-
-    @property
-    def outer_typed_dict_key(self) -> str | None:
-        return None
-
     def _call_args(
         self, deps: AgentDeps, args_dict: dict[str, Any], message: messages.ToolCall
     ) -> tuple[list[Any], dict[str, Any]]:
         if self._single_arg_name:
             args_dict = {self._single_arg_name: args_dict}
 
-        args = [RunContext(deps, self._current_retry, message.tool_name)] if self.takes_ctx else []
+        args = [RunContext(deps, self.current_retry, message.tool_name)] if self.takes_ctx else []
         for positional_field in self._positional_fields:
             args.append(args_dict.pop(positional_field))
         if self._var_positional_field:
@@ -225,8 +278,8 @@ def _call_args(
         return args, args_dict
 
     def _on_error(self, exc: ValidationError | ModelRetry, call_message: messages.ToolCall) -> messages.RetryPrompt:
-        self._current_retry += 1
-        if self.max_retries is None or self._current_retry > self.max_retries:
+        self.current_retry += 1
+        if self.max_retries is None or self.current_retry > self.max_retries:
             raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {self.max_retries}') from exc
         else:
             if isinstance(exc, ValidationError):
@@ -238,3 +291,35 @@ def _on_error(self, exc: ValidationError | ModelRetry, call_message: messages.To
                 content=content,
                 tool_id=call_message.tool_id,
             )
+
+
+ObjectJsonSchema: TypeAlias = dict[str, Any]
+"""Type representing JSON schema of an object, e.g. where `"type": "object"`.
+
+This type is used to define tools parameters (aka arguments) in [ToolDefinition][pydantic_ai.tools.ToolDefinition].
+
+With PEP-728 this should be a TypedDict with `type: Literal['object']`, and `extra_items=Any`
+"""
+
+
+@dataclass
+class ToolDefinition:
+    """Definition of a tool passed to a model.
+
+    This is used for both function tools result tools.
+    """
+
+    name: str
+    """The name of the tool."""
+
+    description: str
+    """The description of the tool."""
+
+    parameters_json_schema: ObjectJsonSchema
+    """The JSON schema for the tool's parameters."""
+
+    outer_typed_dict_key: str | None = None
+    """The key in the outer [TypedDict] that wraps a result tool.
+
+    This will only be set for result tools which don't have an `object` JSON schema.
+    """
diff --git a/pyproject.toml b/pyproject.toml
index b1758d5c..40742777 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -152,7 +152,6 @@ show_missing = true
 ignore_errors = true
 precision = 2
 exclude_lines = [
-    'def __repr__',
     'pragma: no cover',
     'raise NotImplementedError',
     'if TYPE_CHECKING:',


Test patch:
diff --git a/docs/testing-evals.md b/docs/testing-evals.md
index a297753d..63195c0c 100644
--- a/docs/testing-evals.md
+++ b/docs/testing-evals.md
@@ -209,8 +209,6 @@ def call_weather_forecast(  # (1)!
 ) -> ModelAnyResponse:
     if len(messages) == 2:
         # first call, call the weather forecast tool
-        assert set(info.function_tools.keys()) == {'weather_forecast'}
-
         user_prompt = messages[1]
         m = re.search(r'\d{4}-\d{2}-\d{2}', user_prompt.content)
         assert m is not None
diff --git a/pydantic_ai_slim/pydantic_ai/models/test.py b/pydantic_ai_slim/pydantic_ai/models/test.py
index 6658368f..8ab873c6 100644
--- a/pydantic_ai_slim/pydantic_ai/models/test.py
+++ b/pydantic_ai_slim/pydantic_ai/models/test.py
@@ -2,7 +2,7 @@
 
 import re
 import string
-from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable, Iterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import date, datetime, timedelta
@@ -21,8 +21,8 @@
     ToolReturn,
 )
 from ..result import Cost
+from ..tools import ToolDefinition
 from . import (
-    AbstractToolDefinition,
     AgentModel,
     EitherStreamedResponse,
     Model,
@@ -55,25 +55,38 @@ class TestModel(Model):
     """If set, these args will be passed to the result tool."""
     seed: int = 0
     """Seed for generating random data."""
-    # these fields are set when the model is called by the agent
-    agent_model_tools: Mapping[str, AbstractToolDefinition] | None = field(default=None, init=False)
+    agent_model_function_tools: list[ToolDefinition] | None = field(default=None, init=False)
+    """Definition of function tools passed to the model.
+
+    This is set when the model is called, so will reflect the function tools from the last step of the last run.
+    """
     agent_model_allow_text_result: bool | None = field(default=None, init=False)
-    agent_model_result_tools: list[AbstractToolDefinition] | None = field(default=None, init=False)
+    """Whether plain text responses from the model are allowed.
+
+    This is set when the model is called, so will reflect the value from the last step of the last run.
+    """
+    agent_model_result_tools: list[ToolDefinition] | None = field(default=None, init=False)
+    """Definition of result tools passed to the model.
+
+    This is set when the model is called, so will reflect the result tools from the last step of the last run.
+    """
 
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
-        self.agent_model_tools = function_tools
+        self.agent_model_function_tools = function_tools
         self.agent_model_allow_text_result = allow_text_result
-        self.agent_model_result_tools = list(result_tools) if result_tools is not None else None
+        self.agent_model_result_tools = result_tools
 
         if self.call_tools == 'all':
-            tool_calls = [(r.name, r) for r in function_tools.values()]
+            tool_calls = [(r.name, r) for r in function_tools]
         else:
-            tools_to_call = (function_tools[name] for name in self.call_tools)
+            function_tools_lookup = {t.name: t for t in function_tools}
+            tools_to_call = (function_tools_lookup[name] for name in self.call_tools)
             tool_calls = [(r.name, r) for r in tools_to_call]
 
         if self.custom_result_text is not None:
@@ -90,11 +103,12 @@ async def agent_model(
                 result = _utils.Either(right=self.custom_result_args)
         elif allow_text_result:
             result = _utils.Either(left=None)
-        elif result_tools is not None:
+        elif result_tools:
             result = _utils.Either(right=None)
         else:
             result = _utils.Either(left=None)
-        return TestAgentModel(tool_calls, result, self.agent_model_result_tools, self.seed)
+
+        return TestAgentModel(tool_calls, result, result_tools, self.seed)
 
     def name(self) -> str:
         return 'test-model'
@@ -107,13 +121,11 @@ class TestAgentModel(AgentModel):
     # NOTE: Avoid test discovery by pytest.
     __test__ = False
 
-    tool_calls: list[tuple[str, AbstractToolDefinition]]
+    tool_calls: list[tuple[str, ToolDefinition]]
     # left means the text is plain text; right means it's a function call
     result: _utils.Either[str | None, Any | None]
-    result_tools: list[AbstractToolDefinition] | None
+    result_tools: list[ToolDefinition]
     seed: int
-    step: int = 0
-    last_message_count: int = 0
 
     async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
         return self._request(messages), Cost()
@@ -127,18 +139,19 @@ async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherS
         else:
             yield TestStreamStructuredResponse(msg, cost)
 
-    def gen_tool_args(self, tool_def: AbstractToolDefinition) -> Any:
-        return _JsonSchemaTestData(tool_def.json_schema, self.seed).generate()
+    def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
+        return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
 
     def _request(self, messages: list[Message]) -> ModelAnyResponse:
-        if self.step == 0 and self.tool_calls:
+        # if there are tools, the first thing we want to do is call all of them
+        if self.tool_calls and not any(m.role == 'model-structured-response' for m in messages):
             calls = [ToolCall.from_dict(name, self.gen_tool_args(args)) for name, args in self.tool_calls]
-            self.step += 1
-            self.last_message_count = len(messages)
             return ModelStructuredResponse(calls=calls)
 
-        new_messages = messages[self.last_message_count :]
-        self.last_message_count = len(messages)
+        # get messages since the last model response
+        new_messages = _get_new_messages(messages)
+
+        # check if there are any retry prompts, if so retry them
         new_retry_names = {m.tool_name for m in new_messages if isinstance(m, RetryPrompt)}
         if new_retry_names:
             calls = [
@@ -146,34 +159,42 @@ def _request(self, messages: list[Message]) -> ModelAnyResponse:
                 for name, args in self.tool_calls
                 if name in new_retry_names
             ]
-            self.step += 1
             return ModelStructuredResponse(calls=calls)
-        else:
-            if response_text := self.result.left:
-                self.step += 1
-                if response_text.value is None:
-                    # build up details of tool responses
-                    output: dict[str, Any] = {}
-                    for message in messages:
-                        if isinstance(message, ToolReturn):
-                            output[message.tool_name] = message.content
-                    if output:
-                        return ModelTextResponse(content=pydantic_core.to_json(output).decode())
-                    else:
-                        return ModelTextResponse(content='success (no tool calls)')
+
+        if response_text := self.result.left:
+            if response_text.value is None:
+                # build up details of tool responses
+                output: dict[str, Any] = {}
+                for message in messages:
+                    if isinstance(message, ToolReturn):
+                        output[message.tool_name] = message.content
+                if output:
+                    return ModelTextResponse(content=pydantic_core.to_json(output).decode())
                 else:
-                    return ModelTextResponse(content=response_text.value)
+                    return ModelTextResponse(content='success (no tool calls)')
             else:
-                assert self.result_tools is not None, 'No result tools provided'
-                custom_result_args = self.result.right
-                result_tool = self.result_tools[self.seed % len(self.result_tools)]
-                if custom_result_args is not None:
-                    self.step += 1
-                    return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, custom_result_args)])
-                else:
-                    response_args = self.gen_tool_args(result_tool)
-                    self.step += 1
-                    return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, response_args)])
+                return ModelTextResponse(content=response_text.value)
+        else:
+            assert self.result_tools, 'No result tools provided'
+            custom_result_args = self.result.right
+            result_tool = self.result_tools[self.seed % len(self.result_tools)]
+            if custom_result_args is not None:
+                return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, custom_result_args)])
+            else:
+                response_args = self.gen_tool_args(result_tool)
+                return ModelStructuredResponse(calls=[ToolCall.from_dict(result_tool.name, response_args)])
+
+
+def _get_new_messages(messages: list[Message]) -> list[Message]:
+    last_model_index = None
+    for i, m in enumerate(messages):
+        if m.role in ('model-structured-response', 'model-text-response'):
+            last_model_index = i
+
+    if last_model_index is not None:
+        return messages[last_model_index + 1 :]
+    else:
+        return []
 
 
 @dataclass
diff --git a/tests/models/test_gemini.py b/tests/models/test_gemini.py
index fb24e975..adb342b8 100644
--- a/tests/models/test_gemini.py
+++ b/tests/models/test_gemini.py
@@ -12,7 +12,7 @@
 from pydantic import BaseModel
 from typing_extensions import Literal, TypeAlias
 
-from pydantic_ai import Agent, ModelRetry, UnexpectedModelBehavior, UserError, _utils
+from pydantic_ai import Agent, ModelRetry, UnexpectedModelBehavior, UserError
 from pydantic_ai.messages import (
     ArgsDict,
     ModelStructuredResponse,
@@ -42,6 +42,7 @@
     _GeminiUsageMetaData,
 )
 from pydantic_ai.result import Cost
+from pydantic_ai.tools import ToolDefinition
 
 from ..conftest import ClientWithHandler, IsNow, TestEnv
 
@@ -76,7 +77,7 @@ def test_api_key_empty(env: TestEnv):
 
 async def test_agent_model_simple(allow_model_requests: None):
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    agent_model = await m.agent_model({}, True, None)
+    agent_model = await m.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
     assert isinstance(agent_model.http_client, httpx.AsyncClient)
     assert agent_model.model_name == 'gemini-1.5-flash'
     assert isinstance(agent_model.auth, ApiKeyAuth)
@@ -85,24 +86,15 @@ async def test_agent_model_simple(allow_model_requests: None):
     assert agent_model.tool_config is None
 
 
-@dataclass
-class TestToolDefinition:
-    __test__ = False
-    name: str
-    description: str
-    json_schema: _utils.ObjectJsonSchema
-    outer_typed_dict_key: str | None = None
-
-
 async def test_agent_model_tools(allow_model_requests: None):
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    tools = {
-        'foo': TestToolDefinition(
+    tools = [
+        ToolDefinition(
             'foo',
             'This is foo',
             {'type': 'object', 'title': 'Foo', 'properties': {'bar': {'type': 'number', 'title': 'Bar'}}},
         ),
-        'apple': TestToolDefinition(
+        ToolDefinition(
             'apple',
             'This is apple',
             {
@@ -112,13 +104,13 @@ async def test_agent_model_tools(allow_model_requests: None):
                 },
             },
         ),
-    }
-    result_tool = TestToolDefinition(
+    ]
+    result_tool = ToolDefinition(
         'result',
         'This is the tool for the final Result',
         {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}, 'required': ['spam']},
     )
-    agent_model = await m.agent_model(tools, True, [result_tool])
+    agent_model = await m.agent_model(function_tools=tools, allow_text_result=True, result_tools=[result_tool])
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -152,12 +144,12 @@ async def test_agent_model_tools(allow_model_requests: None):
 
 async def test_require_response_tool(allow_model_requests: None):
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    result_tool = TestToolDefinition(
+    result_tool = ToolDefinition(
         'result',
         'This is the tool for the final Result',
         {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}},
     )
-    agent_model = await m.agent_model({}, False, [result_tool])
+    agent_model = await m.agent_model(function_tools=[], allow_text_result=False, result_tools=[result_tool])
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -209,12 +201,12 @@ class Locations(BaseModel):
     )
 
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    result_tool = TestToolDefinition(
+    result_tool = ToolDefinition(
         'result',
         'This is the tool for the final Result',
         json_schema,
     )
-    agent_model = await m.agent_model({}, True, [result_tool])
+    agent_model = await m.agent_model(function_tools=[], allow_text_result=True, result_tools=[result_tool])
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -255,12 +247,12 @@ class Locations(BaseModel):
     json_schema = Locations.model_json_schema()
 
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    result_tool = TestToolDefinition(
+    result_tool = ToolDefinition(
         'result',
         'This is the tool for the final Result',
         json_schema,
     )
-    agent_model = await m.agent_model({}, True, [result_tool])
+    agent_model = await m.agent_model(function_tools=[], allow_text_result=True, result_tools=[result_tool])
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -318,13 +310,13 @@ class Location(BaseModel):
     )
 
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    result_tool = TestToolDefinition(
+    result_tool = ToolDefinition(
         'result',
         'This is the tool for the final Result',
         json_schema,
     )
     with pytest.raises(UserError, match=r'Recursive `\$ref`s in JSON Schema are not supported by Gemini'):
-        await m.agent_model({}, True, [result_tool])
+        await m.agent_model(function_tools=[], allow_text_result=True, result_tools=[result_tool])
 
 
 @dataclass
diff --git a/tests/models/test_model_function.py b/tests/models/test_model_function.py
index 4fdb422e..b4940784 100644
--- a/tests/models/test_model_function.py
+++ b/tests/models/test_model_function.py
@@ -87,7 +87,7 @@ def test_simple(set_event_loop: None):
 
 async def weather_model(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:  # pragma: no cover
     assert info.allow_text_result
-    assert info.function_tools.keys() == {'get_location', 'get_weather'}
+    assert {t.name for t in info.function_tools} == {'get_location', 'get_weather'}
     last = messages[-1]
     if last.role == 'user':
         return ModelStructuredResponse(
@@ -225,7 +225,7 @@ def test_var_args(set_event_loop: None):
 async def call_tool(messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
     if len(messages) == 1:
         assert len(info.function_tools) == 1
-        tool_id = next(iter(info.function_tools.keys()))
+        tool_id = info.function_tools[0].name
         return ModelStructuredResponse(calls=[ToolCall.from_json(tool_id, '{}')])
     else:
         return ModelTextResponse('final response')
@@ -350,7 +350,7 @@ def test_call_all(set_event_loop: None):
 def test_retry_str(set_event_loop: None):
     call_count = 0
 
-    async def try_again(messages: list[Message], _: AgentInfo) -> ModelAnyResponse:
+    async def try_again(msgs_: list[Message], _agent_info: AgentInfo) -> ModelAnyResponse:
         nonlocal call_count
         call_count += 1
 
diff --git a/tests/models/test_vertexai.py b/tests/models/test_vertexai.py
index ff509867..73635945 100644
--- a/tests/models/test_vertexai.py
+++ b/tests/models/test_vertexai.py
@@ -33,7 +33,7 @@ async def test_init_service_account(tmp_path: Path, allow_model_requests: None):
     assert model.url is None
     assert model.auth is None
 
-    await model.agent_model({}, False, None)
+    await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
 
     assert model.url == snapshot(
         'https://us-central1-aiplatform.googleapis.com/v1/projects/my-project-id/locations/us-central1/'
@@ -58,7 +58,7 @@ async def test_init_env(mocker: MockerFixture, allow_model_requests: None):
 
     assert patch.call_count == 0
 
-    await model.agent_model({}, False, None)
+    await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
 
     assert patch.call_count == 1
 
@@ -69,7 +69,7 @@ async def test_init_env(mocker: MockerFixture, allow_model_requests: None):
     assert model.auth is not None
     assert model.name() == snapshot('vertexai:gemini-1.5-flash')
 
-    await model.agent_model({}, False, None)
+    await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
     assert model.url is not None
     assert model.auth is not None
     assert patch.call_count == 1
@@ -83,7 +83,7 @@ async def test_init_right_project_id(tmp_path: Path, allow_model_requests: None)
     assert model.url is None
     assert model.auth is None
 
-    await model.agent_model({}, False, None)
+    await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
 
     assert model.url == snapshot(
         'https://us-central1-aiplatform.googleapis.com/v1/projects/my-project-id/locations/us-central1/'
@@ -99,7 +99,7 @@ async def test_init_service_account_wrong_project_id(tmp_path: Path):
     model = VertexAIModel('gemini-1.5-flash', service_account_file=service_account_path, project_id='different')
 
     with pytest.raises(UserError) as exc_info:
-        await model.agent_model({}, False, None)
+        await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
     assert str(exc_info.value) == snapshot(
         "The project_id you provided does not match the one from service account file: 'different' != 'my-project-id'"
     )
@@ -110,7 +110,7 @@ async def test_init_env_wrong_project_id(mocker: MockerFixture):
     model = VertexAIModel('gemini-1.5-flash', project_id='different')
 
     with pytest.raises(UserError) as exc_info:
-        await model.agent_model({}, False, None)
+        await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
     assert str(exc_info.value) == snapshot(
         "The project_id you provided does not match the one from `google.auth.default()`: 'different' != 'my-project-id'"
     )
@@ -124,7 +124,7 @@ async def test_init_env_no_project_id(mocker: MockerFixture):
     model = VertexAIModel('gemini-1.5-flash')
 
     with pytest.raises(UserError) as exc_info:
-        await model.agent_model({}, False, None)
+        await model.agent_model(function_tools=[], allow_text_result=True, result_tools=[])
     assert str(exc_info.value) == snapshot('No project_id provided and none found in `google.auth.default()`')
 
 
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 18690c3c..b8c5f143 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -2,6 +2,7 @@
 from datetime import timezone
 from typing import Any, Callable, Union
 
+import httpx
 import pytest
 from inline_snapshot import snapshot
 from pydantic import BaseModel
@@ -24,6 +25,7 @@
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.result import Cost, RunResult
+from pydantic_ai.tools import ToolDefinition
 
 from .conftest import IsNow, TestEnv
 
@@ -185,34 +187,33 @@ def test_response_tuple(set_event_loop: None):
     result = agent.run_sync('Hello')
     assert result.data == snapshot(('a', 'a'))
 
-    assert m.agent_model_tools == snapshot({})
+    assert m.agent_model_function_tools == snapshot([])
     assert m.agent_model_allow_text_result is False
 
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 1
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'json_schema', 'outer_typed_dict_key'
-    agent_model_result_tool = {f: getattr(m.agent_model_result_tools[0], f) for f in fields}
-    assert agent_model_result_tool == snapshot(
-        {
-            'name': 'final_result',
-            'description': 'The final response which ends this conversation',
-            'json_schema': {
-                'properties': {
-                    'response': {
-                        'maxItems': 2,
-                        'minItems': 2,
-                        'prefixItems': [{'type': 'string'}, {'type': 'string'}],
-                        'title': 'Response',
-                        'type': 'array',
-                    }
+    assert m.agent_model_result_tools == snapshot(
+        [
+            ToolDefinition(
+                name='final_result',
+                description='The final response which ends this conversation',
+                parameters_json_schema={
+                    'properties': {
+                        'response': {
+                            'maxItems': 2,
+                            'minItems': 2,
+                            'prefixItems': [{'type': 'string'}, {'type': 'string'}],
+                            'title': 'Response',
+                            'type': 'array',
+                        }
+                    },
+                    'required': ['response'],
+                    'type': 'object',
                 },
-                'required': ['response'],
-                'type': 'object',
-            },
-            'outer_typed_dict_key': 'response',
-        }
+                outer_typed_dict_key='response',
+            )
+        ]
     )
 
 
@@ -244,27 +245,28 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
     assert result.data == snapshot('success (no tool calls)')
     assert got_tool_call_name == snapshot(None)
 
-    assert m.agent_model_tools == snapshot({})
+    assert m.agent_model_function_tools == snapshot([])
     assert m.agent_model_allow_text_result is True
 
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 1
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'json_schema', 'outer_typed_dict_key'
-    agent_model_result_tool = {f: getattr(m.agent_model_result_tools[0], f) for f in fields}
-    assert agent_model_result_tool == snapshot(
-        {
-            'name': 'final_result',
-            'description': 'The final response which ends this conversation',
-            'json_schema': {
-                'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'string'}},
-                'title': 'Foo',
-                'required': ['a', 'b'],
-                'type': 'object',
-            },
-            'outer_typed_dict_key': None,
-        }
+    assert m.agent_model_result_tools == snapshot(
+        [
+            ToolDefinition(
+                name='final_result',
+                description='The final response which ends this conversation',
+                parameters_json_schema={
+                    'properties': {
+                        'a': {'title': 'A', 'type': 'integer'},
+                        'b': {'title': 'B', 'type': 'string'},
+                    },
+                    'required': ['a', 'b'],
+                    'title': 'Foo',
+                    'type': 'object',
+                },
+            )
+        ]
     )
 
 
@@ -318,21 +320,18 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
     assert result.data == mod.Foo(a=0, b='a')
     assert got_tool_call_name == snapshot('final_result_Foo')
 
-    assert m.agent_model_tools == snapshot({})
+    assert m.agent_model_function_tools == snapshot([])
     assert m.agent_model_allow_text_result is False
 
     assert m.agent_model_result_tools is not None
     assert len(m.agent_model_result_tools) == 2
 
-    # to match the protocol, we just extract the attributes we care about
-    fields = 'name', 'description', 'json_schema', 'outer_typed_dict_key'
-    agent_model_result_tools = [{f: getattr(t, f) for f in fields} for t in m.agent_model_result_tools]
-    assert agent_model_result_tools == snapshot(
+    assert m.agent_model_result_tools == snapshot(
         [
-            {
-                'name': 'final_result_Foo',
-                'description': 'Foo: The final response which ends this conversation',
-                'json_schema': {
+            ToolDefinition(
+                name='final_result_Foo',
+                description='Foo: The final response which ends this conversation',
+                parameters_json_schema={
                     'properties': {
                         'a': {'title': 'A', 'type': 'integer'},
                         'b': {'title': 'B', 'type': 'string'},
@@ -341,19 +340,17 @@ def validate_result(ctx: RunContext[None], r: Any) -> Any:
                     'title': 'Foo',
                     'type': 'object',
                 },
-                'outer_typed_dict_key': None,
-            },
-            {
-                'name': 'final_result_Bar',
-                'description': 'This is a bar model.',
-                'json_schema': {
+            ),
+            ToolDefinition(
+                name='final_result_Bar',
+                description='This is a bar model.',
+                parameters_json_schema={
                     'properties': {'b': {'title': 'B', 'type': 'string'}},
                     'required': ['b'],
                     'title': 'Bar',
                     'type': 'object',
                 },
-                'outer_typed_dict_key': None,
-            },
+            ),
         ]
     )
 
@@ -388,7 +385,6 @@ async def ret_a(x: str) -> str:
     result2 = agent.run_sync('Hello again', message_history=result1.new_messages())
     assert result2 == snapshot(
         RunResult(
-            data='{"ret_a":"a-apple"}',
             _all_messages=[
                 SystemPrompt(content='Foobar'),
                 UserPrompt(content='Hello', timestamp=IsNow(tz=timezone.utc)),
@@ -400,19 +396,15 @@ async def ret_a(x: str) -> str:
                 ModelTextResponse(content='{"ret_a":"a-apple"}', timestamp=IsNow(tz=timezone.utc)),
                 # second call, notice no repeated system prompt
                 UserPrompt(content='Hello again', timestamp=IsNow(tz=timezone.utc)),
-                ModelStructuredResponse(
-                    calls=[ToolCall(tool_name='ret_a', args=ArgsDict(args_dict={'x': 'a'}))],
-                    timestamp=IsNow(tz=timezone.utc),
-                ),
-                ToolReturn(tool_name='ret_a', content='a-apple', timestamp=IsNow(tz=timezone.utc)),
                 ModelTextResponse(content='{"ret_a":"a-apple"}', timestamp=IsNow(tz=timezone.utc)),
             ],
             _new_message_index=5,
+            data='{"ret_a":"a-apple"}',
             _cost=Cost(),
         )
     )
     new_msg_roles = [msg.role for msg in result2.new_messages()]
-    assert new_msg_roles == snapshot(['user', 'model-structured-response', 'tool-return', 'model-text-response'])
+    assert new_msg_roles == snapshot(['user', 'model-text-response'])
     assert result2.new_messages_json().startswith(b'[{"content":"Hello again",')
 
     # if we pass all_messages, system prompt is NOT inserted before the message_history messages,
@@ -433,11 +425,6 @@ async def ret_a(x: str) -> str:
                 ModelTextResponse(content='{"ret_a":"a-apple"}', timestamp=IsNow(tz=timezone.utc)),
                 # second call, notice no repeated system prompt
                 UserPrompt(content='Hello again', timestamp=IsNow(tz=timezone.utc)),
-                ModelStructuredResponse(
-                    calls=[ToolCall(tool_name='ret_a', args=ArgsDict(args_dict={'x': 'a'}))],
-                    timestamp=IsNow(tz=timezone.utc),
-                ),
-                ToolReturn(tool_name='ret_a', content='a-apple', timestamp=IsNow(tz=timezone.utc)),
                 ModelTextResponse(content='{"ret_a":"a-apple"}', timestamp=IsNow(tz=timezone.utc)),
             ],
             _new_message_index=5,
@@ -534,8 +521,12 @@ async def make_request() -> str:
         # raised a `RuntimeError: Event loop is closed` on repeat runs when we used `asyncio.run()`
         client = cached_async_http_client()
         # use this as I suspect it's about the fastest globally available endpoint
-        response = await client.get('https://cloudflare.com/cdn-cgi/trace')
-        return str(response.status_code)
+        try:
+            response = await client.get('https://cloudflare.com/cdn-cgi/trace')
+        except httpx.ConnectError:
+            pytest.skip('offline')
+        else:
+            return str(response.status_code)
 
     for _ in range(2):
         result = agent.run_sync('Hello')
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 3f2c880d..5449f2e0 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -98,8 +98,8 @@ def test_docs_examples(
     # waiting for https://github.com/pydantic/pytest-examples/issues/43
     if 'import DatabaseConn' in example.source:
         ruff_ignore.append('I001')
-    elif 'async def my_tool(' in example.source:
-        # not sure what's going on here, just ignore it for now
+    elif 'async def my_tool(' in example.source or 'async def only_if_42(' in example.source:
+        # until https://github.com/pydantic/pytest-examples/issues/46 is fixed
         ruff_ignore.append('I001')
 
     line_length = 88
diff --git a/tests/test_logfire.py b/tests/test_logfire.py
index 7e22ee3d..9cc6cfed 100644
--- a/tests/test_logfire.py
+++ b/tests/test_logfire.py
@@ -1,11 +1,10 @@
 from __future__ import annotations as _annotations
 
 from dataclasses import dataclass
-from datetime import timezone
 from typing import Any, Callable
 
 import pytest
-from dirty_equals import IsInt, IsJson, IsNow, IsStr
+from dirty_equals import IsInt, IsJson, IsStr
 from inline_snapshot import snapshot
 from typing_extensions import NotRequired, TypedDict
 
@@ -76,14 +75,16 @@ async def my_ret(x: int) -> str:
                 'id': 0,
                 'message': 'my_agent run prompt=Hello',
                 'children': [
-                    {'id': 1, 'message': 'model request -> model-structured-response'},
+                    {'id': 1, 'message': 'preparing model and tools run_step=1'},
+                    {'id': 2, 'message': 'model request -> model-structured-response'},
                     {
-                        'id': 2,
+                        'id': 3,
                         'message': 'handle model response -> tool-return',
-                        'children': [{'id': 3, 'message': "running tools=['my_ret']"}],
+                        'children': [{'id': 4, 'message': "running tools=['my_ret']"}],
                     },
-                    {'id': 4, 'message': 'model request -> model-text-response'},
-                    {'id': 5, 'message': 'handle model response -> final result'},
+                    {'id': 5, 'message': 'preparing model and tools run_step=2'},
+                    {'id': 6, 'message': 'model request -> model-text-response'},
+                    {'id': 7, 'message': 'handle model response -> final result'},
                 ],
             }
         ]
@@ -101,34 +102,15 @@ async def my_ret(x: int) -> str:
                         'custom_result_text': None,
                         'custom_result_args': None,
                         'seed': 0,
-                        'agent_model_tools': {
-                            'my_ret': {
-                                'function': IsStr(regex='<function test_logfire.<locals>.my_ret at 0x.+>'),
-                                'takes_ctx': False,
-                                'max_retries': 1,
-                                'name': 'my_ret',
-                                'description': '',
-                                '_is_async': True,
-                                '_single_arg_name': None,
-                                '_positional_fields': [],
-                                '_var_positional_field': None,
-                                '_json_schema': {
-                                    'properties': {'x': {'title': 'X', 'type': 'integer'}},
-                                    'required': ['x'],
-                                    'type': 'object',
-                                    'additionalProperties': False,
-                                },
-                                '_current_retry': 0,
-                            }
-                        },
-                        'agent_model_allow_text_result': True,
+                        'agent_model_function_tools': None,
+                        'agent_model_allow_text_result': None,
                         'agent_model_result_tools': None,
                     },
                     'name': 'my_agent',
                     'last_run_messages': None,
                 }
             ),
-            'logfire.null_args': ('custom_model',),
+            'mode_selection': 'from-agent',
             'model_name': 'test-model',
             'agent_name': 'my_agent',
             'logfire.msg_template': '{agent_name} run {prompt=}',
@@ -167,29 +149,10 @@ async def my_ret(x: int) -> str:
                             'title': 'Agent',
                             'x-python-datatype': 'dataclass',
                             'properties': {
-                                'model': {
-                                    'type': 'object',
-                                    'title': 'TestModel',
-                                    'x-python-datatype': 'dataclass',
-                                    'properties': {
-                                        'agent_model_tools': {
-                                            'type': 'object',
-                                            'properties': {
-                                                'my_ret': {
-                                                    'type': 'object',
-                                                    'title': 'Tool',
-                                                    'x-python-datatype': 'dataclass',
-                                                    'properties': {
-                                                        'function': {'type': 'object', 'x-python-datatype': 'unknown'}
-                                                    },
-                                                }
-                                            },
-                                        }
-                                    },
-                                }
+                                'model': {'type': 'object', 'title': 'TestModel', 'x-python-datatype': 'dataclass'}
                             },
                         },
-                        'custom_model': {},
+                        'mode_selection': {},
                         'model_name': {},
                         'agent_name': {},
                         'all_messages': {
@@ -250,48 +213,9 @@ async def my_ret(x: int) -> str:
             'code.function': 'run',
             'code.lineno': IsInt(),
             'run_step': 1,
-            'logfire.msg_template': 'model request {run_step=}',
+            'logfire.msg_template': 'preparing model and tools {run_step=}',
             'logfire.span_type': 'span',
-            'response': IsJson(
-                {
-                    'calls': [{'tool_name': 'my_ret', 'args': {'args_dict': {'x': 0}}, 'tool_id': None}],
-                    'timestamp': IsStr() & IsNow(iso_string=True, tz=timezone.utc),
-                    'role': 'model-structured-response',
-                }
-            ),
-            'cost': IsJson({'request_tokens': None, 'response_tokens': None, 'total_tokens': None, 'details': None}),
-            'logfire.msg': 'model request -> model-structured-response',
-            'logfire.json_schema': IsJson(
-                {
-                    'type': 'object',
-                    'properties': {
-                        'run_step': {},
-                        'response': {
-                            'type': 'object',
-                            'title': 'ModelStructuredResponse',
-                            'x-python-datatype': 'dataclass',
-                            'properties': {
-                                'calls': {
-                                    'type': 'array',
-                                    'items': {
-                                        'type': 'object',
-                                        'title': 'ToolCall',
-                                        'x-python-datatype': 'dataclass',
-                                        'properties': {
-                                            'args': {
-                                                'type': 'object',
-                                                'title': 'ArgsDict',
-                                                'x-python-datatype': 'dataclass',
-                                            }
-                                        },
-                                    },
-                                },
-                                'timestamp': {'type': 'string', 'format': 'date-time'},
-                            },
-                        },
-                        'cost': {'type': 'object', 'title': 'Cost', 'x-python-datatype': 'dataclass'},
-                    },
-                }
-            ),
+            'logfire.msg': 'preparing model and tools run_step=1',
+            'logfire.json_schema': '{"type":"object","properties":{"run_step":{}}}',
         }
     )
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
index 3102d0b6..fe5316e8 100644
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -172,7 +172,7 @@ async def stream_structured_function(
         if len(messages) == 1:
             assert agent_info.function_tools is not None
             assert len(agent_info.function_tools) == 1
-            name = next(iter(agent_info.function_tools))
+            name = agent_info.function_tools[0].name
             first = messages[0]
             assert isinstance(first, UserPrompt)
             json_string = json.dumps({'x': first.content})
diff --git a/tests/test_tools.py b/tests/test_tools.py
index 7f561596..df24e46b 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -1,5 +1,6 @@
 import json
-from typing import Annotated
+from dataclasses import dataclass
+from typing import Annotated, Any, Callable, Union
 
 import pytest
 from inline_snapshot import snapshot
@@ -10,6 +11,7 @@
 from pydantic_ai.messages import Message, ModelAnyResponse, ModelTextResponse
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
+from pydantic_ai.tools import ToolDefinition
 
 
 def test_tool_no_ctx():
@@ -70,8 +72,8 @@ async def google_style_docstring(foo: int, bar: str) -> str:  # pragma: no cover
 
 async def get_json_schema(_messages: list[Message], info: AgentInfo) -> ModelAnyResponse:
     assert len(info.function_tools) == 1
-    r = next(iter(info.function_tools.values()))
-    return ModelTextResponse(json.dumps(r.json_schema))
+    r = info.function_tools[0]
+    return ModelTextResponse(json.dumps(r.parameters_json_schema))
 
 
 def test_docstring_google(set_event_loop: None):
@@ -218,7 +220,6 @@ class Foo(BaseModel):
 def test_takes_just_model(set_event_loop: None):
     agent = Agent()
 
-
     @agent.tool_plain
     def takes_just_model(model: Foo) -> str:
         return f'{model.x} {model.y}'
@@ -241,10 +242,6 @@ def takes_just_model(model: Foo) -> str:
 def test_takes_model_and_int(set_event_loop: None):
     agent = Agent()
 
-    class Foo(BaseModel):
-        x: int
-        y: str
-
     @agent.tool_plain
     def takes_just_model(model: Foo, z: int) -> str:
         return f'{model.x} {model.y} {z}'
@@ -286,7 +283,7 @@ def plain_tool(x: int) -> int:
         call_args.append(x)
         return x + 1
 
-    agent = Agent('test', tools=[Tool(plain_tool, False)], retries=7)
+    agent = Agent('test', tools=[Tool(plain_tool)], retries=7)
     result = agent.run_sync('foobar')
     assert result.data == snapshot('{"plain_tool":1}')
     assert call_args == snapshot([0])
@@ -307,7 +304,7 @@ def ctx_tool(ctx: RunContext[int], x: int) -> int:
 
 # pyright: reportPrivateUsage=false
 def test_init_tool_ctx(set_event_loop: None):
-    agent = Agent('test', tools=[Tool(ctx_tool, True, max_retries=3)], deps_type=int, retries=7)
+    agent = Agent('test', tools=[Tool(ctx_tool, takes_ctx=True, max_retries=3)], deps_type=int, retries=7)
     result = agent.run_sync('foobar', deps=5)
     assert result.data == snapshot('{"ctx_tool":5}')
     assert agent._function_tools['ctx_tool'].takes_ctx is True
@@ -321,7 +318,7 @@ def test_init_tool_ctx(set_event_loop: None):
 
 def test_repeat_tool():
     with pytest.raises(UserError, match="Tool name conflicts with existing tool: 'ctx_tool'"):
-        Agent('test', tools=[Tool(ctx_tool, True), ctx_tool], deps_type=int)
+        Agent('test', tools=[Tool(ctx_tool), ctx_tool], deps_type=int)
 
 
 def test_tool_return_conflict():
@@ -340,12 +337,12 @@ def plain_tool(x: int) -> int:  # pragma: no cover
 
     m = r'First parameter of tools that take context must be annotated with RunContext\[\.\.\.\]'
     with pytest.raises(UserError, match=m):
-        Tool(plain_tool, True)
+        Tool(plain_tool, takes_ctx=True)
 
 
 def test_init_plain_tool_invalid():
     with pytest.raises(UserError, match='RunContext annotations can only be used with tools that take context'):
-        Tool(ctx_tool, False)
+        Tool(ctx_tool, takes_ctx=False)
 
 
 def test_return_pydantic_model(set_event_loop: None):
@@ -393,3 +390,79 @@ def return_pydantic_model() -> Foobar:
 
     with pytest.raises(PydanticSerializationError, match='Unable to serialize unknown type:'):
         agent.run_sync('')
+
+
+def test_dynamic_cls_tool(set_event_loop: None):
+    @dataclass
+    class MyTool(Tool[int]):
+        spam: int
+
+        def __init__(self, spam: int = 0, **kwargs: Any):
+            self.spam = spam
+            kwargs.update(function=self.tool_function, takes_ctx=False)
+            super().__init__(**kwargs)
+
+        def tool_function(self, x: int, y: str) -> str:
+            return f'{self.spam} {x} {y}'
+
+        async def prepare_tool_def(self, ctx: RunContext[int]) -> Union[ToolDefinition, None]:
+            if ctx.deps != 42:
+                return await super().prepare_tool_def(ctx)
+
+    agent = Agent('test', tools=[MyTool(spam=777)], deps_type=int)
+    r = agent.run_sync('', deps=1)
+    assert r.data == snapshot('{"tool_function":"777 0 a"}')
+
+    r = agent.run_sync('', deps=42)
+    assert r.data == snapshot('success (no tool calls)')
+
+
+def test_dynamic_plain_tool_decorator(set_event_loop: None):
+    agent = Agent('test', deps_type=int)
+
+    async def prepare_tool_def(ctx: RunContext[int], tool_def: ToolDefinition) -> Union[ToolDefinition, None]:
+        if ctx.deps != 42:
+            return tool_def
+
+    @agent.tool_plain(prepare=prepare_tool_def)
+    def foobar(x: int, y: str) -> str:
+        return f'{x} {y}'
+
+    r = agent.run_sync('', deps=1)
+    assert r.data == snapshot('{"foobar":"0 a"}')
+
+    r = agent.run_sync('', deps=42)
+    assert r.data == snapshot('success (no tool calls)')
+
+
+def test_dynamic_tool_decorator(set_event_loop: None):
+    agent = Agent('test', deps_type=int)
+
+    async def prepare_tool_def(ctx: RunContext[int], tool_def: ToolDefinition) -> Union[ToolDefinition, None]:
+        if ctx.deps != 42:
+            return tool_def
+
+    @agent.tool(prepare=prepare_tool_def)
+    def foobar(ctx: RunContext[int], x: int, y: str) -> str:
+        return f'{ctx.deps} {x} {y}'
+
+    r = agent.run_sync('', deps=1)
+    assert r.data == snapshot('{"foobar":"1 0 a"}')
+
+    r = agent.run_sync('', deps=42)
+    assert r.data == snapshot('success (no tool calls)')
+
+
+def test_future_run_context(set_event_loop: None, create_module: Callable[[str], Any]):
+    mod = create_module("""
+from __future__ import annotations
+
+from pydantic_ai import Agent, RunContext
+
+def ctx_tool(ctx: RunContext[int], x: int) -> int:
+    return x + ctx.deps
+
+agent = Agent('test', tools=[ctx_tool], deps_type=int)
+    """)
+    result = mod.agent.run_sync('foobar', deps=5)
+    assert result.data == snapshot('{"ctx_tool":5}')
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 25e13c65..0bcf88ab 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -7,7 +7,7 @@
 from inline_snapshot import snapshot
 
 from pydantic_ai import UserError
-from pydantic_ai._utils import check_object_json_schema, group_by_temporal
+from pydantic_ai._utils import Either, check_object_json_schema, group_by_temporal
 
 pytestmark = pytest.mark.anyio
 
@@ -43,3 +43,10 @@ def test_check_object_json_schema():
     array_schema = {'type': 'array', 'items': {'type': 'string'}}
     with pytest.raises(UserError, match='^Schema must be an object$'):
         check_object_json_schema(array_schema)
+
+
+def test_either():
+    assert repr(Either[int, int](left=123)) == 'Either(left=123)'
+    assert Either(left=123).whichever() == 123
+    assert repr(Either[int, int](right=456)) == 'Either(right=456)'
+    assert Either(right=456).whichever() == 456
diff --git a/tests/typed_agent.py b/tests/typed_agent.py
index 1daa92c3..8cb7ebec 100644
--- a/tests/typed_agent.py
+++ b/tests/typed_agent.py
@@ -7,6 +7,7 @@
 
 from pydantic_ai import Agent, ModelRetry, RunContext, Tool
 from pydantic_ai.result import RunResult
+from pydantic_ai.tools import ToolDefinition
 
 
 @dataclass
@@ -55,6 +56,30 @@ async def ok_tool(ctx: RunContext[MyDeps], x: str) -> str:
 assert_type(ok_tool, Callable[[RunContext[MyDeps], str], str])  # type: ignore[assert-type]
 
 
+async def prep_ok(ctx: RunContext[MyDeps], tool_def: ToolDefinition) -> ToolDefinition | None:
+    if ctx.deps.foo == 42:
+        return None
+    else:
+        return tool_def
+
+
+@typed_agent.tool(prepare=prep_ok)
+def ok_tool_prepare(ctx: RunContext[MyDeps], x: int, y: str) -> str:
+    return f'{ctx.deps.foo} {x} {y}'
+
+
+async def prep_wrong_type(ctx: RunContext[int], tool_def: ToolDefinition) -> ToolDefinition | None:
+    if ctx.deps == 42:
+        return None
+    else:
+        return tool_def
+
+
+@typed_agent.tool(prepare=prep_wrong_type)  # type: ignore[arg-type]
+def wrong_tool_prepare(ctx: RunContext[MyDeps], x: int, y: str) -> str:
+    return f'{ctx.deps.foo} {x} {y}'
+
+
 @typed_agent.tool_plain
 def ok_tool_plain(x: str) -> dict[str, str]:
     return {'x': x}
@@ -161,20 +186,42 @@ def foobar_plain(x: str, y: int) -> str:
     return f'{x} {y}'
 
 
-Tool(foobar_ctx, True)
-Tool(foobar_plain, False)
+Tool(foobar_ctx, takes_ctx=True)
+Tool(foobar_ctx)
+Tool(foobar_plain, takes_ctx=False)
+Tool(foobar_plain)
 
 # unfortunately we can't type check these cases, since from a typing perspect `foobar_ctx` is valid as a plain tool
-Tool(foobar_ctx, False)
-Tool(foobar_plain, True)
+Tool(foobar_ctx, takes_ctx=False)
+Tool(foobar_plain, takes_ctx=True)
 
 Agent('test', tools=[foobar_ctx], deps_type=int)
 Agent('test', tools=[foobar_plain], deps_type=int)
 Agent('test', tools=[foobar_plain])
-Agent('test', tools=[Tool(foobar_ctx, True)], deps_type=int)
-Agent('test', tools=[Tool(foobar_plain, False)], deps_type=int)
-Agent('test', tools=[Tool(foobar_ctx, True), foobar_ctx, foobar_plain], deps_type=int)
+Agent('test', tools=[Tool(foobar_ctx)], deps_type=int)
+Agent('test', tools=[Tool(foobar_plain)], deps_type=int)
+Agent('test', tools=[Tool(foobar_ctx), foobar_ctx, foobar_plain], deps_type=int)
 
 Agent('test', tools=[foobar_ctx], deps_type=str)  # pyright: ignore[reportArgumentType]
 Agent('test', tools=[foobar_ctx])  # pyright: ignore[reportArgumentType]
-Agent('test', tools=[Tool(foobar_ctx, True)])  # pyright: ignore[reportArgumentType]
+Agent('test', tools=[Tool(foobar_ctx)])  # pyright: ignore[reportArgumentType]
+
+# prepare example from docs:
+
+
+def greet(name: str) -> str:
+    return f'hello {name}'
+
+
+async def prepare_greet(ctx: RunContext[str], tool_def: ToolDefinition) -> ToolDefinition | None:
+    d = f'Name of the {ctx.deps} to greet.'
+    tool_def.parameters_json_schema['properties']['name']['description'] = d
+    return tool_def
+
+
+greet_tool = Tool(greet, prepare=prepare_greet)
+assert_type(greet_tool, Tool[str])
+greet_agent = Agent[str, str]('test', tools=[greet_tool], deps_type=str)
+
+result = greet_agent.run_sync('testing...', deps='human')
+assert result.data == '{"greet":"hello a"}'