furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 10:55:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Keep framework tools native in CodeMode (#296)

Browse Source 
* Keep framework tools native in CodeMode

* test(code_mode): cover load_capability bootstrap through Agent

The existing unit test exercises the native/sandbox split via a static
toolset. Add an Agent-level test for the actual issue #276 scenario: a
deferred capability plus CodeMode(tools='all'), asserting load_capability
stays a native call and the deferred member tool keeps defer_loading=True
instead of being folded into run_code.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor(code_mode): route framework tools native off tool_kind alone

The name fallback (`name == _SEARCH_TOOLS_NAME`) was redundant: pydantic-ai
has set `tool_kind` on the `search_tools` ToolDefinition since 1.95.0, and the
dependency floor is 1.105.0, so `tool_kind is not None` already catches it.
Routing keys purely off the framework discriminator now.

Also switch the #276 bootstrap test from FunctionModel to TestModel (per the
test guideline) and rescope its assertions to provider-agnostic facts -- the
old assertion only held because FunctionModel's profile happens to support
native tool search.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(code_mode): mark unreached deferred-tool body no cover

The #276 bootstrap test keeps `demo_tool` hidden, so its body never runs.
Mark it `# pragma: no cover` to satisfy the 100% coverage gate, matching the
sibling deferred-tool tests.

---------

