Skip to content

Commit 20d5469

Browse files
seratchhabema
andcommitted
feat: #2228 persist tool origin metadata in run items
Co-authored-by: Hassan Abu Alhaj <136383052+habema@users.noreply.github.com>
1 parent 0a100fb commit 20d5469

12 files changed

Lines changed: 515 additions & 12 deletions

File tree

src/agents/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@
160160
ShellToolLocalSkill,
161161
ShellToolSkillReference,
162162
Tool,
163+
ToolOrigin,
164+
ToolOriginType,
163165
ToolOutputFileContent,
164166
ToolOutputFileContentDict,
165167
ToolOutputImage,
@@ -393,6 +395,8 @@ def enable_verbose_stdout_logging():
393395
"MCPApprovalResponseItem",
394396
"ToolCallItem",
395397
"ToolCallOutputItem",
398+
"ToolOrigin",
399+
"ToolOriginType",
396400
"ReasoningItem",
397401
"ItemHelpers",
398402
"RunHooks",

src/agents/agent.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@
4646
FunctionToolResult,
4747
Tool,
4848
ToolErrorFunction,
49+
ToolOrigin,
50+
ToolOriginType,
4951
_build_handled_function_tool_error_handler,
5052
_build_wrapped_function_tool,
5153
_log_function_tool_invocation,
@@ -886,6 +888,11 @@ async def dispatch_stream_events() -> None:
886888
strict_json_schema=True,
887889
is_enabled=is_enabled,
888890
needs_approval=needs_approval,
891+
tool_origin=ToolOrigin(
892+
type=ToolOriginType.AGENT_AS_TOOL,
893+
agent_name=self.name,
894+
agent_tool_name=tool_name_resolved,
895+
),
889896
)
890897
run_agent_tool._is_agent_tool = True
891898
run_agent_tool._agent_instance = self

