fix review comments

Author: seratch

src/agents/items.py

@@ -500,6 +500,9 @@ class ToolApprovalItem(RunItemBase[Any]):
     tool_namespace: str | None = None
     """Optional Responses API namespace for function-tool approvals."""
 
+    tool_origin: ToolOrigin | None = None
+    """Optional metadata describing where the approved tool call came from."""
+
     tool_lookup_key: FunctionToolLookupKey | None = field(
         default=None,
         kw_only=True,

src/agents/run_internal/run_loop.py

@@ -67,7 +67,14 @@
     RawResponsesStreamEvent,
     RunItemStreamEvent,
 )
-from ..tool import FunctionTool, Tool, dispose_resolved_computers, get_function_tool_origin
+from ..tool import (
+    FunctionTool,
+    Tool,
+    ToolOrigin,
+    ToolOriginType,
+    dispose_resolved_computers,
+    get_function_tool_origin,
+)
 from ..tracing import Span, SpanError, agent_span, get_current_trace
 from ..tracing.model_tracing import get_model_tracing_impl
 from ..tracing.span_data import AgentSpanData
@@ -1374,6 +1381,10 @@ async def rewind_model_request() -> None:
                         if metadata is not None:
                             tool_description = metadata.description
                             tool_title = metadata.title
+                        tool_origin = ToolOrigin(
+                            type=ToolOriginType.MCP,
+                            mcp_server_name=output_item.server_label,
+                        )
                     elif matched_tool is not None:
                         tool_description = getattr(matched_tool, "description", None)
                         tool_title = getattr(matched_tool, "_mcp_title", None)

src/agents/run_internal/tool_execution.py

@@ -995,6 +995,7 @@ async def resolve_approval_status(
     context_wrapper: RunContextWrapper[Any],
     tool_namespace: str | None = None,
     tool_lookup_key: FunctionToolLookupKey | None = None,
+    tool_origin: ToolOrigin | None = None,
     on_approval: Callable[[RunContextWrapper[Any], ToolApprovalItem], Any] | None = None,
 ) -> tuple[bool | None, ToolApprovalItem]:
     """Build approval item, run on_approval hook if needed, and return latest approval status."""
