fix: preserve plain run context for run-context tool wrappers (#2634)

Author: seratch

src/agents/tool.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import ast
 import asyncio
 import copy
 import dataclasses
@@ -9,8 +10,10 @@
 import weakref
 from collections.abc import Awaitable, Mapping
 from dataclasses import dataclass, field
+from types import UnionType
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     Callable,
     Generic,
@@ -19,6 +22,9 @@
     TypeVar,
     Union,
     cast,
+    get_args,
+    get_origin,
+    get_type_hints,
     overload,
 )
 
@@ -1373,19 +1379,131 @@ async def maybe_invoke_function_tool_failure_error_function(
     return result
 
 
+def _annotation_expr_name(expr: ast.expr) -> str | None:
+    """Return the unqualified type name for a string annotation expression node."""
+    if isinstance(expr, ast.Name):
+        return expr.id
+    if isinstance(expr, ast.Attribute):
+        return expr.attr
+    return None
+
+
+def _string_annotation_mentions_context_type(annotation: str, *, type_name: str) -> bool:
+    """Return True when a string annotation structurally references the given context type."""
+    try:
+        expression = ast.parse(annotation, mode="eval").body
+    except SyntaxError:
+        return False
+
+    return _annotation_expr_mentions_context_type(expression, type_name=type_name)
+
+
+def _annotation_expr_mentions_context_type(expr: ast.expr, *, type_name: str) -> bool:
+    """Return True when an annotation expression structurally references the given context type."""
+    if isinstance(expr, ast.Constant) and isinstance(expr.value, str):
+        return _string_annotation_mentions_context_type(expr.value, type_name=type_name)
+
+    if _annotation_expr_name(expr) == type_name:
+        return True
+
+    if isinstance(expr, ast.BinOp) and isinstance(expr.op, ast.BitOr):
+        return _annotation_expr_mentions_context_type(
+            expr.left, type_name=type_name
+        ) or _annotation_expr_mentions_context_type(expr.right, type_name=type_name)
+
+    if isinstance(expr, ast.Subscript):
+        wrapper_name = _annotation_expr_name(expr.value)
+        args = expr.slice.elts if isinstance(expr.slice, ast.Tuple) else (expr.slice,)
+
+        if wrapper_name == "Annotated":
+            return bool(args) and _annotation_expr_mentions_context_type(
+                args[0], type_name=type_name
+            )
+
+        if wrapper_name in {"Optional", "Union"}:
+            return any(
+                _annotation_expr_mentions_context_type(arg, type_name=type_name) for arg in args
+            )
+
+        return _annotation_expr_mentions_context_type(expr.value, type_name=type_name)
+
+    return False
+
+
+def _annotation_mentions_context_type(annotation: Any, *, context_type: type[Any]) -> bool:
+    """Return True when an annotation structurally references the given context type."""
+    if annotation is inspect.Signature.empty:
+        return False
+
+    if isinstance(annotation, str):
+        return _string_annotation_mentions_context_type(annotation, type_name=context_type.__name__)
+
+    origin = get_origin(annotation)
+
+    if annotation is context_type or origin is context_type:
+        return True
+
+    if origin is Annotated:
+        args = get_args(annotation)
+        return bool(args) and _annotation_mentions_context_type(args[0], context_type=context_type)
+
+    if origin in (Union, UnionType):
+        return any(
+            _annotation_mentions_context_type(arg, context_type=context_type)
+            for arg in get_args(annotation)
+        )
+
+    return False
+
+
+def _get_function_tool_invoke_context(
+    function_tool: FunctionTool,
+    context: ToolContext[Any],
+) -> ToolContext[Any] | RunContextWrapper[Any]:
+    """Choose the runtime context object to pass into a function tool wrapper.
+
+    Third-party wrappers may declare a narrower `RunContextWrapper` contract and then serialize
+    that object downstream. In those cases, passing the richer `ToolContext` can leak runtime-only
+    metadata such as agents or run config into incompatible serializers. When the wrapper
+    explicitly declares `RunContextWrapper`, preserve only the base context state.
+    """
+    try:
+        parameters = tuple(inspect.signature(function_tool.on_invoke_tool).parameters.values())
+    except (TypeError, ValueError):
+        return context
+
+    if not parameters:
+        return context
+
+    context_annotation = parameters[0].annotation
+    try:
+        resolved_annotations = get_type_hints(function_tool.on_invoke_tool, include_extras=True)
+    except Exception:
+        pass
+    else:
+        context_annotation = resolved_annotations.get(parameters[0].name, context_annotation)
+
+    if _annotation_mentions_context_type(context_annotation, context_type=ToolContext):
+        return context
+    if _annotation_mentions_context_type(context_annotation, context_type=RunContextWrapper):
+        return context._fork_with_tool_input(context.tool_input)
+    return context
+
+
 async def invoke_function_tool(
     *,
     function_tool: FunctionTool,
     context: ToolContext[Any],
     arguments: str,
 ) -> Any:
     """Invoke a function tool, enforcing timeout configuration when provided."""
+    invoke_context = _get_function_tool_invoke_context(function_tool, context)
     timeout_seconds = function_tool.timeout_seconds
     if timeout_seconds is None:
-        return await function_tool.on_invoke_tool(context, arguments)
+        return await function_tool.on_invoke_tool(cast(Any, invoke_context), arguments)
 
     tool_task: asyncio.Future[Any] = asyncio.ensure_future(
-        function_tool.on_invoke_tool(context, arguments)
+        function_tool.on_invoke_tool(cast(Any, invoke_context), arguments)
     )
     try:
         return await asyncio.wait_for(tool_task, timeout=timeout_seconds)

tests/test_tool_context.py

@@ -1,10 +1,14 @@
+from typing import Annotated, Any, cast
+
 import pytest
 from openai.types.responses import ResponseFunctionToolCall
 
 from agents import Agent
 from agents.run_config import RunConfig
 from agents.run_context import RunContextWrapper
+from agents.tool import FunctionTool, invoke_function_tool
 from agents.tool_context import ToolContext
+from agents.usage import Usage
 from tests.utils.hitl import make_context_wrapper
 
 
@@ -198,3 +202,154 @@ def test_tool_context_from_agent_context_prefers_explicit_run_config() -> None:
     )
 
     assert tool_ctx.run_config is explicit_run_config
