From 4fabb35b7da012a21cc102d56df5816325ef3f9a Mon Sep 17 00:00:00 2001 From: barry <91018388+barry166@users.noreply.github.com> Date: Fri, 2 Oct 2026 00:13:29 +0800 Subject: [PATCH 1/2] Python: preserve method tool invocation limits --- .../packages/core/agent_framework/_tools.py | 22 ++++++++---- python/packages/core/tests/core/test_tools.py | 35 +++++++++++++++++++ 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 602da9dc4cb..f50460f5d74 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -12,6 +12,7 @@ import sys import typing import warnings +import weakref from collections import deque from collections.abc import ( AsyncIterable, @@ -95,6 +96,8 @@ logger = logging.getLogger("agent_framework") +_BOUND_TOOL_CACHE: weakref.WeakKeyDictionary[Any, weakref.WeakKeyDictionary[Any, Any]] = weakref.WeakKeyDictionary() + def _generate_function_call_occurrence_id() -> str: """Generate an Agent Framework identity for one function-call occurrence.""" @@ -730,7 +733,7 @@ def __get__(self, obj: Any, objtype: type | None = None) -> FunctionTool: 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 @@ -741,12 +744,17 @@ 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 - - bound_func = copy.copy(self) - bound_func._instance = obj - return bound_func + try: + bound_tools = _BOUND_TOOL_CACHE[self] + except KeyError: + bound_tools = _BOUND_TOOL_CACHE[self] = weakref.WeakKeyDictionary() + try: + return bound_tools[obj] + except KeyError: + bound_func = copy.copy(self) + bound_func._instance = obj + bound_tools[obj] = bound_func + return bound_func return self diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 3071e191513..14cc19f63f0 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -709,6 +709,41 @@ def test_tool(self, x: int, y: int) -> int: assert test_tool(1, 2) == 3 +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) + + def test_tool_with_literal_type_parameter(): """Test tool decorator with Literal type parameter (issue #2891).""" From 556da45b33b8b2a6070f871fa4fbad949ec93b79 Mon Sep 17 00:00:00 2001 From: barry <91018388+barry166@users.noreply.github.com> Date: Fri, 9 Oct 2026 14:52:10 +0800 Subject: [PATCH 2/2] Fix method tool caching for valid owner types and serialization Keep bound tools on their owner so unhashable and equal instances remain isolated and the cache does not retain owners globally. Recreate runtime bindings after pickle or copying, and preserve slots-only compatibility. Tested: 590 related tests passed, 1 skipped; Ruff and Pyright passed Constraint: Slots-only owners must retain the returned tool for persistent limits --- .../packages/core/agent_framework/_tools.py | 49 +++-- python/packages/core/tests/core/test_tools.py | 168 ++++++++++++++++++ 2 files changed, 202 insertions(+), 15 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index f50460f5d74..d412ab8fcb4 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -12,7 +12,6 @@ import sys import typing import warnings -import weakref from collections import deque from collections.abc import ( AsyncIterable, @@ -96,8 +95,6 @@ logger = logging.getLogger("agent_framework") -_BOUND_TOOL_CACHE: weakref.WeakKeyDictionary[Any, weakref.WeakKeyDictionary[Any, Any]] = weakref.WeakKeyDictionary() - def _generate_function_call_occurrence_id() -> str: """Generate an Agent Framework identity for one function-call occurrence.""" @@ -498,6 +495,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. @@ -726,7 +738,9 @@ 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. @@ -744,17 +758,22 @@ 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"}: - try: - bound_tools = _BOUND_TOOL_CACHE[self] - except KeyError: - bound_tools = _BOUND_TOOL_CACHE[self] = weakref.WeakKeyDictionary() - try: - return bound_tools[obj] - except KeyError: - bound_func = copy.copy(self) - bound_func._instance = obj - bound_tools[obj] = bound_func - return bound_func + 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 return self diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 14cc19f63f0..b2544cfbd14 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/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 typing import Annotated, Any, Literal, get_args, get_origin from unittest.mock import Mock @@ -709,6 +713,16 @@ 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 @@ -744,6 +758,160 @@ def fail(self) -> None: 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)."""