Python: Reuse registered tools in Monty callbacks

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:

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):

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

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.