+
+
+@pytest.mark.asyncio
+async def test_invoke_function_tool_passes_plain_run_context_when_requested() -> None:
+    captured_context: RunContextWrapper[str] | None = None
+
+    async def on_invoke_tool(ctx: RunContextWrapper[str], _input: str) -> str:
+        nonlocal captured_context
+        captured_context = ctx
+        return ctx.context
+
+    function_tool = FunctionTool(
+        name="plain_context_tool",
+        description="test",
+        params_json_schema={"type": "object", "properties": {}},
+        on_invoke_tool=on_invoke_tool,
+    )
+    tool_context = ToolContext(
+        context="Stormy",
+        usage=Usage(),
+        tool_name="plain_context_tool",
+        tool_call_id="call-1",
+        tool_arguments="{}",
+        agent=Agent(name="agent"),
+        run_config=RunConfig(model="gpt-4.1-mini"),
+        tool_input={"city": "Tokyo"},
+    )
+
+    result = await invoke_function_tool(
+        function_tool=function_tool,
+        context=tool_context,
+        arguments="{}",
+    )
+
+    assert result == "Stormy"
+    assert captured_context is not None
+    assert not isinstance(captured_context, ToolContext)
+    assert captured_context.context == "Stormy"
+    assert captured_context.usage is tool_context.usage
+    assert captured_context.tool_input == {"city": "Tokyo"}
+
+
+@pytest.mark.asyncio
+async def test_invoke_function_tool_preserves_tool_context_when_requested() -> None:
+    captured_context: ToolContext[str] | None = None
+
+    async def on_invoke_tool(ctx: ToolContext[str], _input: str) -> str:
+        nonlocal captured_context
+        captured_context = ctx
+        return ctx.tool_name
+
+    function_tool = FunctionTool(
+        name="tool_context_tool",
+        description="test",
+        params_json_schema={"type": "object", "properties": {}},
+        on_invoke_tool=on_invoke_tool,
+    )
+    tool_context = ToolContext(
+        context="Stormy",
+        usage=Usage(),
+        tool_name="tool_context_tool",
+        tool_call_id="call-2",
+        tool_arguments="{}",
+        agent=Agent(name="agent"),
+        run_config=RunConfig(model="gpt-4.1-mini"),
+    )
+
+    result = await invoke_function_tool(
+        function_tool=function_tool,
+        context=tool_context,
+        arguments="{}",
+    )
+
+    assert result == "tool_context_tool"
+    assert captured_context is tool_context
+
+
+@pytest.mark.asyncio
+async def test_invoke_function_tool_ignores_context_name_substrings_in_string_annotations() -> None:
+    captured_context: object | None = None
+
+    class MyRunContextWrapper:
+        pass
+
+    async def on_invoke_tool(ctx: "MyRunContextWrapper", _input: str) -> str:
+        nonlocal captured_context
+        captured_context = ctx
+        return "ok"
+
+    function_tool = FunctionTool(
+        name="substring_context_tool",
+        description="test",
+        params_json_schema={"type": "object", "properties": {}},
+        on_invoke_tool=cast(Any, on_invoke_tool),
+    )
+    tool_context = ToolContext(
+        context="Stormy",
+        usage=Usage(),
+        tool_name="substring_context_tool",
+        tool_call_id="call-3",
+        tool_arguments="{}",
+    )
+
+    result = await invoke_function_tool(
+        function_tool=function_tool,
+        context=tool_context,
+        arguments="{}",
+    )
+
+    assert result == "ok"
+    assert captured_context is tool_context
+
+
+@pytest.mark.asyncio
+async def test_invoke_function_tool_ignores_annotated_string_metadata_when_matching_context() -> (
+    None
+):
+    captured_context: ToolContext[str] | RunContextWrapper[str] | None = None
+
+    async def on_invoke_tool(
+        ctx: Annotated[RunContextWrapper[str], "ToolContext note"], _input: str
+    ) -> str:
+        nonlocal captured_context
+        captured_context = ctx
+        return ctx.context
+
+    function_tool = FunctionTool(
+        name="annotated_string_context_tool",
+        description="test",
+        params_json_schema={"type": "object", "properties": {}},
+        on_invoke_tool=on_invoke_tool,
+    )
+    tool_context = ToolContext(
+        context="Stormy",
+        usage=Usage(),
+        tool_name="annotated_string_context_tool",
+        tool_call_id="call-4",
+        tool_arguments="{}",
+        tool_input={"city": "Tokyo"},
+    )
+
+    result = await invoke_function_tool(
+        function_tool=function_tool,
+        context=tool_context,
+        arguments="{}",
+    )
+
+    assert result == "Stormy"
+    assert captured_context is not None
+    assert not isinstance(captured_context, ToolContext)
+    assert captured_context.tool_input == {"city": "Tokyo"}