Python: Preserve method tool invocation limits

python/packages/core/agent_framework/_tools.py

@@ -583,6 +583,21 @@ def _format_tool_parameters(  # pyright: ignore[reportUnusedFunction]
 ClassT = TypeVar("ClassT", bound="SerializationMixin")
 
 
+class _BoundToolCache:
+    """Keep bound tools on their owner without serializing runtime bindings."""
+
+    def __init__(self, owner: Any = None) -> None:
+        self.owner = owner
+        self.tools: dict[FunctionTool, FunctionTool] = {}
+
+    def __reduce__(self) -> tuple[type[_BoundToolCache], tuple[()]]:
+        return type(self), ()
+
+    def to_dict(self) -> dict[str, Any]:
+        """Exclude runtime bindings when the owner is converted to JSON."""
+        return {}
+
+
 class FunctionTool(SerializationMixin):
     """A tool that wraps a Python function to make it callable by AI models.
 
@@ -643,6 +658,7 @@ class WeatherArgs(BaseModel):
         "_schema_supplied",
         "_invoke_sync_on_event_loop",
     }
+    _PICKLE_OMIT_FIELDS = SerializationMixin._PICKLE_OMIT_FIELDS | {"_invocation_duration_histogram"}
 
     def __init__(
         self,
@@ -765,6 +781,11 @@ def __init__(
         self.type: Literal["function_tool"] = "function_tool"
         self.result_parser = result_parser
 
+    def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None:
+        """Restore pickle state and recreate runtime-only metrics."""
+        super().__setstate__(state)
+        self._invocation_duration_histogram = _default_histogram()
+
     def _discover_injected_parameters(self) -> None:
         """Inspect the wrapped function for runtime injection parameters."""
         func = self.func.func if isinstance(self.func, FunctionTool) else self.func
@@ -811,14 +832,16 @@ def __get__(self, obj: Any, objtype: type | None = None) -> FunctionTool:
         """Implement the descriptor protocol to support bound methods.
 
         When a FunctionTool is accessed as an attribute of a class instance,
-        this method is called to bind the instance to the function.
+        this method is called to bind the instance to the function. Instances
+        with a ``__dict__`` cache their bound tools across attribute accesses.
+        For slots-only instances, retain the returned tool to preserve its counters.
 
         Args:
             obj: The instance that owns the descriptor, or None for class access.
             objtype: The type that owns the descriptor.
 
         Returns:
-            A new FunctionTool with the instance bound to the wrapped function.
+            A FunctionTool with the instance bound to the wrapped function.
         """
         if obj is None:
             # Accessed from the class, not an instance
@@ -829,9 +852,19 @@ def __get__(self, obj: Any, objtype: type | None = None) -> FunctionTool:
             sig = inspect.signature(self.func)
             params = list(sig.parameters.keys())
             if params and params[0] in {"self", "cls"}:
-                # Create a new FunctionTool with the bound method
-                import copy
-
+                instance_dict: dict[str, Any] | None = getattr(obj, "__dict__", None)
+                if isinstance(instance_dict, dict):
+                    cache = instance_dict.get("_agent_framework_bound_tools")
+                    # A shallow-copied owner must not reuse the original owner's bindings.
+                    if not isinstance(cache, _BoundToolCache) or cache.owner is not obj:
+                        cache = instance_dict["_agent_framework_bound_tools"] = _BoundToolCache(obj)
+                    if self not in cache.tools:
+                        bound_func = copy.copy(self)
+                        bound_func._instance = obj
+                        cache.tools[self] = bound_func
+                    return cache.tools[self]
+
+                # shortcut: slots-only owners cannot store bindings; retain the returned tool for persistent limits.
                 bound_func = copy.copy(self)
                 bound_func._instance = obj
                 return bound_func

python/packages/core/tests/core/test_tools.py

@@ -1,9 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 import asyncio
 import copy
+import gc
 import logging
+import pickle
 import threading
+import weakref
 from contextvars import ContextVar
+from dataclasses import dataclass, field
 from datetime import date, datetime
 from typing import Annotated, Any, Literal, get_args, get_origin
 from unittest.mock import Mock
@@ -799,6 +803,205 @@ def test_tool(self, x: int, y: int) -> int:
     assert test_tool(1, 2) == 3
 
 
+def _pickleable_tool_value(self: Any) -> str:
+    return self.value
+
+
+@dataclass(eq=False)
+class _PickleableToolOwner:
+    value: str = "pong"
+    ping = FunctionTool(name="ping", func=_pickleable_tool_value, input_model={}, max_invocations=1)
+
+
+async def test_method_tool_max_invocations_persists_across_attribute_access() -> None:
+    """Method tool invocation limits apply across repeated attribute access."""
+    from agent_framework.exceptions import ToolException
+
+    class Tools:
+        @tool(max_invocations=1)
+        def ping(self) -> str:
+            return "pong"
+
+    tools = Tools()
+
+    assert await tools.ping.invoke(skip_parsing=True) == "pong"
+    assert tools.ping.invocation_count == 1
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await tools.ping.invoke(skip_parsing=True)
+
+
+async def test_method_tool_max_invocation_exceptions_persists_across_attribute_access() -> None:
+    """Method tool exception limits apply across repeated attribute access."""
+    from agent_framework.exceptions import ToolException
+
+    class Tools:
+        @tool(max_invocation_exceptions=1)
+        def fail(self) -> None:
+            raise RuntimeError("boom")
+
+    tools = Tools()
+
+    with pytest.raises(RuntimeError, match="boom"):
+        await tools.fail.invoke(skip_parsing=True)
+    assert tools.fail.invocation_exception_count == 1
+    with pytest.raises(ToolException, match="maximum exception limit"):
+        await tools.fail.invoke(skip_parsing=True)
+
+
+@pytest.mark.parametrize("hashable", [False, True])
+async def test_method_tool_equal_owners_have_independent_limits(hashable: bool) -> None:
+    from agent_framework.exceptions import ToolException
+
+    @dataclass(unsafe_hash=hashable)  # type: ignore[misc]
+    class Tools:
+        value: str = field(compare=False)
+
+        @tool(max_invocations=1)
+        def ping(self) -> str:
+            return self.value
+
+    first, second = Tools("first"), Tools("second")
+    assert first == second
+    assert first.ping is first.ping
+    assert first.ping is not second.ping
+    assert await first.ping.invoke(skip_parsing=True) == "first"
+    assert await second.ping.invoke(skip_parsing=True) == "second"
+    for owner in (first, second):
+        with pytest.raises(ToolException, match="maximum invocation limit"):
+            await owner.ping.invoke(skip_parsing=True)
+
+
+@pytest.mark.parametrize("has_dict", [False, True])
+async def test_method_tool_owner_without_weakref_support(has_dict: bool) -> None:
+    from agent_framework.exceptions import ToolException
+
+    class Tools:
+        __slots__ = ("__dict__",) if has_dict else ()
+
+        @tool(max_invocations=1)
+        def ping(self) -> str:
+            return "pong"
+
+    owner = Tools()
+    with pytest.raises(TypeError):
+        weakref.ref(owner)
+    bound_tool = owner.ping
+    assert await bound_tool.invoke(skip_parsing=True) == "pong"
+    if has_dict:
+        assert owner.ping is bound_tool
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await bound_tool.invoke(skip_parsing=True)
+
+
+def test_method_tool_cache_releases_owner() -> None:
+    owner = _PickleableToolOwner()
+    owner_ref = weakref.ref(owner)
+    assert owner.ping is owner.ping
+
+    del owner
+    gc.collect()
+
+    assert owner_ref() is None
+
+
+def test_method_tool_owner_json_serialization_excludes_bindings() -> None:
+    from agent_framework._serialization import make_json_safe
+
+    class Tools:
+        def __init__(self) -> None:
+            self.value = "pong"
+
+        @tool
+        def ping(self) -> str:
+            return self.value
+
+    owner = Tools()
+    assert owner.ping() == "pong"
+    assert make_json_safe(owner) == {"value": "pong", "_agent_framework_bound_tools": {}}
+
+
+async def test_retained_method_tool_keeps_owner_alive() -> None:
+    owner = _PickleableToolOwner()
+    owner_ref = weakref.ref(owner)
+    bound_tool = owner.ping
+    del owner
+    gc.collect()
+
+    assert owner_ref() is not None
+    assert await bound_tool.invoke(skip_parsing=True) == "pong"
+
+    del bound_tool
+    gc.collect()
+    assert owner_ref() is None
+
+
+@pytest.mark.parametrize("copy_owner", [copy.copy, copy.deepcopy])
+async def test_method_tool_copied_owner_rebinds(copy_owner) -> None:
+    from agent_framework.exceptions import ToolException
+
+    original = _PickleableToolOwner("original")
+    assert await original.ping.invoke(skip_parsing=True) == "original"
+
+    copied = copy_owner(original)
+    copied.value = "copied"
+    assert copied.ping is not original.ping
+    assert copied.ping is copied.ping
+    assert await copied.ping.invoke(skip_parsing=True) == "copied"
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await original.ping.invoke(skip_parsing=True)
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await copied.ping.invoke(skip_parsing=True)
+
+
+@pytest.mark.parametrize("bind_first", [False, True])
+async def test_function_tool_pickle_after_descriptor_access(bind_first: bool) -> None:
+    descriptor = FunctionTool(name="ping", func=_pickleable_tool_value, input_model={}, max_invocations=1)
+    owner = _PickleableToolOwner()
+    if bind_first:
+        assert await descriptor.__get__(owner).invoke(skip_parsing=True) == "pong"
+
+    restored = pickle.loads(pickle.dumps(descriptor))
+
+    assert restored.name == descriptor.name
+    assert await restored.__get__(owner).invoke(skip_parsing=True) == "pong"
+
+
+async def test_bound_function_tool_pickle_preserves_owner_and_count() -> None:
+    from agent_framework.exceptions import ToolException
+
+    descriptor = FunctionTool(name="ping", func=_pickleable_tool_value, input_model={}, max_invocations=2)
+    owner = _PickleableToolOwner("original")
+    bound_tool = descriptor.__get__(owner)
+    assert await bound_tool.invoke(skip_parsing=True) == "original"
+
+    restored = pickle.loads(pickle.dumps(bound_tool))
+    owner.value = "changed"
+
+    assert restored.invocation_count == 1
+    assert await restored.invoke(skip_parsing=True) == "original"
+    assert restored.invocation_count == 2
+    assert bound_tool.invocation_count == 1
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await restored.invoke(skip_parsing=True)
+
+
+async def test_method_tool_pickled_owner_rebuilds_cache() -> None:
+    from agent_framework.exceptions import ToolException
+
+    original = _PickleableToolOwner("original")
+    assert await original.ping.invoke(skip_parsing=True) == "original"
+
+    restored = pickle.loads(pickle.dumps(original))
+    restored.value = "restored"
+
+    assert restored.ping is restored.ping
+    assert await restored.ping.invoke(skip_parsing=True) == "restored"
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await original.ping.invoke(skip_parsing=True)
+    with pytest.raises(ToolException, match="maximum invocation limit"):
+        await restored.ping.invoke(skip_parsing=True)
+
+
 def test_tool_with_literal_type_parameter():
     """Test tool decorator with Literal type parameter (issue #2891)."""