diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index eb40fec7ee7..8cbd8451739 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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,20 @@ 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. + + Raises: + ValueError: If the instance already stores a non-cache value in + ``_agent_framework_bound_tools``. """ if obj is None: # Accessed from the class, not an instance @@ -829,9 +856,24 @@ 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") + if "_agent_framework_bound_tools" in instance_dict and not isinstance(cache, _BoundToolCache): + raise ValueError( + "Cannot cache bound tools: '_agent_framework_bound_tools' is already in use. " + "Rename the existing instance attribute to use method 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 diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 8dd1da6f053..95ea3a8b158 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 datetime import date, datetime from typing import Annotated, Any, Literal, get_args, get_origin from unittest.mock import Mock @@ -799,6 +803,224 @@ 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) + + +@pytest.mark.parametrize("existing_value", [None, {}, {"application_state": "keep"}, "keep"]) +def test_method_tool_cache_preserves_existing_owner_attribute(existing_value: Any) -> None: + class Tools: + def __init__(self) -> None: + self._agent_framework_bound_tools = existing_value + + @tool + def ping(self) -> str: + return "pong" + + owner = Tools() + + with pytest.raises(ValueError, match="_agent_framework_bound_tools"): + owner.ping() + + assert owner._agent_framework_bound_tools is existing_value + assert vars(owner) == {"_agent_framework_bound_tools": existing_value} + + +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)."""