@@ -1003,6 +1004,7 @@ async def resolve_approval_status(
         raw_item=raw_item,
         tool_name=tool_name,
         tool_namespace=tool_namespace,
+        tool_origin=tool_origin,
         tool_lookup_key=tool_lookup_key,
     )
     approval_status = context_wrapper.get_approval_status(
@@ -1506,6 +1508,7 @@ async def _maybe_execute_tool_approval(
                 raw_item=raw_tool_call,
                 tool_name=func_tool.name,
                 tool_namespace=tool_namespace,
+                tool_origin=get_function_tool_origin(func_tool),
                 tool_lookup_key=tool_lookup_key,
                 _allow_bare_name_alias=should_allow_bare_name_approval_alias(
                     func_tool,

src/agents/run_internal/tool_planning.py

@@ -22,7 +22,7 @@
     ToolCallOutputItem,
 )
 from ..run_context import RunContextWrapper
-from ..tool import FunctionTool, MCPToolApprovalRequest
+from ..tool import FunctionTool, MCPToolApprovalRequest, get_function_tool_origin
 from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
 from .run_steps import (
     ToolRunApplyPatchCall,
@@ -410,11 +410,17 @@ async def _collect_runs_by_approval(
         if approval_status is True:
             approved_runs.append(run)
         else:
+            function_tool = get_mapping_or_attr(run, "function_tool")
             pending_item = existing_pending or ToolApprovalItem(
                 agent=agent,
                 raw_item=get_mapping_or_attr(run, "tool_call"),
                 tool_name=tool_name,
                 tool_namespace=get_tool_call_namespace(get_mapping_or_attr(run, "tool_call")),
+                tool_origin=(
+                    get_function_tool_origin(function_tool)
+                    if isinstance(function_tool, FunctionTool)
+                    else None
+                ),
                 tool_lookup_key=get_function_tool_lookup_key_for_call(
                     get_mapping_or_attr(run, "tool_call")
                 ),

src/agents/run_internal/turn_resolution.py

@@ -78,6 +78,8 @@
     LocalShellTool,
     ShellTool,
     Tool,
+    ToolOrigin,
+    ToolOriginType,
     get_function_tool_origin,
 )
 from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
@@ -1031,6 +1033,7 @@ def _add_unmatched_pending(approval: ToolApprovalItem) -> None:
             raw_item=run.tool_call,
             tool_name=run.function_tool.name,
             tool_namespace=get_tool_call_namespace(run.tool_call),
+            tool_origin=get_function_tool_origin(run.function_tool),
             tool_lookup_key=get_function_tool_lookup_key_for_call(run.tool_call),
             _allow_bare_name_alias=should_allow_bare_name_approval_alias(
                 run.function_tool,
@@ -1520,6 +1523,10 @@ def _dump_output_item(raw_item: Any) -> dict[str, Any]:
                     agent=agent,
                     description=metadata.description if metadata is not None else None,
                     title=metadata.title if metadata is not None else None,
+                    tool_origin=ToolOrigin(
+                        type=ToolOriginType.MCP,
+                        mcp_server_name=output.server_label,
+                    ),
                 )
             )
             tools_used.append("mcp")

src/agents/run_state.py

@@ -1214,6 +1214,8 @@ def _serialize_tool_approval_interruption(
         interruption_dict["tool_name"] = interruption.tool_name
     if interruption.tool_namespace is not None:
         interruption_dict["tool_namespace"] = interruption.tool_namespace
+    if interruption.tool_origin is not None:
+        interruption_dict["tool_origin"] = interruption.tool_origin.to_json_dict()
     tool_lookup_key = serialize_function_tool_lookup_key(
         getattr(interruption, "tool_lookup_key", None)
     )
@@ -1885,6 +1887,7 @@ def _deserialize_tool_approval_item(
 
     tool_name = item_data.get("tool_name")
     tool_namespace = item_data.get("tool_namespace")
+    tool_origin = _deserialize_tool_origin(item_data.get("tool_origin"))
     tool_lookup_key = deserialize_function_tool_lookup_key(item_data.get("tool_lookup_key"))
     allow_bare_name_alias = item_data.get("allow_bare_name_alias") is True
     raw_item = _deserialize_tool_approval_raw_item(raw_item_data)
@@ -1893,6 +1896,7 @@ def _deserialize_tool_approval_item(
         raw_item=raw_item,
         tool_name=tool_name,
         tool_namespace=tool_namespace,
+        tool_origin=tool_origin,
         tool_lookup_key=tool_lookup_key,
         _allow_bare_name_alias=allow_bare_name_alias,
     )

tests/test_tool_origin.py

@@ -4,14 +4,16 @@
 import json
 import weakref
 from collections.abc import Sequence
-from typing import Any, TypeVar
+from typing import Any, TypeVar, cast
 
 import pytest
 from mcp import Tool as MCPTool
+from openai.types.responses.response_output_item import McpCall, McpListTools, McpListToolsTool
 from pydantic import BaseModel
 
 from agents import (
     Agent,
+    HostedMCPTool,
     ModelResponse,
     RunConfig,
     RunContextWrapper,
@@ -25,6 +27,7 @@
     Usage,
     function_tool,
 )
+from agents.items import MCPListToolsItem, ToolApprovalItem
 from agents.mcp import MCPUtil
 from agents.run_internal import run_loop
 from agents.run_internal.run_loop import get_output_schema
@@ -48,6 +51,22 @@ class StructuredOutputPayload(BaseModel):
     status: str
 
 
+def _make_hosted_mcp_list_tools(server_label: str, tool_name: str) -> McpListTools:
+    return McpListTools(
+        id=f"list_{server_label}",
+        server_label=server_label,
+        tools=[
+            McpListToolsTool(
+                name=tool_name,
+                input_schema={},
+                description="Search the docs.",
+                annotations={"title": "Search Docs"},
+            )
+        ],
+        type="mcp_list_tools",
+    )
+
+
 @pytest.mark.asyncio
 async def test_runner_attaches_function_tool_origin_to_call_and_output_items() -> None:
     model = FakeModel()
@@ -184,6 +203,104 @@ async def test_streamed_tool_call_item_includes_local_mcp_origin() -> None:
     )
 
 
+def test_process_model_response_attaches_hosted_mcp_tool_origin() -> None:
+    agent = Agent(name="hosted-mcp")
+    hosted_tool = HostedMCPTool(
+        tool_config=cast(
+            Any,
+            {
+                "type": "mcp",
+                "server_label": "docs_server",
+                "server_url": "https://example.com/mcp",
+            },
+        )
+    )
+    existing_items = [
+        MCPListToolsItem(
+            agent=agent,
+            raw_item=_make_hosted_mcp_list_tools("docs_server", "search_docs"),
+        )
+    ]
+    response = ModelResponse(
+        output=[
+            McpCall(
+                id="mcp_call_1",
+                arguments="{}",
+                name="search_docs",
+                server_label="docs_server",
+                type="mcp_call",
+                status="completed",
+            )
+        ],
+        usage=Usage(),
+        response_id="resp_hosted_mcp",
+    )
+
+    processed = run_loop.process_model_response(
+        agent=agent,
+        all_tools=[hosted_tool],
+        response=response,
+        output_schema=None,
+        handoffs=[],
+        existing_items=existing_items,
+    )
+
+    tool_call_item = _first_item(processed.new_items, ToolCallItem)
+    assert tool_call_item.tool_origin == ToolOrigin(
+        type=ToolOriginType.MCP,
+        mcp_server_name="docs_server",
+    )
+
+
+@pytest.mark.asyncio
+async def test_streamed_tool_call_item_includes_hosted_mcp_origin() -> None:
+    model = FakeModel()
+    hosted_tool = HostedMCPTool(
+        tool_config=cast(
+            Any,
+            {
+                "type": "mcp",
+                "server_label": "docs_server",
+                "server_url": "https://example.com/mcp",
+            },
+        )
+    )
+    agent = Agent(name="stream-hosted-mcp", model=model, tools=[hosted_tool])
+    model.add_multiple_turn_outputs(
+        [
+            [
+                _make_hosted_mcp_list_tools("docs_server", "search_docs"),
+                McpCall(
+                    id="mcp_call_stream_1",
+                    arguments="{}",
+                    name="search_docs",
+                    server_label="docs_server",
+                    type="mcp_call",
+                    status="completed",
+                ),
+            ],
+            [get_text_message("done")],
+        ]
+    )
+
+    result = Runner.run_streamed(agent, input="hello")
+    seen_tool_item: ToolCallItem | None = None
+    async for event in result.stream_events():
+        if (
+            event.type == "run_item_stream_event"
+            and isinstance(event.item, ToolCallItem)
+            and isinstance(event.item.raw_item, McpCall)
+        ):
+            seen_tool_item = event.item
+            break
+
+    assert seen_tool_item is not None
+    assert seen_tool_item.tool_origin == ToolOrigin(
+        type=ToolOriginType.MCP,
+        mcp_server_name="docs_server",
+    )
+
+
 def test_local_mcp_tool_origin_does_not_retain_server_object() -> None:
     server = FakeMCPServer(server_name="docs_server")
     function_tool = MCPUtil.to_function_tool(
@@ -332,3 +449,52 @@ async def test_run_state_from_json_reads_legacy_1_5_without_tool_origin() -> Non
     restored_item = _first_item(restored._generated_items, ToolCallItem)
     assert restored_item.description == "Legacy tool"
     assert restored_item.tool_origin is None
+
+
+@pytest.mark.asyncio
+async def test_run_state_roundtrip_preserves_tool_origin_on_approval_interruptions() -> None:
+    agent = Agent(name="approval-origin")
+    state: RunState[Any, Agent[Any]] = make_run_state(agent)
+    state._generated_items.append(
+        ToolApprovalItem(
+            agent=agent,
+            raw_item=make_tool_call(call_id="call_approval", name="approval_tool"),
+            tool_name="approval_tool",
+            tool_origin=ToolOrigin(type=ToolOriginType.FUNCTION),
+        )
+    )
+
+    restored = await roundtrip_state(agent, state)
+
+    approval_item = _first_item(restored._generated_items, ToolApprovalItem)
+    assert approval_item.tool_origin == ToolOrigin(type=ToolOriginType.FUNCTION)
+
+
+@pytest.mark.asyncio
+async def test_run_state_from_json_reads_legacy_1_6_approval_without_tool_origin() -> None:
+    agent = Agent(name="approval-origin-legacy")
+    state: RunState[Any, Agent[Any]] = make_run_state(agent)
+    state._generated_items.append(
+        ToolApprovalItem(
+            agent=agent,
+            raw_item=make_tool_call(call_id="call_legacy_approval", name="approval_tool"),
+            tool_name="approval_tool",
+            tool_origin=ToolOrigin(type=ToolOriginType.FUNCTION),
+        )
+    )
+
+    restored = await roundtrip_state(
+        agent,
+        state,
+        mutate_json=lambda data: {
+            **data,
+            "$schemaVersion": "1.6",
+            "generated_items": [
+                {key: value for key, value in item.items() if key != "tool_origin"}
+                for item in data["generated_items"]
+            ],
+        },
+    )
+
+    approval_item = _first_item(restored._generated_items, ToolApprovalItem)
+    assert approval_item.tool_origin is None