pydantic/pydantic-ai
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai.git synced 2026-10-12 02:21:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

Re-resolve tools when capability availability changes mid-step so prepare_tools still governs them (#8071)

Browse Source
This commit is contained in:
Douwe Maan
2026-09-16 19:50:58 -06:00
committed by GitHub
parent 97035aece8
commit 3732f445fa
8 changed files with 555 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/references/ON-DEMAND-CAPABILITIES.md
+2 -2
View File
@@ -125,9 +125,9 @@ When `load_capability` succeeds:
Use `ctx.is_tool_available(tool_def)` when a wrapping toolset needs to decide whether a definition it holds is currently visible. The definition form remains reliable inside `get_tools`; the name form looks in the current resolved `ctx.tools` snapshot and is intended for model-request hooks and tool execution.
Message history matters. Loaded capability state is reconstructed from matching `LoadCapabilityCallPart` and `LoadCapabilityReturnPart` pairs, while revealed function-tool state is reconstructed from `ToolAvailabilityDeltaPart` entries. A history processor must preserve the deltas or the complete capability-load pairs from which Pydantic AI can reconstruct them. If it removes both representations, those tools become hidden again.
Message history matters. Loaded capability state is reconstructed from matching `LoadCapabilityCallPart` and `LoadCapabilityReturnPart` pairs, while revealed function-tool state is reconstructed from `ToolAvailabilityDeltaPart` entries. A history processor must preserve the deltas or the complete capability-load pairs from which Pydantic AI can reconstruct them. If it removes both representations, those tools become hidden again. A processor that *adds* a complete load pair activates that capability for the request it is processing, not the next one: its `prepare_tools` runs before any of its tools are dispatched, so a capability using `prepare_tools` as a permission filter still governs them.
A `CompactionPart` resets both forms of prospective derived state at its exact position, so future requests load and reveal capability tools again. For a call in the response currently being dispatched, pre-boundary evidence still counts when the serving provider did not honor that boundary on the request wire; otherwise a call without visible load evidence is refused with a "not available yet" retry naming the capability to load.
A `CompactionPart` resets both forms of prospective derived state at its exact position, so future requests load and reveal capability tools again. For a call in the response currently being dispatched, pre-boundary evidence still counts when the serving provider did not honor that boundary on the request wire; otherwise a call without visible load evidence is refused with a "not available yet" retry naming the capability to load. When that pre-boundary evidence does admit the call, the owning capability counts as active for that dispatch, so its `prepare_tools` and tool hooks run over the call — a capability filtering its own tools governs them on whatever evidence authorized them. `ctx.loaded_capability_ids` still reads the conservative window, so it can be empty there.
## Dynamic Descriptions and Instructions
pydantic_ai_slim/pydantic_ai/_agent_graph.py
+23 -4
View File
@@ -433,8 +433,10 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
# identity survives `replace(ctx, ...)`, which shallow-copies) and are only ever mutated in
# place — never reassigned. The per-step refresh relies on that shared identity for both, and
# `discovered_tool_names` additionally on the in-step reveals written by tool execution.
# `loaded_capability_ids` is refreshed from history only: a capability loaded during a step
# lands from the next one, so nothing writes it mid-step. Reassigning either (here, or by
# `loaded_capability_ids` is refreshed from history only: a capability the *model* loads during
# a step lands from the next one, since its load return only reaches history at the step's end.
# It is still refreshed a second time within the step, after history processing has rewritten
# that history — the only way its contents move mid-step. Reassigning either (here, or by
# passing it to a `replace(ctx, ...=...)`) would silently break in-step tool reveals.
loaded_capability_ids: set[str]
discovered_tool_names: set[str]
@@ -1647,6 +1649,14 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
ctx.deps.resumed_request_index = shifted if shifted >= 0 else None
# `ctx.state.message_history` is the same list used by `capture_run_messages`, so we should replace its contents, not the reference
ctx.state.message_history[:] = messages
# Processing may have added or removed `load_capability` exchanges, and the durable history it
# just rewrote is what the rest of the step reads availability from. Refresh so the execution
# gate agrees with the reveal state `_with_outgoing_reveal_state` derives below from the same
# messages; `ToolManager.for_run_step` re-resolves off this set at dispatch, so a capability
# that became active here still governs its own tools through `prepare_tools`.
_refresh_loaded_capability_ids(ctx)
# Update the new message index to ensure `result.new_messages()` returns the correct messages
ctx.deps.new_message_index = _first_new_message_index(
messages,
@@ -1855,6 +1865,12 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
# replace its contents (dropping the suspended response) rather than the reference;
# `_finish_handling` then appends the final merged response after the base history.
ctx.state.message_history[:] = base_messages
# Same reason as in `_prepare_request`: processing may have changed which capabilities the
# durable history shows as loaded, and the tool calls this continuation comes back with are
# dispatched against that history.
_refresh_loaded_capability_ids(ctx)
ctx.deps.new_message_index = _first_new_message_index(
base_messages,
ctx.state.run_id,
@@ -2256,8 +2272,11 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
# This will raise errors for any tool name conflicts
ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
# The manager was already prepared for this same run step before the model request, so
# `for_run_step` deliberately returns it unchanged, keeping the retries it accumulated —
# which is why the evidence lands field by field rather than by swapping in `run_context`.
# `for_run_step` normally returns it unchanged, keeping the retries it accumulated — which is
# why the evidence lands field by field rather than by swapping in `run_context`. (It does
# re-resolve when capability availability moved since preparation, e.g. a history processor
# injected a `load_capability` exchange; that path carries the same retries through and ends
# up holding `run_context`, whose evidence this assignment then re-applies harmlessly.)
# Only the retrospective evidence is carried: replacing the prospective shared sets would
# affect the next request's reveal pruning and search ranking.
assert ctx.deps.tool_manager.ctx is not None
pydantic_ai_slim/pydantic_ai/_run_context.py
+32 -5
View File
@@ -360,7 +360,8 @@ class RunContext(Generic[RunContextAgentDepsT]):
_anchored_evidence: AnchoredEvidence = field(default_factory=lambda: AnchoredEvidence(), repr=False)
"""Evidence the serving provider could still see that the conservative window dropped.
Set at tool-call dispatch and read only by `is_tool_available`. Private because the sets above
Set at tool-call dispatch and read only through `_dispatch_active_capability_ids` (plus the
reveal half, read by `is_tool_available` directly). Private because the sets above
stay the answer for everything that feeds a *future* request, whose provider isn't knowable yet;
this one is the answer for a call the model has already made, where it is. See `AnchoredEvidence`.
"""
@@ -496,6 +497,32 @@ class RunContext(Generic[RunContextAgentDepsT]):
id for id, cap in self.capabilities.items() if cap.defer_loading is not True
} | self.loaded_capability_ids
@property
def _dispatch_active_capability_ids(self) -> set[str]:
"""`active_capability_ids`, widened by the anchored evidence for the response being dispatched.
The single answer to "may this capability act on the call being dispatched right now?", and
it has to be single: `is_tool_available` authorizes the call from it, and the `prepare_tools`
dispatch gate decides from it whether the owning capability's filter runs over that tool. If
only the first consulted the evidence, a tool authorized through a load the conservative
window dropped would execute with its owner's `prepare_tools` never having run — the
capability is not active, so nothing dispatched to it.
The evidence is narrowed to the run's configured deferred ids — the shape every load record
has, since only a deferred capability can be loaded. Inert for both predicates above, which
look up an id that came from a registered capability either way; it is there so that a
history naming a capability this run no longer configures doesn't leave a permanent
difference in `ToolManager.resolved_capability_ids` and rebuild the tools every dispatch.
`_deferred_capability_ids` rather than `capabilities` deliberately: it crosses the Temporal
activity boundary, where reading the live registry raises.
Outside tool-call dispatch `_anchored_evidence` is empty, so this is exactly
`active_capability_ids`.
"""
return self.active_capability_ids | (
self._anchored_evidence.loaded_capability_ids & self._deferred_capability_ids
)
@property
@deprecated(
'`available_capability_ids` is deprecated, use `active_capability_ids` instead: for a '
@@ -513,8 +540,8 @@ class RunContext(Generic[RunContextAgentDepsT]):
def _deferred_capability_ids(self) -> set[str]:
"""IDs of the capabilities configured to load on demand.
Private, and read only by `is_tool_available`, which needs the *configured* shape rather
than the runtime one: `loaded_capability_ids` records what history says was loaded, which
Private, and read only by `is_tool_available` and `_dispatch_active_capability_ids`, which
need the *configured* shape rather than the runtime one: `loaded_capability_ids` records what history says was loaded, which
can name a capability that has since been reconfigured as always-on. Overridden in
`TemporalRunContext` with the snapshot serialized at activity dispatch, since the
`capabilities` registry this reads does not cross that boundary.
@@ -591,7 +618,7 @@ class RunContext(Generic[RunContextAgentDepsT]):
and capability_id in self._deferred_capability_ids
and capability_id in self.loaded_capability_ids | evidence.loaded_capability_ids
):
return capability_id in self.active_capability_ids | evidence.loaded_capability_ids
return capability_id in self._dispatch_active_capability_ids
if tool_def.name not in self.discovered_tool_names | evidence.discovered_tool_names:
return False
# A run holds to load, then reveal, then call. `discovered_tool_names` is raw history
@@ -599,7 +626,7 @@ class RunContext(Generic[RunContextAgentDepsT]):
# never loaded — a history no real run produces, and one that would skip the instructions
# written to be read first. Checking the owner here keeps this predicate in step with what
# `ToolManager` will run, so "available" means one thing everywhere it is asked.
return capability_id is None or capability_id in (self.active_capability_ids | evidence.loaded_capability_ids)
return capability_id is None or capability_id in self._dispatch_active_capability_ids
@property
def tools(self) -> dict[str, ToolDefinition]:
pydantic_ai_slim/pydantic_ai/capabilities/combined.py
+6 -1
View File
@@ -1109,6 +1109,11 @@ def _capability_active(capability: AbstractCapability[AgentDepsT], ctx: RunConte
Activity, not loading: an always-on capability is active for the whole run without ever being
loaded, which is why the deferred branch below is the only one that consults history.
During tool-call dispatch that history question is answered against the provider that served the
response, via `_dispatch_active_capability_ids` — the same set `RunContext.is_tool_available`
authorizes the call from. Reading the narrower `active_capability_ids` here would let a tool the
gate admits on anchored evidence execute without its owner's `prepare_tools` ever running.
"""
if capability.defer_loading is not True:
return True
@@ -1116,4 +1121,4 @@ def _capability_active(capability: AbstractCapability[AgentDepsT], ctx: RunConte
# Deferred capabilities are required to have an explicit `id` (enforced in
# `_build_run_capabilities`), which is also the key they're registered under, so we read
# it directly rather than resolving the instance back to its run-local registry id.
return capability.id is not None and capability.id in ctx.active_capability_ids
return capability.id is not None and capability.id in ctx._dispatch_active_capability_ids # pyright: ignore[reportPrivateUsage]
pydantic_ai_slim/pydantic_ai/tool_manager.py
+56 -14
View File
@@ -167,6 +167,25 @@ class ToolManager(Generic[AgentDepsT]):
"""
default_max_retries: int = 1
"""Default number of times to retry a tool"""
resolved_capability_ids: frozenset[str] | None = None
"""Snapshot of the capability-activity set the cached `tools` were resolved against.
Availability is an input to tool *resolution*, not just to the execution gate: a deferred
capability only has its `prepare_tools` dispatched once it's active, so a tool set resolved
while its owner was inactive is ungoverned by that capability's filter. Caching on `run_step`
alone would hand that ungoverned set to a dispatch that has since become available — a
permission filter that silently doesn't run. Comparing this against the incoming context makes
`for_run_step` re-resolve when availability moved mid-step instead.
Availability moves two ways within a step, and both have to be in the key: history processing
rewrites the loaded set feeding
[`active_capability_ids`][pydantic_ai.tools.RunContext.active_capability_ids], and dispatch
anchors its evidence to the provider that served the response, which can see a load the
conservative window dropped. So this snapshots the union the execution gate itself authorizes
from, not `active_capability_ids` alone.
`None` before the manager has been prepared for a run step.
"""
@classmethod
@contextmanager
@@ -186,25 +205,43 @@ class ToolManager(Generic[AgentDepsT]):
_parallel_execution_mode_ctx_var.reset(token)
async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
"""Build a new tool manager for the next run step, carrying over the retries from the current run step."""
"""Build a new tool manager for the next run step, carrying over the retries from the current run step.
Re-resolves within a step as well, when capability availability moved since the cached tools
were resolved (see `resolved_capability_ids`), so a capability that became active mid-step
still gets its `prepare_tools` say over its own tools.
"""
resolved_capability_ids = frozenset(ctx._dispatch_active_capability_ids) # pyright: ignore[reportPrivateUsage]
same_step = False
if self.ctx is not None:
if ctx.run_step == self.ctx.run_step:
return self
if resolved_capability_ids == self.resolved_capability_ids:
return self
retries = {
tool_name: count
for tool_name, count in self.ctx.retries.items()
if tool_name not in self.succeeded_tools
}
retries.update(
{
failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1
for failed_tool_name in self.failed_tools
# Same step, so the retry carry-over below must not run again: it increments every
# failed tool's count, and doing that twice within one step would burn two units of
# a tool's budget for a single failure. The step's counts stand as they are.
same_step = True
ctx = replace(ctx, retries=self.ctx.retries)
else:
retries = {
tool_name: count
for tool_name, count in self.ctx.retries.items()
if tool_name not in self.succeeded_tools
}
)
ctx = replace(ctx, retries=retries)
retries.update(
{
failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1
for failed_tool_name in self.failed_tools
}
)
ctx = replace(ctx, retries=retries)
toolset = await self.toolset.for_run_step(ctx)
# `AbstractToolset.for_run_step` marks a step boundary, not a re-resolution: `DynamicToolset`
# exits and re-enters its inner toolset there, re-running the factory and discarding whatever
# state the entered toolset held. Within a step we reuse the already-transitioned instance and
# only re-run `get_tools`, which is what re-runs `prepare_tools`.
toolset = self.toolset if same_step else await self.toolset.for_run_step(ctx)
new_tm = self.__class__(
toolset=toolset,
@@ -213,6 +250,11 @@ class ToolManager(Generic[AgentDepsT]):
tools=await toolset.get_tools(ctx),
default_max_retries=self.default_max_retries,
availability_refused=self.availability_refused,
resolved_capability_ids=resolved_capability_ids,
# Per-step accumulators: they belong to the step being rebuilt, and dropping them would
# under-count the retries the next step's carry-over derives from them.
failed_tools=self.failed_tools if same_step else set[str](),
succeeded_tools=self.succeeded_tools if same_step else set[str](),
)
# Make the prepared ToolManager accessible from RunContext so that
# wrapper toolsets (e.g. CodeModeToolset) can dispatch tool calls
tests/test_capability_deferred_loading.py
+235
View File
@@ -1249,6 +1249,241 @@ async def test_processed_history_determines_request_reveal_state() -> None:
assert seen == [set()]
def _inject_load(capability_id: str) -> Callable[[list[ModelMessage]], list[ModelMessage]]:
"""Build a processor that injects a complete `load_capability` exchange for `capability_id`."""
def processor(messages: list[ModelMessage]) -> list[ModelMessage]:
if any(True for _ in iter_message_parts(messages, ModelResponse, LoadCapabilityCallPart)):
return messages
return [
messages[0],
ModelResponse(parts=[LoadCapabilityCallPart(args={'id': capability_id}, tool_call_id='injected')]),
ModelRequest(parts=[LoadCapabilityReturnPart(content={}, tool_call_id='injected')]),
*messages[1:],
]
return processor
def _secret_op_returns(messages: list[ModelMessage]) -> list[object]:
return [
part.content
for part in iter_message_parts(messages, ModelRequest, ToolReturnPart)
if part.tool_name == 'secret_op'
]
def _call_secret_op_once(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
"""Call `secret_op` on the first request, then finish however that call was answered."""
answered = any(
part.tool_name == 'secret_op' for part in iter_message_parts(messages, ModelRequest, ToolReturnPart)
) or any(True for _ in iter_message_parts(messages, ModelRequest, RetryPromptPart))
if answered:
return ModelResponse(parts=[TextPart('done')])
return ModelResponse(parts=[ToolCallPart('secret_op', {}, tool_call_id='call-1')])
@pytest.mark.parametrize('inject_load', [False, True])
async def test_processor_injected_load_lets_capability_prepare_tools_govern_its_tools(inject_load: bool) -> None:
"""A capability made active by history processing still has `prepare_tools` over its own tools.
Availability is an input to tool *resolution*, not just to the execution gate: the tool set
resolved at the start of the step was built while the capability was inactive, so its
`prepare_tools` never ran and its tools sat in that set ungoverned. Caching the resolved set on
`run_step` alone would then hand that set to a dispatch that reads the (now changed)
availability live, and a `prepare_tools` used as a permission filter would be silently bypassed.
Both directions are pinned: with no injection the tool is refused as not-yet-available (the
filter is never reached, and must not need to be), and with the injection it is refused because
the filter removed it — never executed either way.
"""
prepare_tools_calls: list[list[str]] = []
loaded_ids_seen: list[list[str]] = []
class FilteringSecrets(Capability[object]):
"""Uses `prepare_tools` as a permission filter over its own tool."""
async def prepare_tools(self, ctx: RunContext[object], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
prepare_tools_calls.append(sorted(tool_def.name for tool_def in tool_defs))
loaded_ids_seen.append(sorted(ctx.loaded_capability_ids))
return [tool_def for tool_def in tool_defs if tool_def.name != 'secret_op']
secrets_toolset = FunctionToolset[object]()
@secrets_toolset.tool_plain
def secret_op() -> str: # pragma: no cover
return 'EXECUTED'
capabilities: list[AbstractCapability[object]] = [
FilteringSecrets(id='secrets', description='Secret tools.', toolsets=[secrets_toolset], defer_loading=True)
]
if inject_load:
capabilities.append(ProcessHistory[object](processor=_inject_load('secrets')))
agent = Agent(FunctionModel(_call_secret_op_once), capabilities=capabilities)
result = await agent.run('call secret_op')
assert _secret_op_returns(result.all_messages()) == []
refusals = [str(part.content) for part in iter_message_parts(result.all_messages(), ModelRequest, RetryPromptPart)]
if inject_load:
# The capability is active, so its own filter decides — and it removed the tool.
assert prepare_tools_calls[0] == snapshot(['secret_op'])
# The run's own prospective state agrees with the history it is sending, so the documented
# "has the runbook been loaded?" check (`id in ctx.loaded_capability_ids`) answers yes on
# this step, rather than the capability being active only through dispatch-time evidence.
assert loaded_ids_seen[0] == snapshot(['secrets'])
assert refusals == snapshot(["Unknown tool name: 'secret_op'. Available tools: 'load_capability'"])
else:
# The capability never loaded, so the availability gate refuses before any filter is reached.
assert prepare_tools_calls == snapshot([])
assert refusals == snapshot(
[
"Tool 'secret_op' is not available yet: it belongs to capability 'secrets'. Call `load_capability` for it first, then call the tool again once you've read the capability's instructions."
]
)
async def test_processor_injected_load_makes_capability_tool_callable() -> None:
"""Without a filter, a processor-injected load makes the capability's tool callable that same step.
The mirror of the test above, and the reason availability is refreshed after processing at all:
the reveal state the request ships is derived from the processed history, so the execution gate
has to read the same history or the model is offered a tool that comes back refused.
"""
secrets_toolset = FunctionToolset[object]()
@secrets_toolset.tool_plain
def secret_op() -> str:
return 'EXECUTED'
agent = Agent(
FunctionModel(_call_secret_op_once),
capabilities=[
Capability[object](
id='secrets', description='Secret tools.', toolsets=[secrets_toolset], defer_loading=True
),
ProcessHistory[object](processor=_inject_load('secrets')),
],
)
result = await agent.run('call secret_op')
assert _secret_op_returns(result.all_messages()) == snapshot(['EXECUTED'])
assert [
str(part.content) for part in iter_message_parts(result.all_messages(), ModelRequest, RetryPromptPart)
] == snapshot([])
async def test_processor_removed_load_leaves_advertisement_and_gate_in_agreement() -> None:
"""Refreshing availability from processed history must not advertise a tool the gate will refuse.
The mirror of the injection tests, and the direction where making the request and the gate agree
could have gone wrong: the refresh takes the capability *out* of the loaded set for a step whose
tool set was resolved while it was in. `_with_outgoing_reveal_state` derives the request's reveal
state from the same processed messages, so the tool is withheld from the wire rather than shipped
and then refused — the model is never offered something `ToolManager` would reject.
"""
secrets_toolset = FunctionToolset[object]()
@secrets_toolset.tool_plain
def secret_op() -> str: # pragma: no cover
return 'EXECUTED'
advertised: list[list[str]] = []
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
advertised.append(sorted(tool.name for tool in info.function_tools))
return _call_secret_op_once(messages, info)
def drop_load_exchange(messages: list[ModelMessage]) -> list[ModelMessage]:
kept: list[ModelMessage] = []
for message in messages:
parts = [
part
for part in message.parts
if not isinstance(part, (LoadCapabilityCallPart, LoadCapabilityReturnPart))
]
if parts:
kept.append(replace(message, parts=parts))
return kept
agent = Agent(
FunctionModel(model_fn),
capabilities=[
Capability[object](
id='secrets', description='Secret tools.', toolsets=[secrets_toolset], defer_loading=True
),
ProcessHistory[object](processor=drop_load_exchange),
],
)
result = await agent.run(
'call secret_op',
message_history=[
ModelRequest(parts=[UserPromptPart(content='load it')]),
ModelResponse(parts=[LoadCapabilityCallPart(args={'id': 'secrets'}, tool_call_id='l1')]),
ModelRequest(parts=[LoadCapabilityReturnPart(content={}, tool_call_id='l1')]),
],
)
# Withheld from the wire, so the model was never offered it...
assert advertised[0] == snapshot(['load_capability'])
# ...and the call the stub makes anyway is refused, rather than running ungoverned.
assert _secret_op_returns(result.all_messages()) == []
assert [
str(part.content) for part in iter_message_parts(result.all_messages(), ModelRequest, RetryPromptPart)
] == snapshot(
[
"Tool 'secret_op' is not available yet: it belongs to capability 'secrets'. Call `load_capability` for it first, then call the tool again once you've read the capability's instructions."
]
)
async def test_processor_injected_load_is_governed_when_resuming_a_suspended_response() -> None:
"""The same governance holds on the request that resumes a provider-suspended turn.
That path prepares its request separately from the normal one — the history already ends in the
suspended response, so there is no new `ModelRequest` to append — but it runs the same history
processing and dispatches the continuation's tool calls against the result, so it needs the same
post-processing availability refresh.
"""
prepare_tools_calls: list[list[str]] = []
loaded_ids_seen: list[list[str]] = []
class FilteringSecrets(Capability[object]):
"""Uses `prepare_tools` as a permission filter over its own tool."""
async def prepare_tools(self, ctx: RunContext[object], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
prepare_tools_calls.append(sorted(tool_def.name for tool_def in tool_defs))
loaded_ids_seen.append(sorted(ctx.loaded_capability_ids))
return [tool_def for tool_def in tool_defs if tool_def.name != 'secret_op']
secrets_toolset = FunctionToolset[object]()
@secrets_toolset.tool_plain
def secret_op() -> str: # pragma: no cover
return 'EXECUTED'
agent = Agent(
FunctionModel(_call_secret_op_once),
capabilities=[
FilteringSecrets(id='secrets', description='Secret tools.', toolsets=[secrets_toolset], defer_loading=True),
ProcessHistory[object](processor=_inject_load('secrets')),
],
)
result = await agent.run(
message_history=[
ModelRequest(parts=[UserPromptPart(content='call secret_op')]),
ModelResponse(parts=[TextPart('paused')], state='suspended'),
]
)
assert _secret_op_returns(result.all_messages()) == []
assert prepare_tools_calls[0] == snapshot(['secret_op'])
assert loaded_ids_seen[0] == snapshot(['secrets'])
assert [
str(part.content) for part in iter_message_parts(result.all_messages(), ModelRequest, RetryPromptPart)
] == snapshot(["Unknown tool name: 'secret_op'. Available tools: 'load_capability'"])
async def test_orphaned_reveal_evidence_stripped_by_cleanup_does_not_count_as_revealed() -> None:
"""Evidence orphaned by a history processor is stripped before reveal derivation.
tests/test_tool_availability.py
+73 -1
View File
@@ -43,7 +43,7 @@ from pydantic_ai.messages import (
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import RunContext
from pydantic_ai.tools import RunContext, ToolDefinition
from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai.toolsets._deferred_capability_loader import (
LOAD_CAPABILITY_TOOL_NAME,
@@ -154,6 +154,78 @@ async def test_capability_tool_call_uses_serving_providers_wire_window(boundary:
assert result.output == 'done'
@pytest.mark.parametrize(('boundary', 'provider_name'), _INVALID_WIRE_BOUNDARIES)
async def test_capability_prepare_tools_governs_a_tool_the_wire_window_admits(
boundary: CompactionPart, provider_name: str
):
"""A tool admitted on anchored evidence still answers to its owner's `prepare_tools`.
The mirror of `test_capability_tool_call_uses_serving_providers_wire_window`, which pins that
such a call is *not* refused. Admitting it is only half an answer: authorization and filtering
have to read the same set. The conservative window drops the load, so `active_capability_ids`
alone reports the capability inactive and nothing dispatches to it — while the gate, reading the
provider-exact window, admits its tool. A capability using `prepare_tools` as a permission
filter would be bypassed, the same fail-open as a processor-injected load reached from the other
direction.
The test explicitly stamps `ModelResponse.provider_name`; a bare `FunctionModel` would exercise
the missing-provenance fallback instead, where the two windows agree and nothing is admitted.
"""
prepare_tools_calls: list[list[str]] = []
loaded_ids_seen: list[list[str]] = []
class FilteringGuarded(Capability[Any]):
"""Uses `prepare_tools` as a permission filter over its own tool."""
async def prepare_tools(self, ctx: RunContext[Any], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
prepare_tools_calls.append(sorted(tool_def.name for tool_def in tool_defs))
loaded_ids_seen.append(sorted(ctx.loaded_capability_ids))
return [tool_def for tool_def in tool_defs if tool_def.name != 'guarded_tool']
def guarded_tool() -> str: # pragma: no cover
return 'ran'
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
if any(True for _ in iter_message_parts(messages, ModelRequest, RetryPromptPart)):
return _provider_response([make_text_response('done').parts[0]], provider_name)
return _provider_response([ToolCallPart(tool_name='guarded_tool', args={}, tool_call_id='g1')], provider_name)
capability = FilteringGuarded(
id='guarded', description='Guarded.', toolsets=[FunctionToolset([guarded_tool])], defer_loading=True
)
agent = Agent(FunctionModel(model_fn), capabilities=[capability])
history = [
ModelResponse(
parts=[LoadCapabilityCallPart(args={'id': 'guarded'}, tool_call_id='load')], provider_name=provider_name
),
ModelRequest(
parts=[
LoadCapabilityReturnPart(content={}, tool_call_id='load'),
ToolAvailabilityDeltaPart(tools_added=['guarded_tool']),
]
),
ModelResponse(parts=[boundary]),
]
result = await agent.run('go', message_history=history)
executed = [
part
for part in iter_message_parts(result.all_messages(), ModelRequest, ToolReturnPart)
if part.tool_name == 'guarded_tool'
]
assert executed == []
# The filter ran with the capability treated as active, and it removed the tool.
assert prepare_tools_calls[0] == snapshot(['guarded_tool'])
# And it ran off the retrospective supplement alone: the conservative window legitimately
# dropped the load, so the prospective set does not name the capability. That gap is the whole
# point — the filter has to follow whatever the gate authorizes from, not the narrower set.
assert loaded_ids_seen[0] == snapshot([])
assert [
str(part.content) for part in iter_message_parts(result.all_messages(), ModelRequest, RetryPromptPart)
] == snapshot(["Unknown tool name: 'guarded_tool'. Available tools: 'load_capability'"])
async def test_compaction_inside_serving_response_does_not_reset_tool_evidence():
"""Anthropic may compact mid-response; explicit provenance exercises the strict-before rule."""
toolset = FunctionToolset[Any]()
tests/test_toolsets.py
+128
View File
@@ -749,6 +749,134 @@ async def test_tool_manager_reuse_self():
assert tool_manager != updated_tool_manager
async def test_tool_manager_rebuilds_within_step_when_capability_availability_changes():
"""Tools are re-resolved mid-step when capability availability moved, without re-running the step transition.
Invariants an agent run can't reach today, so they're pinned directly rather than through
`Agent.run`. The retry carry-over sits behind the `run_step` short-circuit: a same-step rebuild
that dropped the accumulated counts would hand every tool a fresh budget mid-run, and one that
re-ran the carry-over would charge two units of that budget for a single failure. The per-step
accumulators have to survive it too, in both directions — `failed_tools` so the next step's
carry-over still charges the failure, `succeeded_tools` so it still clears the count of a tool
that recovered. And `AbstractToolset.for_run_step` is a step boundary with real lifecycle
effects (`DynamicToolset` exits and re-enters its inner toolset there, re-running its factory),
so a same-step rebuild has to reuse the already-transitioned toolset and only re-run
`get_tools`.
"""
factory_calls = 0
flaky_calls = 0
def toolset_func(ctx: RunContext[None]) -> AbstractToolset[None]:
nonlocal factory_calls
factory_calls += 1
toolset = FunctionToolset[None](max_retries=3)
@toolset.tool_plain
def failing_tool() -> int:
raise ModelRetry('This tool always fails')
@toolset.tool_plain
def flaky_tool() -> int:
nonlocal flaky_calls
flaky_calls += 1
if flaky_calls == 1:
raise ModelRetry('Not yet')
return flaky_calls
return toolset
secrets = Capability[None](id='secrets', description='Secret tools.', defer_loading=True)
step_1 = replace(build_run_context(None, run_step=1), capabilities={'secrets': secrets})
async with DynamicToolset(toolset_func) as dynamic_toolset:
tool_manager = await ToolManager[None](dynamic_toolset).for_run_step(step_1)
assert factory_calls == 1
for tool_name in ('failing_tool', 'flaky_tool'):
with pytest.raises(ToolRetryError):
await tool_manager.handle_call(ToolCallPart(tool_name=tool_name, args={}))
# A step advance charges both failures, so the same-step rebuild below has counts to carry.
tool_manager = await tool_manager.for_run_step(replace(step_1, run_step=2))
assert factory_calls == 2
assert tool_manager.ctx is not None
assert tool_manager.ctx.retries == snapshot({'failing_tool': 1, 'flaky_tool': 1})
with pytest.raises(ToolRetryError):
await tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={}))
assert await tool_manager.handle_call(ToolCallPart(tool_name='flaky_tool', args={})) == 2
assert tool_manager.failed_tools == {'failing_tool'}
assert tool_manager.succeeded_tools == {'flaky_tool'}
# Availability moves mid-step: same `run_step`, but the deferred capability is now loaded.
# (A fresh context carries no retries of its own, as the graph's does not.)
step_2_loaded = replace(step_1, run_step=2, loaded_capability_ids={'secrets'})
rebuilt = await tool_manager.for_run_step(step_2_loaded)
assert rebuilt is not tool_manager
# The step transition is not re-run, so the dynamic toolset's factory doesn't fire again.
assert factory_calls == 2
assert rebuilt.ctx is not None
# Retries carry through untouched: neither reset by the fresh context nor charged twice.
assert rebuilt.ctx.retries == snapshot({'failing_tool': 1, 'flaky_tool': 1})
# Both per-step accumulators survive, so the next step's carry-over sees this step's outcomes.
assert rebuilt.failed_tools == {'failing_tool'}
assert rebuilt.succeeded_tools == {'flaky_tool'}
# Availability unchanged: no rebuild, the same manager is returned.
assert await rebuilt.for_run_step(step_2_loaded) is rebuilt
# Advancing the step charges the failure exactly once and clears the recovered tool's count,
# and does re-run the transition.
advanced = await rebuilt.for_run_step(replace(step_2_loaded, run_step=3))
assert factory_calls == 3
assert advanced.ctx is not None
assert advanced.ctx.retries == snapshot({'failing_tool': 2})
assert advanced.failed_tools == set()
assert advanced.succeeded_tools == set()
async def test_tool_manager_availability_snapshot_is_independent_of_the_shared_set():
"""The recorded availability has to be a snapshot, and removal has to invalidate as well as addition.
The run's `loaded_capability_ids` is a single mutable set shared by reference with every
`RunContext` of the run: `_refresh_loaded_capability_ids` clears and updates it in place rather
than handing over a replacement, precisely so the copies stay in sync. A manager that recorded
the live set instead of a snapshot of it would compare it against itself and never invalidate.
Exercised here by mutating the very set the manager resolved against, which a test that passes
a fresh set per step does not do.
"""
toolset = FunctionToolset[None]()
@toolset.tool_plain
def plain_tool() -> int: # pragma: no cover
return 1
secrets = Capability[None](id='secrets', description='Secret tools.', defer_loading=True)
loaded_capability_ids: set[str] = set()
ctx = replace(
build_run_context(None, run_step=1),
capabilities={'secrets': secrets},
loaded_capability_ids=loaded_capability_ids,
)
assert ctx.loaded_capability_ids is loaded_capability_ids
tool_manager = await ToolManager[None](toolset).for_run_step(ctx)
assert tool_manager.resolved_capability_ids == snapshot(frozenset())
# Mutated in place, exactly as the run's own refresh does it.
loaded_capability_ids.add('secrets')
rebuilt = await tool_manager.for_run_step(ctx)
assert rebuilt is not tool_manager
assert rebuilt.resolved_capability_ids == snapshot(frozenset({'secrets'}))
# And back out: a processor can drop a load pair as easily as add one.
loaded_capability_ids.remove('secrets')
re_rebuilt = await rebuilt.for_run_step(ctx)
assert re_rebuilt is not rebuilt
assert re_rebuilt.resolved_capability_ids == snapshot(frozenset())
async def test_tool_manager_retry_logic():
"""Test the retry logic with failed_tools and for_run_step method."""