src/agents/items.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
from .exceptions import AgentsException, ModelBehaviorError
5555
from .logger import logger
5656
from .tool import (
57+
ToolOrigin,
5758
ToolOutputFileContent,
5859
ToolOutputImage,
5960
ToolOutputText,
@@ -358,6 +359,9 @@ class ToolCallItem(RunItemBase[Any]):
358359
title: str | None = None
359360
"""Optional short display label if known at item creation time."""
360361

362+
tool_origin: ToolOrigin | None = None
363+
"""Optional metadata describing the source of a function-tool-backed item."""
364+
361365

362366
ToolCallOutputTypes: TypeAlias = (
363367
FunctionCallOutput
@@ -382,6 +386,9 @@ class ToolCallOutputItem(RunItemBase[Any]):
382386

383387
type: Literal["tool_call_output_item"] = "tool_call_output_item"
384388

389+
tool_origin: ToolOrigin | None = None
390+
"""Optional metadata describing the source of a function-tool-backed item."""
391+
385392
def to_input_item(self) -> TResponseInputItem:
386393
"""Converts the tool output into an input item for the next model turn.
387394

src/agents/mcp/util.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
FunctionTool,
2828
Tool,
2929
ToolErrorFunction,
30+
ToolOrigin,
31+
ToolOriginType,
3032
ToolOutputImageDict,
3133
ToolOutputTextDict,
3234
_build_handled_function_tool_error_handler,
@@ -314,6 +316,10 @@ def to_function_tool(
314316
strict_json_schema=is_strict,
315317
needs_approval=needs_approval,
316318
mcp_title=resolve_mcp_tool_title(tool),
319+
tool_origin=ToolOrigin(
320+
type=ToolOriginType.MCP,
321+
mcp_server_name=server.name,
322+
),
317323
)
318324
return function_tool
319325

src/agents/run_internal/approvals.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
from ..agent import Agent
1515
from ..items import ItemHelpers, RunItem, ToolApprovalItem, ToolCallOutputItem, TResponseInputItem
16+
from ..tool import ToolOrigin
1617
from .items import ReasoningItemIdPolicy, run_item_to_input_item
1718

1819
# --------------------------
@@ -28,6 +29,7 @@ def append_approval_error_output(
2829
tool_name: str,
2930
call_id: str | None,
3031
message: str,
32+
tool_origin: ToolOrigin | None = None,
3133
) -> None:
3234
"""Emit a synthetic tool output so users see why an approval failed."""
3335
error_tool_call = _build_function_tool_call_for_approval_error(tool_call, tool_name, call_id)
@@ -36,6 +38,7 @@ def append_approval_error_output(
3638
output=message,
3739
raw_item=ItemHelpers.tool_call_output_item(error_tool_call, message),
3840
agent=agent,
41+
tool_origin=tool_origin,
3942
)
4043
)
4144

src/agents/run_internal/items.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,7 @@ def function_rejection_item(
308308
*,
309309
rejection_message: str = REJECTION_MESSAGE,
310310
scope_id: str | None = None,
311+
tool_origin: Any = None,
311312
) -> ToolCallOutputItem:
312313
"""Build a ToolCallOutputItem representing a rejected function tool call."""
313314
if isinstance(tool_call, ResponseFunctionToolCall):
@@ -316,6 +317,7 @@ def function_rejection_item(
316317
output=rejection_message,
317318
raw_item=ItemHelpers.tool_call_output_item(tool_call, rejection_message),
318319
agent=agent,
320+
tool_origin=tool_origin,
319321
)
320322

321323

src/agents/run_internal/run_loop.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,12 @@
1111
from collections.abc import Awaitable, Callable, Mapping
1212
from typing import Any, TypeVar, cast
1313

14-
from openai.types.responses import Response, ResponseCompletedEvent, ResponseOutputItemDoneEvent
14+
from openai.types.responses import (
15+
Response,
16+
ResponseCompletedEvent,
17+
ResponseFunctionToolCall,
18+
ResponseOutputItemDoneEvent,
19+
)
1520
from openai.types.responses.response_output_item import McpCall, McpListTools
1621
from openai.types.responses.response_prompt_param import ResponsePromptParam
1722
from openai.types.responses.response_reasoning_item import ResponseReasoningItem
@@ -64,7 +69,7 @@
6469
RawResponsesStreamEvent,
6570
RunItemStreamEvent,
6671
)
67-
from ..tool import FunctionTool, Tool, dispose_resolved_computers
72+
from ..tool import FunctionTool, Tool, dispose_resolved_computers, get_function_tool_origin
6873
from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span
6974
from ..tracing.model_tracing import get_model_tracing_impl
7075
from ..tracing.span_data import AgentSpanData, TaskSpanData
@@ -138,6 +143,7 @@
138143
from .streaming import stream_step_items_to_queue, stream_step_result_to_queue
139144
from .tool_actions import ApplyPatchAction, ComputerAction, LocalShellAction, ShellAction
140145
from .tool_execution import (
146+
build_litellm_json_tool_call,
141147
coerce_shell_call,
142148
execute_apply_patch_calls,
143149
execute_computer_actions,
@@ -1540,8 +1546,16 @@ async def rewind_model_request() -> None:
15401546
matched_tool = (
15411547
tool_map.get(tool_lookup_key) if tool_lookup_key is not None else None
15421548
)
1549+
if (
1550+
matched_tool is None
1551+
and output_schema is not None
1552+
and isinstance(output_item, ResponseFunctionToolCall)
1553+
and output_item.name == "json_tool_call"
1554+
):
1555+
matched_tool = build_litellm_json_tool_call(output_item)
15431556
tool_description: str | None = None
15441557
tool_title: str | None = None
1558+
tool_origin = None
15451559
if isinstance(output_item, McpCall):
15461560
metadata = hosted_mcp_tool_metadata.get(
15471561
(output_item.server_label, output_item.name)
@@ -1552,12 +1566,14 @@ async def rewind_model_request() -> None:
15521566
elif matched_tool is not None:
15531567
tool_description = getattr(matched_tool, "description", None)
15541568
tool_title = getattr(matched_tool, "_mcp_title", None)
1569+
tool_origin = get_function_tool_origin(matched_tool)
15551570

15561571
tool_item = ToolCallItem(
15571572
raw_item=cast(ToolCallItemTypes, output_item),
15581573
agent=public_agent,
15591574
description=tool_description,
15601575
title=tool_title,
1576+
tool_origin=tool_origin,
15611577
)
15621578
streamed_result._event_queue.put_nowait(
15631579
RunItemStreamEvent(item=tool_item, name="tool_called")

src/agents/run_internal/tool_execution.py

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,8 @@
7171
ShellCallOutcome,
7272
ShellCommandOutput,
7373
Tool,
74+
ToolOrigin,
75+
get_function_tool_origin,
7476
invoke_function_tool,
7577
maybe_invoke_function_tool_failure_error_function,
7678
resolve_computer,
@@ -1050,6 +1052,7 @@ async def on_invoke_tool(_ctx: ToolContext[Any], value: Any) -> Any:
10501052
on_invoke_tool=on_invoke_tool,
10511053
strict_json_schema=True,
10521054
is_enabled=True,
1055+
_emit_tool_origin=False,
10531056
)
10541057

10551058

@@ -1649,6 +1652,7 @@ async def _maybe_execute_tool_approval(
16491652
tool_call,
16501653
rejection_message=rejection_message,
16511654
scope_id=self.tool_state_scope_id,
1655+
tool_origin=get_function_tool_origin(func_tool),
16521656
),
16531657
)
16541658

@@ -1844,6 +1848,7 @@ def _build_function_tool_results(self) -> list[FunctionToolResult]:
18441848
output=result,
18451849
raw_item=ItemHelpers.tool_call_output_item(tool_run.tool_call, result),
18461850
agent=self.public_agent,
1851+
tool_origin=get_function_tool_origin(tool_run.function_tool),
18471852
)
18481853
else:
18491854
# Skip tool output until nested interruptions are resolved.
@@ -2060,14 +2065,22 @@ async def execute_approved_tools(
20602065
if isinstance(tool_name, str) and tool_name:
20612066
tool_map[tool_name] = tool
20622067

2063-
def _append_error(message: str, *, tool_call: Any, tool_name: str, call_id: str) -> None:
2068+
def _append_error(
2069+
message: str,
2070+
*,
2071+
tool_call: Any,
2072+
tool_name: str,
2073+
call_id: str,
2074+
tool_origin: ToolOrigin | None = None,
2075+
) -> None:
20642076
append_approval_error_output(
20652077
message=message,
20662078
tool_call=tool_call,
20672079
tool_name=tool_name,
20682080
call_id=call_id,
20692081
generated_items=generated_items,
20702082
agent=agent,
2083+
tool_origin=tool_origin,
20712084
)
20722085

20732086
async def _resolve_tool_run(
@@ -2095,14 +2108,25 @@ async def _resolve_tool_run(
20952108

20962109
call_id = extract_tool_call_id(tool_call)
20972110
if not call_id:
2111+
resolved_tool = tool_map.get(approval_key) if approval_key is not None else None
2112+
if resolved_tool is None and tool_namespace is None:
2113+
resolved_tool = tool_map.get(tool_name)
20982114
_append_error(
20992115
message="Tool approval item missing call ID.",
21002116
tool_call=tool_call,
21012117
tool_name=tool_name,
21022118
call_id="unknown",
2119+
tool_origin=(
2120+
get_function_tool_origin(resolved_tool)
2121+
if isinstance(resolved_tool, FunctionTool)
2122+
else None
2123+
),
21032124
)
21042125
return None
21052126

2127+
resolved_tool = tool_map.get(approval_key) if approval_key is not None else None
2128+
if resolved_tool is None and tool_namespace is None:
2129+
resolved_tool = tool_map.get(tool_name)
21062130
approval_status = context_wrapper.get_approval_status(
21072131
tool_name,
21082132
call_id,
@@ -2111,9 +2135,6 @@ async def _resolve_tool_run(
21112135
tool_lookup_key=tool_lookup_key,
21122136
)
21132137
if approval_status is False:
2114-
resolved_tool = tool_map.get(approval_key) if approval_key is not None else None
2115-
if resolved_tool is None and tool_namespace is None:
2116-
resolved_tool = tool_map.get(tool_name)
21172138
message = REJECTION_MESSAGE
21182139
if isinstance(resolved_tool, FunctionTool):
21192140
message = await resolve_approval_rejection_message(
@@ -2131,6 +2152,11 @@ async def _resolve_tool_run(
21312152
tool_call=tool_call,
21322153
tool_name=tool_name,
21332154
call_id=call_id,
2155+
tool_origin=(
2156+
get_function_tool_origin(resolved_tool)
2157+
if isinstance(resolved_tool, FunctionTool)
2158+
else None
2159+
),
21342160
)
21352161
return None
21362162

@@ -2140,12 +2166,15 @@ async def _resolve_tool_run(
21402166
tool_call=tool_call,
21412167
tool_name=tool_name,
21422168
call_id=call_id,
2169+
tool_origin=(
2170+
get_function_tool_origin(resolved_tool)
2171+
if isinstance(resolved_tool, FunctionTool)
2172+
else None
2173+
),
21432174
)
21442175
return None
21452176

2146-
tool = tool_map.get(approval_key) if approval_key is not None else None
2147-
if tool is None and tool_namespace is None:
2148-
tool = tool_map.get(tool_name)
2177+
tool = resolved_tool
21492178
if tool is None:
21502179
_append_error(
21512180
message=f"Tool '{display_tool_name}' not found.",

src/agents/run_internal/turn_resolution.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
LocalShellTool,
8181
ShellTool,
8282
Tool,
83+
get_function_tool_origin,
8384
)
8485
from ..tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
8586
from ..tracing import SpanError, handoff_span
@@ -777,6 +778,7 @@ async def _record_function_rejection(
777778
tool_call,
778779
rejection_message=rejection_message,
779780
scope_id=tool_state_scope_id,
781+
tool_origin=get_function_tool_origin(function_tool),
780782
)
781783
)
782784
if isinstance(call_id, str):
@@ -1791,11 +1793,19 @@ def _dump_output_item(raw_item: Any) -> dict[str, Any]:
17911793
func_tool = function_map.get(lookup_key) if lookup_key is not None else None
17921794
if func_tool is None:
17931795
if output_schema is not None and output.name == "json_tool_call":
1794-
items.append(ToolCallItem(raw_item=output, agent=agent))
1796+
synthetic_tool = build_litellm_json_tool_call(output)
1797+
items.append(
1798+
ToolCallItem(
1799+
raw_item=output,
1800+
agent=agent,
1801+
description=synthetic_tool.description,
1802+
tool_origin=get_function_tool_origin(synthetic_tool),
1803+
)
1804+
)
17951805
functions.append(
17961806
ToolRunFunction(
17971807
tool_call=output,
1798-
function_tool=build_litellm_json_tool_call(output),
1808+
function_tool=synthetic_tool,
17991809
)
18001810
)
18011811
continue
@@ -1816,6 +1826,7 @@ def _dump_output_item(raw_item: Any) -> dict[str, Any]:
18161826
agent=agent,
18171827
description=func_tool.description,
18181828
title=func_tool._mcp_title,
1829+
tool_origin=get_function_tool_origin(func_tool),
18191830
)
18201831
)
18211832
functions.append(

0 commit comments

Comments
 (0)