From da284fb8a4dd133648e7a7ddbb3d6524f605e87d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 15 Sep 2026 14:15:22 +0200 Subject: [PATCH] Python: Reuse registered tools in Monty callbacks --- python/packages/monty/README.md | 8 ++ .../_execute_code_tool.py | 3 +- .../monty/tests/monty/test_monty_codeact.py | 113 ++++++++++++++++++ .../monty/test_monty_codeact_integration.py | 55 ++++++++- 4 files changed, 176 insertions(+), 3 deletions(-) diff --git a/python/packages/monty/README.md b/python/packages/monty/README.md index 47b781098f4..d6f5a8475a4 100644 --- a/python/packages/monty/README.md +++ b/python/packages/monty/README.md @@ -105,6 +105,14 @@ agent = Agent( ) ``` +### Host tool lifetime + +Registered `FunctionTool` instances retain their invocation and exception counters +across `execute_code` calls and provider runs. Their `max_invocations` and +`max_invocation_exceptions` limits use the same counters as direct invocations of +those instances. A provider's run-scoped snapshot captures tool membership; it +does not reset host tool counters. + ### File mounts and resource limits Mount host directories into the sandbox and cap execution resources: diff --git a/python/packages/monty/agent_framework_monty/_execute_code_tool.py b/python/packages/monty/agent_framework_monty/_execute_code_tool.py index 4d5e957e4a8..18e0ab85bd0 100644 --- a/python/packages/monty/agent_framework_monty/_execute_code_tool.py +++ b/python/packages/monty/agent_framework_monty/_execute_code_tool.py @@ -15,7 +15,6 @@ import json import mimetypes from collections.abc import Callable, Iterator, Sequence -from copy import copy from functools import partial from pathlib import Path, PurePosixPath from typing import Any, cast @@ -161,7 +160,7 @@ def _make_tool_callback(tool_obj: FunctionTool) -> Callable[..., Any]: receives real Python objects. ``FunctionTool.invoke`` accepts direct keyword arguments and handles both sync and async underlying functions internally. """ - return partial(copy(tool_obj).invoke, skip_parsing=True) + return partial(tool_obj.invoke, skip_parsing=True) class MontyExecuteCodeTool(FunctionTool): diff --git a/python/packages/monty/tests/monty/test_monty_codeact.py b/python/packages/monty/tests/monty/test_monty_codeact.py index b1ccb3828e0..35548130ae4 100644 --- a/python/packages/monty/tests/monty/test_monty_codeact.py +++ b/python/packages/monty/tests/monty/test_monty_codeact.py @@ -9,6 +9,7 @@ from __future__ import annotations +import asyncio import json import sys import types @@ -21,6 +22,7 @@ import pytest from agent_framework import Content, FunctionTool, Message, tool from agent_framework._sessions import SessionContext +from agent_framework.exceptions import ToolException from agent_framework_monty import MontyCodeActProvider, MontyExecuteCodeTool from agent_framework_monty import _execute_code_tool as execute_code_module @@ -851,3 +853,114 @@ def adder(a: Annotated[int, ""], b: Annotated[int, ""]) -> int: cid, payload = await bridge._invoke_tool(7, "adder", {"a": 6, "b": 7}) assert cid == 7 assert payload == {"return_value": 13}, payload + + +async def test_tool_callbacks_share_registered_invocation_limit() -> None: + @tool(max_invocations=1) + def limited() -> int: + return 42 + + first = execute_code_module._make_tool_callback(limited) + second = execute_code_module._make_tool_callback(limited) + assert await first() == 42 + assert limited.invocation_count == 1 + for callback in (first, second, execute_code_module._make_tool_callback(limited)): + with pytest.raises(ToolException, match="maximum invocation limit"): + await callback() + assert limited.invocation_count == 1 + + limited.invocation_count = 0 + assert await second() == 42 + assert limited.invocation_count == 1 + + +@pytest.mark.parametrize("is_async", [False, True]) +async def test_tool_callbacks_share_registered_exception_limit(is_async: bool) -> None: + def fail() -> int: + raise ValueError("expected failure") + + async def async_fail() -> int: + await asyncio.sleep(0) + return fail() + + failing = FunctionTool(name="failing", func=async_fail if is_async else fail, max_invocation_exceptions=1) + with pytest.raises(ValueError, match="expected failure"): + await execute_code_module._make_tool_callback(failing)() + with pytest.raises(ToolException, match="maximum exception limit"): + await execute_code_module._make_tool_callback(failing)() + assert failing.invocation_count == 1 + assert failing.invocation_exception_count == 1 + + +async def test_concurrent_tool_callbacks_share_registered_invocation_limit() -> None: + @tool(max_invocations=1) + async def limited() -> int: + await asyncio.sleep(0) + return 42 + + bridges = [ + bridge_module.InlineCodeBridge({"limited": execute_code_module._make_tool_callback(limited)}) for _ in range(8) + ] + results = await asyncio.gather(*(bridge._invoke_tool(index, "limited", {}) for index, bridge in enumerate(bridges))) + + assert sum(payload == {"return_value": 42} for _, payload in results) == 1 + assert sum(payload.get("exc_type") == "ToolException" for _, payload in results) == 7 + assert limited.invocation_count == 1 + + +async def test_tool_callback_preserves_bound_instance_and_result_parser() -> None: + parser = MagicMock(return_value=[Content.from_text("parsed")]) + + class Counter: + def __init__(self) -> None: + self.value = 0 + + @tool(max_invocations=1, result_parser=parser) + def increment(self) -> int: + self.value += 1 + return self.value + + first_owner = Counter() + second_owner = Counter() + first_tool = first_owner.increment + second_tool = second_owner.increment + + assert await execute_code_module._make_tool_callback(first_tool)() == 1 + with pytest.raises(ToolException, match="maximum invocation limit"): + await execute_code_module._make_tool_callback(first_tool)() + assert await execute_code_module._make_tool_callback(second_tool)() == 1 + assert first_owner.value == second_owner.value == 1 + assert first_tool.invocation_count == second_tool.invocation_count == 1 + assert first_tool.result_parser is parser + parser.assert_not_called() + + first_tool.invocation_count = 0 + assert await first_tool.invoke() == [Content.from_text("parsed")] + parser.assert_called_once_with(2) + + +async def test_provider_runs_share_registered_tool_invocation_limit() -> None: + @tool(max_invocations=1) + def limited() -> int: + return 42 + + provider = MontyCodeActProvider(tools=[limited]) + for run_index in range(2): + context = SessionContext(input_messages=[]) + await provider.before_run(agent=MagicMock(), session=None, context=context, state={}) + run_tool = context.tools[0] + assert isinstance(run_tool, MontyExecuteCodeTool) + assert run_tool.get_tools()[0] is limited + script = _set_script( + _FakeFunctionSnapshot(function_name="limited", call_id=1), + _FakeFutureSnapshot(pending_call_ids=[1]), + _FakeMontyComplete(), + ) + await run_tool._run_code(code="await limited()") + payload = script.resume_log[-1][2][1] + if run_index == 0: + assert payload == {"return_value": 42} + else: + assert payload["exc_type"] == "ToolException" + assert "maximum invocation limit" in payload["message"] + assert limited.invocation_count == 1 diff --git a/python/packages/monty/tests/monty/test_monty_codeact_integration.py b/python/packages/monty/tests/monty/test_monty_codeact_integration.py index 2eeffb02a67..c52fb0e00ee 100644 --- a/python/packages/monty/tests/monty/test_monty_codeact_integration.py +++ b/python/packages/monty/tests/monty/test_monty_codeact_integration.py @@ -17,7 +17,7 @@ from unittest.mock import MagicMock import pytest -from agent_framework import Agent, Content, Message, tool +from agent_framework import Agent, Content, FunctionTool, Message, tool from agent_framework._sessions import SessionContext from agent_framework_monty import MontyCodeActProvider, MontyExecuteCodeTool @@ -138,11 +138,64 @@ async def test_async_host_tool_is_awaited() -> None: assert any("ping" in text for text in texts) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("call", ["await limited()", "await call_tool('limited')"]) +async def test_host_tool_invocation_limit_across_code_executions(is_async: bool, call: str) -> None: + calls: list[int] = [] + + def limited() -> int: + calls.append(42) + return 42 + + async def async_limited() -> int: + await asyncio.sleep(0) + return limited() + + host_tool = FunctionTool(name="limited", func=async_limited if is_async else limited, max_invocations=1) + monty_tool = MontyExecuteCodeTool(tools=[host_tool]) + + first = await monty_tool._run_code(code=call) + assert _text_outputs(first) == ["42"] + for _ in range(2): + result = await monty_tool._run_code(code=call) + assert any(content.type == "error" for content in result) + assert calls == [42] + assert host_tool.invocation_count == 1 + + host_tool.invocation_count = 0 + result = await monty_tool._run_code(code=f"{call}\n{call}") + assert any(content.type == "error" for content in result) + assert calls == [42, 42] + assert host_tool.invocation_count == 1 + + # --------------------------------------------------------------------------- # Concurrency # --------------------------------------------------------------------------- +@pytest.mark.parametrize("separate_executions", [False, True]) +async def test_concurrent_code_calls_share_host_tool_invocation_limit(separate_executions: bool) -> None: + calls: list[int] = [] + + @tool(max_invocations=1) + async def limited() -> int: + calls.append(42) + await asyncio.sleep(0) + return 42 + + monty_tool = MontyExecuteCodeTool(tools=[limited]) + if separate_executions: + results = await asyncio.gather(*(monty_tool._run_code(code="await limited()") for _ in range(4))) + assert sum(_text_outputs(result) == ["42"] for result in results) == 1 + assert sum(any(content.type == "error" for content in result) for result in results) == 3 + else: + result = await monty_tool._run_code(code="await asyncio.gather(limited(), limited(), limited(), limited())") + assert any(content.type == "error" for content in result) + assert calls == [42] + assert limited.invocation_count == 1 + + async def test_asyncio_gather_fans_out_tool_calls_concurrently() -> None: """Two async tools dispatched via ``asyncio.gather`` should run on the event loop in parallel.