diff --git a/python/packages/core/agent_framework/_agent_hooks.py b/python/packages/core/agent_framework/_agent_hooks.py index ea5ff016e6..9d18b2a883 100644 --- a/python/packages/core/agent_framework/_agent_hooks.py +++ b/python/packages/core/agent_framework/_agent_hooks.py @@ -862,23 +862,99 @@ def _agent_updates_from_response(response: AgentResponse[Any]) -> list[AgentResp return updates -def _tool_names(context: AgentContext) -> list[str]: - """Project the registered tool names for ``agent_startup`` (spec ``tools_registered``).""" - from ._tools import _get_tool_name, normalize_tools # type: ignore[reportPrivateUsage] +def _normalized_tools(tools: Any, *, point: str) -> list[Any]: + """Normalize a tools value for projection; empty (with a warning) when it cannot be. + + Everything — including the emptiness check inside ``normalize_tools`` — runs under + the guard, so a tools container whose ``__bool__``/``__len__``/``__iter__`` raises + degrades to omitting the projection instead of aborting the emission mid-run. + """ + from ._tools import normalize_tools - tools: Any = context.tools if context.tools is not None else getattr(context.agent, "tools", None) - if tools is None: + try: + return list(normalize_tools(tools)) + except Exception: + logger.warning("agent-hooks could not normalize the tools for the %s projection.", point) return [] + + +def _projected_tool_name(item: Any) -> str: + """Project a tool's display name; hosted-tool mappings fall back to their ``type``.""" + from ._tools import _get_tool_name # type: ignore[reportPrivateUsage] + + fallback = type(item).__name__ + name = _get_tool_name(item) + if name: + return name + if isinstance(item, Mapping): + # Hosted-tool mappings carry top-level fields (no nested "function"), e.g. + # {"type": "web_search"} or {"type": "web_search_20250305", "name": "web_search"}. + mapping = cast("Mapping[str, Any]", item) + for key in ("name", "type"): + value = mapping.get(key) + if isinstance(value, str) and value: + return value + return fallback + + +def _tool_description(item: Any) -> str | None: + """Extract a tool description from a tool object or dict tool definition.""" + if isinstance(item, Mapping): + mapping = cast("Mapping[str, Any]", item) + function = mapping.get("function") + # Function-tool dicts nest the description; hosted-tool mappings keep it top-level. + description = ( + cast("Mapping[str, Any]", function).get("description") + if isinstance(function, Mapping) + else mapping.get("description") + ) + return description if isinstance(description, str) else None + description = getattr(item, "description", None) + return description if isinstance(description, str) else None + + +def _tool_names(context: AgentContext) -> list[str]: + """Project the run-start tool names for ``agent_startup`` (spec ``tools_registered``). + + The framework owns the run-start resolution — ``AgentContext._resolve_run_start_tools`` + returns the agent-declared tools plus this invocation's run-level tools (the named + ``tools`` parameter takes precedence over a ``tools`` entry in the options mapping), + already normalized — so this helper is projection only. This is deliberately the + run-start snapshot: tools registered dynamically during the run (for example by + context providers during run preparation, or by MCP servers whose functions expand + at connect time) cannot be known at ``agent_startup`` time; they surface in each + ``pre_model_call`` emission's ``tools`` projection (the completed per-call set) and + are bracketed by ``pre_tool_call``/``post_tool_call`` like any other tool when + invoked. An unresolvable set degrades to a warning and an empty snapshot rather + than aborting the emission. + """ try: - normalized = normalize_tools(tools) + resolved = context._resolve_run_start_tools() # pyright: ignore[reportPrivateUsage] except Exception: - logger.warning("agent-hooks could not normalize the run's tools for the agent_startup projection.") + logger.warning("agent-hooks could not resolve the run-start tools for the agent_startup projection.") return [] - names: list[str] = [] - for item in normalized: - name = _get_tool_name(item) - names.append(name if name else type(item).__name__) - return names + return [_projected_tool_name(item) for item in resolved] + + +def _pre_model_call_tools(options: Mapping[str, Any]) -> list[dict[str, Any]] | None: + """Project the per-call effective tool set for ``pre_model_call`` (spec ``tools``). + + This is the completed set for this model call — the run-start tools plus anything + registered during the run (context-provider tools, connected MCP-server functions, + progressive tool exposure) — whereas ``agent_startup``'s ``tools_registered`` is the + run-start snapshot. Entries carry ``{"name", "description"?}`` (description omitted + when the tool has none). Returns ``None`` — the spec's optional field is omitted — + when the call offers no tools or the set cannot be projected, so an unprojectable + set is never misreported as "no tools". + """ + projected: list[dict[str, Any]] = [] + for item in _normalized_tools(options.get("tools"), point="pre_model_call"): + entry: dict[str, Any] = {"name": _projected_tool_name(item)} + description = _tool_description(item) + if description: + entry["description"] = description + projected.append(entry) + return projected or None def _is_host_error(record: InterceptionRecord) -> bool: @@ -1269,7 +1345,7 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ model_id = str(options.get("model") or type(context.client).__name__) before = _ModelRequestCodec.to_wire(context.messages) outcome: EmitOutcome = await state.emitter.emit( - state.builder.pre_model_call(model_id=model_id, messages=before) + state.builder.pre_model_call(model_id=model_id, messages=before, tools=_pre_model_call_tools(options)) ) transformed_messages = _ModelRequestCodec.write_back(context.messages, before, outcome.target) if transformed_messages is not None: diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index d22f6e14e1..42ca1ba03a 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -31,6 +31,7 @@ FunctionInvocationContext, MiddlewareTypes, _as_middleware_list, # pyright: ignore[reportPrivateUsage] + _select_run_level_tools, # pyright: ignore[reportPrivateUsage] categorize_middleware, ) from ._serialization import SerializationMixin @@ -1342,8 +1343,13 @@ async def _prepare_run_context( opts = dict(options) if options else {} existing_additional_args: dict[str, Any] = opts.pop("additional_function_arguments", None) or {} - # Get tools from options or named parameter (named param takes precedence) - tools_ = tools if tools is not None else opts.pop("tools", None) + # Run-level tools: the named parameter takes precedence over an options entry + # (_select_run_level_tools is the framework's single statement of that rule, + # shared with the middleware layer's run-start resolution). The options entry + # is consumed either way, so a losing options["tools"] can never ride the + # remaining options into the request and silently override the resolved list. + tools_ = _select_run_level_tools(tools, opts) + opts.pop("tools", None) input_messages = normalize_messages(messages) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 45a3dcbf82..2e682d84ee 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -7,7 +7,7 @@ import logging import sys from abc import ABC, abstractmethod -from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Iterable, Mapping, Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast, overload @@ -151,6 +151,60 @@ class MiddlewareType(str, Enum): CHAT = "chat" +def _select_run_level_tools(tools: Any, options: Mapping[str, Any] | None) -> Any: + """Select this invocation's run-level tool source. + + The named ``tools`` parameter takes precedence over a ``tools`` entry in the + options mapping. This is the framework's single statement of that rule — the run-start + resolution (:meth:`AgentContext._resolve_run_start_tools`, projected into + ``agent_startup`` by observability middleware) and the run's own setup + (``Agent._prepare_run_context``) both consume it, so the run-start view can never + disagree with what the run executes. + """ + if tools is not None: + return tools + if options is not None: + return options.get("tools") + return None + + +def _materialize_tool_container(tools: Any) -> Any: + """Materialize one-shot iterable tool containers into lists, tools untouched. + + ``normalize_tools`` recursively flattens any iterable tool collection, so + generators and other single-pass iterables are supported containers at every + nesting level — but each observer that iterates one consumes it for everyone + after it. This walks exactly the container shapes that flattening walks and lists + them once, without converting any tool: middleware must keep seeing the caller's + original tool objects, so identity-based policy checks still fire. Leaf shapes — + tools, dict specs, mapping-like collections (flattened via their re-iterable + ``.tools`` attribute), pydantic models, strings, callables — pass through + untouched, and an already re-iterable container whose elements needed no + materialization keeps its identity. + """ + from pydantic import BaseModel + + from ._mcp import MCPTool + from ._tools import FunctionTool + + def is_leaf(value: Any) -> bool: + # Mirror the shapes normalize_tools treats as non-container leaves (or + # flattens without directly iterating the value itself). + if value is None or isinstance(value, (FunctionTool, MCPTool, Mapping, BaseModel, str, bytes, bytearray)): + return True + return callable(value) or not isinstance(value, Iterable) + + def materialize(value: Any) -> Any: + if is_leaf(value): + return value + items = [materialize(item) for item in cast("Iterable[Any]", value)] + if isinstance(value, Sequence) and all(new is old for new, old in zip(items, cast("Sequence[Any]", value))): + return cast("Any", value) + return items + + return materialize(tools) + + class AgentContext: """Context object for agent middleware invocations. @@ -266,6 +320,34 @@ def __init__( # gate binds to that run's identity and never to middleware-initiated runs. self._run_persistence_gate: _RunPersistenceGate | None = None + def _resolve_run_start_tools(self) -> list[ToolTypes]: + """Resolve the run-start tool list for this invocation, normalized. + + This is the framework's one statement of the run-start tool policy, kept next + to the run-option rules it mirrors so they evolve together (middleware such as + agent-hooks reads it instead of re-deriving the precedence): the agent's + declared tools (:class:`~agent_framework.Agent` keeps them in + ``default_options["tools"]``; other agent implementations may expose a + ``tools`` attribute) followed by this invocation's run-level tools, where the + named ``tools`` parameter takes precedence over a ``tools`` entry in the + options mapping — matching the run's own resolution in + ``Agent._prepare_run_context``. Tools registered later in the run (context + providers during run preparation, MCP servers expanding at connect time, + progressive tool exposure) are deliberately not part of the run-start view. + + Normalization errors propagate; callers that must not fail should guard. + """ + from ._tools import normalize_tools + + declared_options = getattr(self.agent, "default_options", None) + declared: Any = ( + cast("Mapping[str, Any]", declared_options).get("tools") if isinstance(declared_options, Mapping) else None + ) + if declared is None: + declared = getattr(self.agent, "tools", None) + run_level = _select_run_level_tools(self.tools, self.options) + return [*normalize_tools(declared), *normalize_tools(run_level)] + class FunctionInvocationContext: """Context object for function middleware invocations. @@ -1535,6 +1617,32 @@ def run( effective_function_invocation_kwargs = ( dict(function_invocation_kwargs) if function_invocation_kwargs is not None else {} ) + # Select the winning run-level tool route first, then materialize only that + # source: the losing route is never iterated — a losing one-shot options + # entry stays untouched for its owner and cannot raise or trigger side + # effects — and it is dropped from the forwarded options (on a copy), so no + # layer below (telemetry serialization, run setup) ever consumes or records a + # source the run will not use. Materialization covers the nested collection + # forms normalize_tools recursively flattens, so every observer of the run — + # middleware pipeline, telemetry, run setup — shares one re-iterable + # structure of the caller's original tool objects: observation never consumes + # the run's tool source, and identity-based policy checks (for example + # rejecting one specific privileged callable) keep seeing exactly what the + # caller supplied. + selected_tools = _select_run_level_tools(tools, options) + materialized_tools = _materialize_tool_container(selected_tools) + if tools is not None: + tools = materialized_tools + if options is not None and "tools" in options: + options = cast( + "ChatOptions[Any]", + {key: item for key, item in options.items() if key != "tools"}, + ) + elif materialized_tools is not selected_tools: + # The options mapping supplied the winner; swap the materialized value in + # on a copy (the caller's mapping is never mutated). + options = cast("ChatOptions[Any]", {**cast("Mapping[str, Any]", options), "tools": materialized_tools}) + # Execute with middleware if available if not pipeline.has_middlewares: return super().run( # type: ignore[misc, no-any-return] diff --git a/python/packages/core/tests/core/test_agent_hooks.py b/python/packages/core/tests/core/test_agent_hooks.py index 96c2ee3883..fb34dd8d5f 100644 --- a/python/packages/core/tests/core/test_agent_hooks.py +++ b/python/packages/core/tests/core/test_agent_hooks.py @@ -371,6 +371,310 @@ def structured_tool(order_id: str) -> dict[str, Any]: assert post_tool["target"] == value +@requires_sdk +async def test_pre_model_call_projects_per_call_effective_tools(chat_client_base: MockBaseChatClient) -> None: + """Every ``pre_model_call`` carries the effective tool set for that call (spec ``tools``). + + Constructor-registered and run-level tools are both part of the effective set the + model is offered, so both must appear — a registration-time-only projection would + hide the run-level tools from auditors (the Python half of #7560; parity with the + .NET per-call ``ChatOptions.Tools`` projection). + """ + + @tool(approval_mode="never_require") + def run_only_tool(city: str) -> str: + """Look up a city.""" + return city + + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=[create_agent_hooks_middleware([guard])], + ) + + await agent.run("Get weather for Seattle", tools=[run_only_tool]) + + pre_models = guard.contexts_for("pre_model_call") + assert len(pre_models) == 2 + for pre_model in pre_models: + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool", "run_only_tool"] + # Descriptions ride along ({name, description?}), so auditors see what the model saw. + assert pre_models[0]["tools"][0]["description"] == "Get the weather for a location." + # agent_startup's tools_registered is the run-start snapshot: both are known at run + # start here, so both appear (constructor tools were previously dropped entirely). + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool", "run_only_tool"] + + +@requires_sdk +async def test_agent_startup_projects_constructor_registered_tools(chat_client_base: MockBaseChatClient) -> None: + """Constructor-registered tools appear in ``tools_registered`` (#7560).""" + guard = AllowGuard() + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[create_agent_hooks_middleware([guard])]) + + await agent.run("hello") + + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool"] + + +@requires_sdk +async def test_pre_model_call_tools_include_provider_contributed_tools(chat_client_base: MockBaseChatClient) -> None: + """Tools registered during run preparation surface in the per-call ``tools`` projection. + + Context providers contribute tools after ``agent_startup`` has been emitted, so the + run-start snapshot cannot know them — the per-call projection is where they become + visible to auditors (same contract as the .NET fix on #7564). + """ + from agent_framework import ContextProvider + + @tool(approval_mode="never_require") + def provider_tool(query: str) -> str: + """A tool contributed by a context provider.""" + return query + + class ToolContextProvider(ContextProvider): + def __init__(self) -> None: + super().__init__(source_id="tool-context") + + async def before_run(self, *, agent: Any, session: Any, context: Any, state: Any) -> None: + context.extend_tools("tool-context", [provider_tool]) + + guard = AllowGuard() + agent = Agent( + client=chat_client_base, + context_providers=[ToolContextProvider()], + middleware=[create_agent_hooks_middleware([guard])], + ) + + await agent.run("hello") + + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == [] + pre_model = guard.contexts_for("pre_model_call")[0] + assert [entry["name"] for entry in pre_model["tools"]] == ["provider_tool"] + + +@requires_sdk +@pytest.mark.parametrize("max_iterations", [1], indirect=True) +async def test_tools_disabled_final_call_still_projects_effective_tools( + chat_client_base: MockBaseChatClient, +) -> None: + """The loop's ``tool_choice="none"`` final call projects its effective options' tools. + + When the iteration budget is exhausted the function-invocation loop requests one + final response with ``tool_choice="none"`` but the tools still in the options — + the projection reflects exactly those effective options (parity with the .NET + per-call ``ChatOptions.Tools`` projection), not a guess about tool availability. + """ + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), tool_call_response()] + agent = Agent( + client=chat_client_base, + tools=[weather_tool], + middleware=[create_agent_hooks_middleware([guard])], + ) + + response = await agent.run("Get weather for Seattle") + + assert "broke out" in response.text + pre_models = guard.contexts_for("pre_model_call") + assert len(pre_models) == 2 + for pre_model in pre_models: + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool"] + + +@requires_sdk +async def test_pre_model_call_omits_tools_when_call_has_none(chat_client_base: MockBaseChatClient) -> None: + """A call with no tools omits the optional ``tools`` field instead of claiming an empty set.""" + guard = AllowGuard() + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) + + await agent.run("hello") + + pre_model = guard.contexts_for("pre_model_call")[0] + assert "tools" not in pre_model + + +@requires_sdk +async def test_hosted_tool_mappings_project_top_level_name_and_description( + chat_client_base: MockBaseChatClient, +) -> None: + """Hosted-tool mappings project their top-level fields, with ``type`` naming unnamed tools. + + The provider factories return plain mappings without a nested ``function`` object + (e.g. OpenAI's ``{"type": "web_search"}``, Anthropic's ``{"type": "web_search_20250305", + "name": "web_search"}``); those must not be projected as ``{"name": "dict"}``. + """ + guard = AllowGuard() + agent = Agent( + client=chat_client_base, + tools=[ + {"type": "web_search"}, + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "custom_hosted", "name": "lookup", "description": "Look things up."}, + ], + middleware=[create_agent_hooks_middleware([guard])], + ) + + await agent.run("hello") + + pre_model = guard.contexts_for("pre_model_call")[0] + assert pre_model["tools"] == [ + {"name": "web_search"}, + {"name": "web_search"}, + {"name": "lookup", "description": "Look things up."}, + ] + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["web_search", "web_search", "lookup"] + + +@requires_sdk +def test_tools_projection_survives_hostile_tools_container(caplog: pytest.LogCaptureFixture) -> None: + """A tools container whose ``__bool__`` raises degrades to omission, never an emission abort.""" + from agent_framework._agent_hooks import _pre_model_call_tools + + class HostileTools: + def __bool__(self) -> bool: + raise RuntimeError("hostile __bool__") + + def __iter__(self) -> Any: + return iter([]) + + with caplog.at_level("WARNING", logger="agent_framework._agent_hooks"): + assert _pre_model_call_tools({"tools": HostileTools()}) is None + assert "could not normalize the tools" in caplog.text + + # The startup snapshot degrades the same way when the run-start resolution raises. + from agent_framework._agent_hooks import _tool_names + + class HostileAgent: + tools = HostileTools() + + context = AgentContext(agent=cast("Any", HostileAgent()), messages=[]) + with caplog.at_level("WARNING", logger="agent_framework._agent_hooks"): + assert _tool_names(context) == [] + assert "could not resolve the run-start tools" in caplog.text + + +@requires_sdk +def test_startup_snapshot_falls_back_to_legacy_tools_attribute() -> None: + """A custom agent whose ``default_options`` mapping has no tools entry keeps its + ``tools``-attribute projection (the pre-existing fallback for non-``Agent`` hosts).""" + from agent_framework._agent_hooks import _tool_names + + class LegacyAgent: + default_options = {"temperature": 0.2} # mapping-valued, but no "tools" key + tools = [weather_tool] + + context = AgentContext(agent=cast("Any", LegacyAgent()), messages=[]) + assert _tool_names(context) == ["weather_tool"] + + +@requires_sdk +async def test_options_route_run_tools_appear_on_both_projections(chat_client_base: MockBaseChatClient) -> None: + """Run-level tools passed as ``options={"tools": ...}`` (instead of the ``tools=`` + keyword) appear in the run-start snapshot and the per-call projection alike.""" + + @tool(approval_mode="never_require") + def options_route_tool(city: str) -> str: + """Arrives through the options dict.""" + return city + + guard = AllowGuard() + agent = Agent(client=chat_client_base, tools=[weather_tool], middleware=[create_agent_hooks_middleware([guard])]) + + await agent.run("hello", options={"tools": [options_route_tool]}) + + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool", "options_route_tool"] + pre_model = guard.contexts_for("pre_model_call")[0] + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool", "options_route_tool"] + + +@requires_sdk +@pytest.mark.parametrize("route", ["tools_kwarg", "options_dict"]) +async def test_one_shot_iterable_run_tools_survive_the_snapshot( + chat_client_base: MockBaseChatClient, route: str +) -> None: + """Observing the tools never consumes them: one-shot iterables stay usable by the run. + + ``normalize_tools`` flattens any iterable tool collection, so a generator is a + supported run-level container. The framework materializes it exactly once when it + builds the middleware context; snapshot, per-call projection, and the run itself + all see the same tools. (Previously the startup projection exhausted the iterable + and enabling agent-hooks silently removed every run-level tool it contained.) + """ + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) + + one_shot = (item for item in [weather_tool]) + if route == "tools_kwarg": + response = await agent.run("Get weather for Seattle", tools=one_shot) + else: + response = await agent.run("Get weather for Seattle", options={"tools": one_shot}) + + # The run kept its tools: the tool call resolved and the loop completed. + assert response.text == "Final response" + assert weather_tool_calls == ["Seattle"] + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool"] + for pre_model in guard.contexts_for("pre_model_call"): + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool"] + + +@requires_sdk +async def test_nested_one_shot_collection_survives_the_snapshot(chat_client_base: MockBaseChatClient) -> None: + """Nested one-shot containers are materialized too, at every level flattening walks. + + ``normalize_tools`` flattens iterable collections recursively, so a generator + nested inside a list is a supported shape; the run-start materialization must + cover it, or the startup projection drains the inner iterator and the run loses + that tool (the same silent degradation one level down). + """ + guard = AllowGuard() + chat_client_base.run_responses = [tool_call_response(), final_response()] + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) + + response = await agent.run("Get weather for Seattle", tools=[(item for item in [weather_tool])]) + + assert response.text == "Final response" + assert weather_tool_calls == ["Seattle"] + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool"] + for pre_model in guard.contexts_for("pre_model_call"): + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool"] + + +@requires_sdk +async def test_snapshot_and_per_call_agree_when_both_run_tool_routes_are_supplied( + chat_client_base: MockBaseChatClient, +) -> None: + """With both ``tools=`` and ``options={"tools": ...}`` supplied, the named parameter + wins everywhere: ``tools_registered`` and the per-call ``tools`` projection report + the same set the run executes. (Previously the losing options entry silently + overrode the request's tools, so the run-start view and the executed set + disagreed.)""" + + @tool(approval_mode="never_require") + def options_entry_tool(x: str) -> str: + """Arrives through the options dict and loses the precedence.""" + return x + + guard = AllowGuard() + agent = Agent(client=chat_client_base, middleware=[create_agent_hooks_middleware([guard])]) + + await agent.run("hello", tools=[weather_tool], options={"tools": [options_entry_tool]}) + + startup = guard.contexts_for("agent_startup")[0] + assert startup["agent_init"]["tools_registered"] == ["weather_tool"] + pre_model = guard.contexts_for("pre_model_call")[0] + assert [entry["name"] for entry in pre_model["tools"]] == ["weather_tool"] + + # endregion # region Deny-before-execution diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index 8676deb69d..7610ce2aaf 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -30,6 +30,7 @@ agent_middleware, chat_middleware, function_middleware, + tool, ) from agent_framework._sessions import InMemoryHistoryProvider @@ -3039,3 +3040,136 @@ async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[ # endregion + +# region Run-level tools as observed by middleware + + +async def test_agent_middleware_observes_and_controls_original_tool_objects() -> None: + """Agent middleware sees the caller's original tool objects, identity intact. + + A guard that enforces policy by identity — for example rejecting one specific + privileged callable — must see exactly what the caller supplied. A one-shot + iterable container is materialized into a list before the pipeline (so observing + the tools does not consume the run's source), but its elements are never + converted, and an identity-based removal by the middleware governs what the run + executes. + """ + invoked: list[str] = [] + + def delete_all_data(target: str) -> str: + """Privileged tool.""" + invoked.append(target) + return f"deleted {target}" + + observed: dict[str, Any] = {} + + class IdentityGuard(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + observed["bare_identity"] = context.tools is delete_all_data + if isinstance(context.tools, list): + tools_list = cast("list[Any]", context.tools) + observed["element_identity"] = any(item is delete_all_data for item in tools_list) + # Identity-based enforcement: strip the privileged callable from the run. + context.tools = [item for item in tools_list if item is not delete_all_data] + await call_next() + + # A bare callable passes through to the pipeline untouched. + agent = Agent(client=MockBaseChatClient(), middleware=[IdentityGuard()]) + await agent.run("hi", tools=delete_all_data) + assert observed["bare_identity"] is True + + # A one-shot iterable is materialized (outer container only): the middleware sees + # the original callable, removes it by identity, and even though the model then + # requests it, the privileged tool is never invoked. + client = MockBaseChatClient() + client.run_responses = [ + ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="call_del", name="delete_all_data", arguments='{"target": "prod"}' + ) + ], + ) + ] + ), + ChatResponse(messages=[Message(role="assistant", contents=["done"])]), + ] + agent2 = Agent(client=client, middleware=[IdentityGuard()]) + await agent2.run("hi", tools=(item for item in [delete_all_data])) + assert observed["element_identity"] is True + assert invoked == [] + + +async def test_named_tools_parameter_wins_over_options_tools_entry() -> None: + """When both run-level routes are supplied, the named parameter wins end to end. + + Previously the losing ``options["tools"]`` entry survived in the remaining options + and rode ``**opts`` into the request, silently overriding the resolved tool list — + the run executed one set while the run-start view reported another. + """ + + @tool(approval_mode="never_require") + def kw_tool(x: str) -> str: + """Named-parameter tool.""" + return x + + @tool(approval_mode="never_require") + def opt_tool(x: str) -> str: + """Options-entry tool.""" + return x + + captured: dict[str, Any] = {} + + class CaptureChatMiddleware(ChatMiddleware): + async def process(self, context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + captured["tools"] = list((context.options or {}).get("tools") or []) + await call_next() + + agent = Agent(client=MockBaseChatClient(), middleware=[CaptureChatMiddleware()]) + await agent.run("hi", tools=[kw_tool], options={"tools": [opt_tool]}) + + assert [getattr(item, "name", None) for item in captured["tools"]] == ["kw_tool"] + + +async def test_losing_run_tool_route_is_never_iterated() -> None: + """Only the winning run-level route is materialized; the loser stays untouched. + + A losing one-shot ``options["tools"]`` source must not be consumed (or allowed to + raise / trigger side effects) when the named parameter wins: the framework drops + the losing entry without iterating it, and its owner can still consume it after + the run. + """ + + @tool(approval_mode="never_require") + def winning_tool(x: str) -> str: + """Named-parameter tool.""" + return x + + @tool(approval_mode="never_require") + def losing_tool(x: str) -> str: + """Options-entry tool.""" + return x + + iterations: list[str] = [] + + def losing_source() -> Any: + iterations.append("iterated") + yield losing_tool + + class Passthrough(AgentMiddleware): + async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + await call_next() + + losing = losing_source() + agent = Agent(client=MockBaseChatClient(), middleware=[Passthrough()]) + await agent.run("hi", tools=[winning_tool], options={"tools": losing}) + + assert iterations == [] + # The source still belongs to its owner, un-drained. + assert list(losing) == [losing_tool] + + +# endregion