Co-authored-by: Aditya Vardhan <adtyavrdhn@gmail.com>
This commit is contained in:
King Star
2026-06-23 18:14:17 +05:30
committed by GitHub
co-authored by Aditya Vardhan
parent 6232e0762b
commit bb6da9948d
2 changed files with 100 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pydantic_ai_harness/code_mode/_toolset.py
+16 -14
View File
@@ -261,11 +261,13 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
callable from the sandbox at runtime. Non-selected tools remain visible
to the model as normal tool calls.
Tools with `defer_loading=True` (Tool Search) are never sandboxed: they stay
native pass-through so the deferred-loading contract is honored, and only get
folded into `run_code` once they've been discovered (`defer_loading=False`).
Tools annotated with `unless_native` likewise stay native so
`Model.prepare_request` can drop them when the provider supports the native tool.
Some tools always stay native rather than being sandboxed:
- Framework control tools (`tool_kind` set: tool search, capability loading).
- `defer_loading=True` tools, until discovery flips them to `defer_loading=False`.
- `unless_native` tools, so `Model.prepare_request` can drop them when the
provider supports the native tool.
To keep a Tool Search corpus native even after discovery (e.g. for prompt-cache
stability), pass a `tool_selector` that excludes tools with `with_native` set.
"""
@@ -357,22 +359,22 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
wrapped_tools = await self.wrapped.get_tools(ctx)
# Split tools into sandboxed vs native based on the selector.
# The search_tools tool (from ToolSearchToolset) is always kept native
# so the model can discover deferred tools alongside run_code.
sandboxed_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
native_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
for name, tool in wrapped_tools.items():
if name == _SEARCH_TOOLS_NAME:
# Framework control tools (tool search, capability loading) stay native to
# drive protocol-level flows. `tool_kind` is the framework's discriminator
# for them; pydantic-ai has set it on `search_tools` since 1.95.0.
if tool.tool_def.tool_kind is not None:
native_tools[name] = tool
elif tool.tool_def.defer_loading:
# Tool Search keeps these out of the model's initial context until discovered.
# Stay native pass-through so `ToolSearchToolset`'s `defer_loading` /
# `with_native` flags reach `Model.prepare_request` unaltered; once a tool is
# discovered it comes back with `defer_loading=False` and is sandboxed from then on.
# Stay native so Tool Search's `defer_loading`/`with_native` flags reach
# `Model.prepare_request` unaltered. Discovery flips `defer_loading` to
# False, and the tool is sandboxed from then on.
native_tools[name] = tool
elif tool.tool_def.unless_native:
# Defer to `Model.prepare_request`'s `unless_native` filtering: keep the local
# fallback native so it can be dropped when the provider supports the native tool.
# Keep the local fallback native so `Model.prepare_request` can drop it
# when the provider supports the native tool.
native_tools[name] = tool
elif await matches_tool_selector(self.tool_selector, ctx, tool.tool_def):
sandboxed_tools[name] = tool
tests/code_mode/test_code_mode.py
+84 -10
View File
@@ -760,6 +760,32 @@ class TestCodeMode:
assert 'async def later' in description
assert 'later' not in tools
async def test_framework_tool_kind_tool_not_sandboxed(self) -> None:
"""Framework control tools with `tool_kind` stay native even when CodeMode wraps all user tools."""
td_loader = ToolDefinition(
name='load_capability',
description='Load a deferred capability.',
parameters_json_schema={
'type': 'object',
'properties': {'capability_id': {'type': 'string'}},
'required': ['capability_id'],
},
return_schema={'type': 'string'},
tool_kind='capability-load',
)
static = _StaticToolset([_make_address_tool_def('get_user', 'Get a user.', 'street'), td_loader])
wrapper = CodeMode[None]().get_wrapper_toolset(static)
assert isinstance(wrapper, CodeModeToolset)
tools = await wrapper.get_tools(build_run_context(None))
description = tools['run_code'].tool_def.description
assert description is not None
assert 'async def get_user' in description
assert 'load_capability' not in description
assert 'load_capability' in tools
assert tools['load_capability'].tool_def.tool_kind == 'capability-load'
async def test_unless_native_tool_not_sandboxed(self) -> None:
"""Tools annotated with `unless_native` stay native so `Model.prepare_request` can filter them."""
td_fallback = ToolDefinition(
@@ -961,6 +987,54 @@ class TestCodeMode:
# The agent's final output reflects the value flowing through the sandbox.
assert result.output == 'sum is 10'
async def test_deferred_capability_loader_stays_native_with_tools_all(self) -> None:
"""Regression for the deferred-capability bootstrap (issue #276).
With `CodeMode(tools='all')` and a deferred capability configured, the
framework-managed `load_capability` tool must reach the model as a native call
(alongside `run_code`) so the model can reveal the capability. The deferred
member tool stays hidden -- it is neither folded into `run_code` nor surfaced as
a plain tool until loaded.
(The native-vs-sandbox split per tool kind is covered directly at the toolset
level by `test_framework_tool_kind_tool_not_sandboxed` and
`test_tool_search_toolset_deferred_tool_not_in_run_code`; this exercises the
end-to-end path through `Agent`.)
"""
from pydantic_ai.capabilities import Capability
capability = Capability[None](
id='demo',
description='Demo deferred capability.',
instructions='Use demo_tool.',
defer_loading=True,
)
@capability.tool_plain
def demo_tool() -> str: # pyright: ignore[reportUnusedFunction]
return 'ok' # pragma: no cover - deferred tool stays hidden, body is not invoked
model = TestModel(call_tools=[])
agent: Agent[None, str] = Agent(
model,
capabilities=[capability, CodeMode[None](tools='all')],
)
await agent.run('inspect tools')
assert model.last_model_request_parameters is not None
by_name = {td.name: td for td in model.last_model_request_parameters.function_tools}
# The bootstrap tool is a native call alongside `run_code`, not buried in the sandbox.
assert 'load_capability' in by_name
assert 'run_code' in by_name
# The deferred member tool stays hidden until loaded: not folded into `run_code`
# and not surfaced as a plain native tool.
assert 'demo_tool' not in by_name
run_code_desc = by_name['run_code'].description or ''
assert 'demo_tool' not in run_code_desc
assert 'load_capability' not in run_code_desc
# ---------------------------------------------------------------------------
# Capability registration
# ---------------------------------------------------------------------------
@@ -2064,7 +2138,7 @@ class TestDynamicCatalog:
await cap.after_tool_execute(
ctx,
call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
tool_def=_search_tool_def_with_kind(),
tool_def=_search_tool_def(),
args={},
result={'discovered_tools': [{'name': 'weather', 'description': '...'}]},
)
@@ -2086,7 +2160,7 @@ class TestDynamicCatalog:
await cap.after_tool_execute(
ctx,
call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
tool_def=_search_tool_def_with_kind(),
tool_def=_search_tool_def(),
args={},
result={'discovered_tools': [{'name': 'weather'}]},
)
@@ -2100,7 +2174,7 @@ class TestDynamicCatalog:
await cap.after_tool_execute(
ctx,
call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
tool_def=_search_tool_def_with_kind(),
tool_def=_search_tool_def(),
args={},
result={'discovered_tools': []},
)
@@ -2133,7 +2207,7 @@ class TestDynamicCatalog:
await cap.after_tool_execute(
ctx,
call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id=cid),
tool_def=_search_tool_def_with_kind(),
tool_def=_search_tool_def(),
args={},
result=result,
)
@@ -2332,11 +2406,6 @@ class TestDynamicCatalog:
assert result.output == 'got 7'
def _search_tool_def_with_kind(name: str = 'search_tools') -> ToolDefinition:
"""A `search_tools` ToolDefinition carrying `tool_kind='tool-search'` (the announcement trigger)."""
return ToolDefinition(name=name, description='', parameters_json_schema={}, tool_kind='tool-search')
def _unused_os_callback(fn: OsFunction, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
"""An `os` callback for tests that only assert description/forwarding, never run code."""
return NOT_HANDLED # pragma: no cover - never invoked by these tests
@@ -2530,13 +2599,18 @@ class TestCodeModeOSAccess:
def _search_tool_def(description: str = 'Search for tools.') -> ToolDefinition:
"""Create a ToolDefinition mimicking the search_tools tool from ToolSearchToolset."""
"""Create a ToolDefinition mimicking the search_tools tool from ToolSearchToolset.
Carries `tool_kind='tool-search'`, matching what pydantic-ai emits (since 1.95.0);
CodeMode routes it native off `tool_kind`, not its name.
"""
from pydantic_ai.toolsets._tool_search import _SEARCH_TOOLS_NAME
return ToolDefinition(
name=_SEARCH_TOOLS_NAME,
description=description,
parameters_json_schema={'type': 'object', 'properties': {'keywords': {'type': 'string'}}},
tool_kind='tool-search',
)