From 1ed9f78e82384bbb818b1ddf3d65df75dc7b036e Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 2 Oct 2026 17:33:15 -0600 Subject: [PATCH] Load a capability once when `load_capability` is called twice for it in one response (#9397) --- docs/capabilities/on-demand.md | 2 +- .../references/ON-DEMAND-CAPABILITIES.md | 1 + .../toolsets/_deferred_capability_loader.py | 33 ++++++++++- tests/test_capability_deferred_loading.py | 56 +++++++++++++++++++ 4 files changed, 89 insertions(+), 3 deletions(-) diff --git a/docs/capabilities/on-demand.md b/docs/capabilities/on-demand.md index 02ed59f829..986451428a 100644 --- a/docs/capabilities/on-demand.md +++ b/docs/capabilities/on-demand.md @@ -46,7 +46,7 @@ The model does not receive the refund instructions or the `refund_status` tool d 2. **Load.** Pydantic AI returns the capability's instructions — *"Always confirm the order ID before issuing a refund."* — as the tool result and exposes the `refund_status` definition on the next request. 3. **Request 2.** The model now sees those instructions in history and `refund_status` in its tool list. It calls `refund_status(order_id='ABC-123')` and answers the user from the result. -Already-loaded capabilities stay loaded for the rest of the run — the model never needs to re-open one. +Already-loaded capabilities stay loaded for the rest of the run — the model never needs to re-open one. A `load_capability` tool call for a capability that is already active, or one that an earlier call in the same response already loads, is answered with a retry prompt instead of a second copy of its instructions. Searching cannot reveal a capability-owned tool: it stays hidden until its capability loads. In runs that also have searchable deferred tools, the catalog explicitly steers the model to load the capability rather than search for its tools; in capability-only runs — where no search surface exists — the catalog omits any mention of searching. diff --git a/pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/references/ON-DEMAND-CAPABILITIES.md b/pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/references/ON-DEMAND-CAPABILITIES.md index 4b405bc755..c71104f555 100644 --- a/pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/references/ON-DEMAND-CAPABILITIES.md +++ b/pydantic_ai_slim/pydantic_ai/.agents/skills/building-pydantic-ai-agents/references/ON-DEMAND-CAPABILITIES.md @@ -120,6 +120,7 @@ When `load_capability` succeeds: - the call is typed as a capability-load message part - the return may include resolved capability instructions and owned toolset instructions - the capability id appears in `ctx.active_capability_ids` from the *next* step onwards, not within the step that loaded it — both sets are derived from message history before each model request +- only the first `load_capability` call for an id in a response loads it; later calls for that id in the same response get a retry prompt, the same way a load of an already-active capability does - tools owned by the loaded capability become visible, and callable, on later steps - `load_capability` remains visible so the tool set stays stable diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/_deferred_capability_loader.py b/pydantic_ai_slim/pydantic_ai/toolsets/_deferred_capability_loader.py index 5c978c6c32..fc023e3216 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/_deferred_capability_loader.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/_deferred_capability_loader.py @@ -5,11 +5,11 @@ from typing import Any from pydantic import TypeAdapter -from pydantic_ai._deferred_capabilities import LoadCapabilityArgs, LoadCapabilityReturn +from pydantic_ai._deferred_capabilities import LoadCapabilityArgs, LoadCapabilityCallPart, LoadCapabilityReturn from pydantic_ai._instructions import resolve_sourced_instructions from pydantic_ai._run_context import AgentDepsT, RunContext from pydantic_ai.exceptions import ModelRetry, UserError -from pydantic_ai.messages import InstructionPart, ToolReturn +from pydantic_ai.messages import InstructionPart, ModelResponse, ToolReturn from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets._capability_owned import CapabilityOwnedToolset from pydantic_ai.toolsets._instruction_collection import collect_toolset_instructions @@ -25,6 +25,10 @@ LOAD_CAPABILITY_ALREADY_ACTIVE_MESSAGE_TEMPLATE = ( 'Capability {capability_id!r} is already active. ' 'Use its existing instructions and any tools it provides; do not call `load_capability` for it again.' ) +LOAD_CAPABILITY_DUPLICATE_CALL_MESSAGE_TEMPLATE = ( + 'Capability {capability_id!r} is already being loaded by an earlier `load_capability` call in this response; ' + 'do not call `load_capability` for it again.' +) _load_capability_args_ta = TypeAdapter(LoadCapabilityArgs) _LOAD_CAPABILITY_SCHEMA = _load_capability_args_ta.json_schema() @@ -78,6 +82,8 @@ class DeferredCapabilityLoaderToolset(WrapperToolset[AgentDepsT]): raise ModelRetry(f'No capability found with id {capability_id!r}.') if capability_id in ctx.active_capability_ids: raise ModelRetry(LOAD_CAPABILITY_ALREADY_ACTIVE_MESSAGE_TEMPLATE.format(capability_id=capability_id)) + if _is_duplicate_load_in_response(ctx, capability_id): + raise ModelRetry(LOAD_CAPABILITY_DUPLICATE_CALL_MESSAGE_TEMPLATE.format(capability_id=capability_id)) # Sourced through `_collect_instructions` rather than `get_instructions` so a loaded # capability's parts carry the same `capability:` keys they would have had if the @@ -112,3 +118,26 @@ class DeferredCapabilityLoaderToolset(WrapperToolset[AgentDepsT]): for ts in owned: parts.extend(await collect_toolset_instructions(ts.wrapped, ctx)) return parts + + +def _is_duplicate_load_in_response(ctx: RunContext[Any], capability_id: str) -> bool: + """Whether an earlier call in the response being executed already loads `capability_id`. + + A load only counts as loaded once its return reaches history at the end of the step, so sibling + calls for the same id in one response can't see each other through `active_capability_ids`. + Ownership is instead derived from the response itself — the last one in history, whose calls are + the ones executing — in model call order, like `_prune_duplicate_tool_reveals` gives the first + call to name a tool its reveal, so which call delivers the instructions doesn't depend on task + scheduling. + """ + response = next((message for message in reversed(ctx.messages) if isinstance(message, ModelResponse)), None) + parts = response.parts if response is not None else () + owner_tool_call_id = next( + ( + part.tool_call_id + for part in parts + if isinstance(part, LoadCapabilityCallPart) and part.capability_id == capability_id + ), + ctx.tool_call_id, + ) + return owner_tool_call_id != ctx.tool_call_id diff --git a/tests/test_capability_deferred_loading.py b/tests/test_capability_deferred_loading.py index a411007cc7..5d6854a9d7 100644 --- a/tests/test_capability_deferred_loading.py +++ b/tests/test_capability_deferred_loading.py @@ -12,6 +12,7 @@ from typing import Any, cast import pytest +from pydantic_ai._deferred_capabilities import parse_loaded_capabilities from pydantic_ai._run_context import RunContext from pydantic_ai.agent import Agent from pydantic_ai.capabilities import ( @@ -57,10 +58,12 @@ from pydantic_ai.native_tools import ( ) from pydantic_ai.native_tools._tool_search import ToolSearchTool from pydantic_ai.settings import ModelSettings as _ModelSettings +from pydantic_ai.tool_manager import ParallelExecutionMode from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets import FunctionToolset from pydantic_ai.toolsets._deferred_capability_loader import ( LOAD_CAPABILITY_ALREADY_ACTIVE_MESSAGE_TEMPLATE, + LOAD_CAPABILITY_DUPLICATE_CALL_MESSAGE_TEMPLATE, LOAD_CAPABILITY_TOOL_NAME, ) from pydantic_ai.usage import RequestUsage, RunUsage @@ -2519,3 +2522,56 @@ async def test_load_capability_retries_when_capability_is_already_loaded() -> No ] assert len(load_returns) == 1 assert load_returns[0].instructions == 'Deferred instructions.' + + +@pytest.mark.parametrize('mode', ['parallel', 'sequential']) +async def test_load_capability_called_twice_in_one_response_loads_once(mode: ParallelExecutionMode) -> None: + """Only the first `load_capability` call for an id in a response delivers its instructions (#7298). + + Sibling calls can't see each other's load through `active_capability_ids` (a load only counts once + its return reaches history), so the duplicate is refused by response position instead, and a load + of a different capability in the same response is unaffected. + """ + + def dyn_tool() -> str: + return 'ok' # pragma: no cover + + dyn = Capability[object]( + id='dyn', description='Dynamic.', instructions='Dyn runbook.', tools=[dyn_tool], defer_loading=True + ) + other = Capability[object](id='other', description='Other.', instructions='Other runbook.', defer_loading=True) + + def model_fn(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse: + if len(messages) == 1: + return ModelResponse( + parts=[ + ToolCallPart(LOAD_CAPABILITY_TOOL_NAME, {'id': 'dyn'}, tool_call_id='first'), + ToolCallPart(LOAD_CAPABILITY_TOOL_NAME, {'id': 'other'}, tool_call_id='other'), + ToolCallPart(LOAD_CAPABILITY_TOOL_NAME, {'id': 'dyn'}, tool_call_id='duplicate'), + ] + ) + return make_text_response('done') + + agent = Agent(FunctionModel(model_fn), capabilities=[dyn, other]) + with Agent.parallel_tool_call_execution_mode(mode): + result = await agent.run('load dyn twice') + + assert result.output == 'done' + messages = result.all_messages() + tool_results = [ + part + for part in messages[2].parts + if isinstance(part, LoadCapabilityReturnPart | RetryPromptPart | ToolAvailabilityDeltaPart) + ] + assert [(type(part).__name__, part.tool_call_id) for part in tool_results] == [ + ('LoadCapabilityReturnPart', 'first'), + ('ToolAvailabilityDeltaPart', 'first'), + ('LoadCapabilityReturnPart', 'other'), + ('RetryPromptPart', 'duplicate'), + ] + first, _, other_return, duplicate = tool_results + assert isinstance(first, LoadCapabilityReturnPart) and first.instructions == 'Dyn runbook.' + assert isinstance(other_return, LoadCapabilityReturnPart) and other_return.instructions == 'Other runbook.' + assert isinstance(duplicate, RetryPromptPart) + assert duplicate.content == LOAD_CAPABILITY_DUPLICATE_CALL_MESSAGE_TEMPLATE.format(capability_id='dyn') + assert parse_loaded_capabilities(messages) == {'dyn', 'other'}