Skip to content

Commit 5237420

Browse files
seratchseratch
authored
fix(core): resolve tools after agent start hooks (#5124)
1 parent 9316424 commit 5237420

18 files changed

Lines changed: 1459 additions & 122 deletions

‎src/agents/_public_agent.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
22

33
from __future__ import annotations
44

5-
from .agent import Agent
5+
from typing import Any, TypeVar, cast
6+
7+
from .agent import Agent, AgentBase
8+
9+
TAgent = TypeVar("TAgent", bound=AgentBase[Any])
610

711
_PUBLIC_AGENT_ATTR = "_agents_public_agent"
812

@@ -13,9 +17,9 @@ def set_public_agent(execution_agent: Agent, public_agent: Agent) -> Agent:
1317
return execution_agent
1418

1519

16-
def get_public_agent(agent: Agent) -> Agent:
20+
def get_public_agent(agent: TAgent) -> TAgent:
1721
"""Return the user-visible agent identity for hooks, tool execution, and results."""
1822
public_agent = getattr(agent, _PUBLIC_AGENT_ATTR, None)
1923
if isinstance(public_agent, Agent):
20-
return public_agent
24+
return cast(TAgent, public_agent)
2125
return agent

‎src/agents/agent.py‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,14 @@ class AgentBase(Generic[TContext]):
212212
mcp_config: MCPConfig = field(default_factory=lambda: MCPConfig())
213213
"""Configuration for MCP servers."""
214214

215+
def __setattr__(self, name: str, value: Any) -> None:
216+
if name == "tools":
217+
from .run_internal.agent_tool_configuration import assign_agent_tools
218+
219+
assign_agent_tools(self, value)
220+
else:
221+
object.__setattr__(self, name, value)
222+
215223
async def _get_mcp_tool_reserved_names(
216224
self, run_context: RunContextWrapper[TContext]
217225
) -> set[str]:
@@ -253,6 +261,8 @@ def _use_mcp_handoff_snapshot(
253261

254262
async def get_mcp_tools(self, run_context: RunContextWrapper[TContext]) -> list[Tool]:
255263
"""Fetches the available tools from the MCP servers."""
264+
from ._public_agent import get_public_agent
265+
256266
convert_schemas_to_strict = self.mcp_config.get("convert_schemas_to_strict", False)
257267
failure_error_function = self.mcp_config.get(
258268
"failure_error_function", default_tool_error_function
@@ -267,14 +277,22 @@ async def get_mcp_tools(self, run_context: RunContextWrapper[TContext]) -> list[
267277
self.mcp_servers,
268278
convert_schemas_to_strict,
269279
run_context,
270-
self,
280+
get_public_agent(self),
271281
failure_error_function=failure_error_function,
272282
include_server_in_tool_names=include_server_in_tool_names,
273283
reserved_tool_names=reserved_tool_names,
274284
)
275285

276286
async def get_all_tools(self, run_context: RunContextWrapper[TContext]) -> list[Tool]:
277287
"""All agent tools, including MCP tools and function tools."""
288+
from ._public_agent import get_public_agent
289+
from .run_internal.agent_tool_configuration import (
290+
register_agent_tool_configuration,
291+
snapshot_agent_tools,
292+
)
293+
294+
register_agent_tool_configuration(get_public_agent(self))
295+
tools = snapshot_agent_tools(self)
278296
mcp_tools = await self.get_mcp_tools(run_context)
279297

280298
async def _check_tool_enabled(tool: Tool) -> bool:
@@ -284,16 +302,16 @@ async def _check_tool_enabled(tool: Tool) -> bool:
284302
attr = tool.is_enabled
285303
if isinstance(attr, bool):
286304
return attr
287-
res = attr(run_context, self)
305+
res = attr(run_context, get_public_agent(self))
288306
if inspect.isawaitable(res):
289307
return bool(await res)
290308
return bool(res)
291309

292-
tools = list(self.tools)
293310
results = await gather_with_cancel(*(_check_tool_enabled(t) for t in tools))
294311
enabled: list[Tool] = [t for t, ok in zip(tools, results, strict=False) if ok]
295312
all_tools: list[Tool] = prune_orphaned_tool_search_tools([*mcp_tools, *enabled])
296313
_validate_codex_tool_name_collisions(all_tools)
314+
snapshot_agent_tools(get_public_agent(self))
297315
return all_tools
298316

299317

@@ -1148,5 +1166,5 @@ async def get_prompt(
11481166
return await PromptUtil.to_model_input(
11491167
self.prompt,
11501168
run_context,
1151-
cast(Agent[TContext], get_public_agent(self)),
1169+
get_public_agent(self),
11521170
)

‎src/agents/lifecycle.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,19 @@ async def on_llm_end(
3737
async def on_agent_start(self, context: AgentHookContext[TContext], agent: TAgent) -> None:
3838
"""Called before the agent is invoked. Called each time the current agent changes.
3939
40+
Keep ``agent.tools`` unchanged when sharing an agent across concurrent runs.
41+
Tool-list replacement while concurrent runs use the same agent raises ``UserError``
42+
before newly resolved tools reach the model. In-place mutations of a shared list are not
43+
isolated. Changes to shared tool objects, including ``FunctionTool.is_enabled`` and
44+
``FunctionTool.needs_approval``, are not isolated or detected. For per-run configuration
45+
changes, use an independent agent, tools list, and tool objects. Copying only the list
46+
still shares its tool objects. Alternatively, keep the tool configuration unchanged and
47+
use context-based enablement and approval callbacks.
48+
49+
When restoring a serialized ``RunState``, rebuild any hook-created tools on the supplied
50+
agents before loading the state. Serialized state does not contain Python tool definitions,
51+
and start hooks are not replayed before pending approval calls are resolved.
52+
4053
Args:
4154
context: The agent hook context.
4255
agent: The agent that is about to be invoked.
@@ -114,6 +127,19 @@ async def on_start(self, context: AgentHookContext[TContext], agent: TAgent) ->
114127
"""Called before the agent is invoked. Called each time the running agent is changed to this
115128
agent.
116129
130+
Keep ``agent.tools`` unchanged when sharing an agent across concurrent runs.
131+
Tool-list replacement while concurrent runs use the same agent raises ``UserError``
132+
before newly resolved tools reach the model. In-place mutations of a shared list are not
133+
isolated. Changes to shared tool objects, including ``FunctionTool.is_enabled`` and
134+
``FunctionTool.needs_approval``, are not isolated or detected. For per-run configuration
135+
changes, use an independent agent, tools list, and tool objects. Copying only the list
136+
still shares its tool objects. Alternatively, keep the tool configuration unchanged and
137+
use context-based enablement and approval callbacks.
138+
139+
When restoring a serialized ``RunState``, rebuild any hook-created tools on the supplied
140+
agents before loading the state. Serialized state does not contain Python tool definitions,
141+
and start hooks are not replayed before pending approval calls are resolved.
142+
117143
Args:
118144
context: The agent hook context.
119145
agent: This agent instance.

‎src/agents/run.py‎

Lines changed: 31 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@
8383
validate_output_guardrails_with_server_managed_conversation,
8484
validate_session_conversation_settings,
8585
)
86+
from .run_internal.agent_tool_configuration import agent_tool_configuration_run
8687
from .run_internal.approvals import approvals_from_step
8788
from .run_internal.blocked_output import (
8889
OUTPUT_GUARDRAIL_BLOCKED_TOOL_OUTPUT,
@@ -120,9 +121,7 @@
120121
_safe_redacted_persistence_error,
121122
cleanup_models_after_run,
122123
finalize_max_turns_handler_output,
123-
get_all_tools,
124124
get_output_schema,
125-
initialize_computer_tools,
126125
resolve_interrupted_turn,
127126
run_input_guardrails,
128127
run_output_guardrails,
@@ -563,7 +562,13 @@ async def run(
563562
redacted_error: BaseException | None = None
564563
try:
565564
try:
566-
return await self._run_impl(starting_agent, input, **kwargs)
565+
configuration_agent = (
566+
input._current_agent
567+
if isinstance(input, RunState) and input._current_agent is not None
568+
else starting_agent
569+
)
570+
with agent_tool_configuration_run(configuration_agent):
571+
return await self._run_impl(starting_agent, input, **kwargs)
567572
except BaseException as error:
568573
if not _is_error_data_redacted(error):
569574
raise
@@ -1482,11 +1487,6 @@ def _mark_response_hooks_started() -> None:
14821487
if not run_state._pending_input:
14831488
run_state._generated_items = list(generated_items)
14841489
run_state._session_items = list(session_items)
1485-
all_tools = await get_all_tools(execution_agent, context_wrapper)
1486-
all_tools = await initialize_computer_tools(
1487-
tools=all_tools, context_wrapper=context_wrapper
1488-
)
1489-
14901490
if current_span is None:
14911491
if (output_schema := get_output_schema(execution_agent)) is not None:
14921492
output_type_name = output_schema.name()
@@ -1699,7 +1699,6 @@ async def _save_max_turns_handler_output(
16991699
model_task = asyncio.create_task(
17001700
run_single_turn(
17011701
bindings=current_bindings,
1702-
all_tools=all_tools,
17031702
original_input=original_input,
17041703
generated_items=items_for_model,
17051704
hooks=hooks,
@@ -1774,7 +1773,6 @@ async def _save_max_turns_handler_output(
17741773
else:
17751774
turn_result = await run_single_turn(
17761775
bindings=current_bindings,
1777-
all_tools=all_tools,
17781776
original_input=original_input,
17791777
generated_items=items_for_model,
17801778
hooks=hooks,
@@ -2601,24 +2599,30 @@ def run_streamed(
26012599

26022600
# Kick off the actual agent loop in the background and return the streamed result object.
26032601
async def run_loop() -> None:
2604-
await start_streaming(
2605-
starting_input=input_for_result,
2606-
streamed_result=streamed_result,
2607-
starting_agent=starting_agent,
2608-
max_turns=max_turns,
2609-
hooks=hooks,
2610-
context_wrapper=context_wrapper,
2611-
run_config=run_config,
2612-
error_handlers=error_handlers,
2613-
previous_response_id=previous_response_id,
2614-
auto_previous_response_id=auto_previous_response_id,
2615-
conversation_id=conversation_id,
2616-
session=session,
2617-
run_state=run_state,
2618-
trace_workflow_name=trace_workflow_name,
2619-
is_resumed_state=is_resumed_state,
2620-
sandbox_runtime=sandbox_runtime,
2602+
configuration_agent = (
2603+
run_state._current_agent
2604+
if run_state is not None and run_state._current_agent is not None
2605+
else starting_agent
26212606
)
2607+
with agent_tool_configuration_run(configuration_agent):
2608+
await start_streaming(
2609+
starting_input=input_for_result,
2610+
streamed_result=streamed_result,
2611+
starting_agent=starting_agent,
2612+
max_turns=max_turns,
2613+
hooks=hooks,
2614+
context_wrapper=context_wrapper,
2615+
run_config=run_config,
2616+
error_handlers=error_handlers,
2617+
previous_response_id=previous_response_id,
2618+
auto_previous_response_id=auto_previous_response_id,
2619+
conversation_id=conversation_id,
2620+
session=session,
2621+
run_state=run_state,
2622+
trace_workflow_name=trace_workflow_name,
2623+
is_resumed_state=is_resumed_state,
2624+
sandbox_runtime=sandbox_runtime,
2625+
)
26222626

26232627
# Keep the outer task frame inside the boundary so it cannot retain run payloads.
26242628
streamed_result.run_loop_task = asyncio.create_task(

‎src/agents/run_internal/agent_tool_configuration.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
from __future__ import annotations
2+
3+
from collections.abc import Iterator
4+
from contextlib import contextmanager
5+
from contextvars import ContextVar
6+
from dataclasses import dataclass, field
7+
from threading import Lock
8+
from typing import TYPE_CHECKING, Any
9+
10+
from ..exceptions import UserError
11+
from ..tool import Tool
12+
13+
if TYPE_CHECKING:
14+
from ..agent import AgentBase
15+
16+
17+
@dataclass(eq=False)
18+
class _ToolResolution:
19+
agent: AgentBase[Any]
20+
initial_tools: tuple[Tool, ...]
21+
tools_assigned: bool = False
22+
conflicted: bool = False
23+
24+
25+
@dataclass
26+
class _ToolConfigurationRun:
27+
resolutions: dict[int, _ToolResolution] = field(default_factory=dict)
28+
closed: bool = False
29+
30+
31+
_current_run: ContextVar[_ToolConfigurationRun | None] = ContextVar(
32+
"agent_tool_configuration_run", default=None
33+
)
34+
_resolutions: dict[int, list[_ToolResolution]] = {}
35+
_resolution_lock = Lock()
36+
37+
38+
def _observe_tools(active: list[_ToolResolution], tools: tuple[Tool, ...]) -> None:
39+
if len(active) > 1 and any(
40+
resolution.tools_assigned
41+
or len(resolution.initial_tools) != len(tools)
42+
or any(
43+
before is not after
44+
for before, after in zip(resolution.initial_tools, tools, strict=False)
45+
)
46+
for resolution in active
47+
):
48+
for resolution in active:
49+
resolution.conflicted = True
50+
51+
52+
def _raise_conflict() -> None:
53+
raise UserError(
54+
"Agent.tools changed while concurrent runs were using the same Agent. "
55+
"Use a separate agent per run, for example agent.clone(tools=[*agent.tools]), "
56+
"or keep Agent.tools unchanged and use context-based tool enablement."
57+
)
58+
59+
60+
@contextmanager
61+
def agent_tool_configuration_run(agent: AgentBase[Any]) -> Iterator[None]:
62+
"""Own configuration checks across all turns without persisting runtime locks or state."""
63+
owner = _ToolConfigurationRun()
64+
token = _current_run.set(owner)
65+
try:
66+
register_agent_tool_configuration(agent)
67+
yield
68+
finally:
69+
with _resolution_lock:
70+
owner.closed = True
71+
for agent_id, resolution in owner.resolutions.items():
72+
active = _resolutions[agent_id]
73+
# Preserve changes from failed/cancelled callbacks for surviving runs.
74+
_observe_tools(active, tuple(resolution.agent.tools))
75+
active.remove(resolution)
76+
if not active:
77+
del _resolutions[agent_id]
78+
owner.resolutions.clear()
79+
_current_run.reset(token)
80+
81+
82+
def register_agent_tool_configuration(agent: AgentBase[Any]) -> None:
83+
owner = _current_run.get()
84+
if owner is None:
85+
return
86+
agent_id = id(agent)
87+
with _resolution_lock:
88+
if owner.closed or agent_id in owner.resolutions:
89+
return
90+
active = _resolutions.setdefault(agent_id, [])
91+
resolution = _ToolResolution(agent, tuple(agent.tools))
92+
active.append(resolution)
93+
owner.resolutions[agent_id] = resolution
94+
_observe_tools(active, resolution.initial_tools)
95+
96+
97+
def assign_agent_tools(agent: AgentBase[Any], tools: list[Tool]) -> None:
98+
"""Retain replacements even if a competing hook restores the original list."""
99+
with _resolution_lock:
100+
active = _resolutions.get(id(agent), [])
101+
if active and tools is not agent.tools:
102+
for resolution in active:
103+
resolution.tools_assigned = True
104+
if len(active) > 1:
105+
for resolution in active:
106+
resolution.conflicted = True
107+
object.__setattr__(agent, "tools", tools)
108+
109+
110+
def snapshot_agent_tools(agent: AgentBase[Any]) -> tuple[Tool, ...]:
111+
"""Capture one tool sequence and reject conflicts with overlapping resolution scopes."""
112+
with _resolution_lock:
113+
tools = tuple(agent.tools)
114+
active = _resolutions.get(id(agent), [])
115+
_observe_tools(active, tools)
116+
conflicted = any(resolution.conflicted for resolution in active)
117+
if conflicted:
118+
_raise_conflict()
119+
return tools
120+
121+
122+
@contextmanager
123+
def guard_agent_tool_configuration(agent: AgentBase[Any]) -> Iterator[None]:
124+
"""Check both sides of hooks/discovery; the run owner retains the record between turns."""
125+
register_agent_tool_configuration(agent)
126+
snapshot_agent_tools(agent)
127+
yield
128+
snapshot_agent_tools(agent)

0 commit comments

Comments
 (0)