Fix nested agent tool run lookup for call_id collisions

Author: seratch

src/agents/agent_tool_state.py

@@ -8,26 +8,52 @@
     from .result import RunResult, RunResultStreaming
 
 # Ephemeral maps linking tool call objects to nested agent results within the same run.
-# Store by object identity to avoid collisions when call IDs repeat across tool calls.
+# Store by object identity, and index by a stable signature to avoid call ID collisions.
 _agent_tool_run_results_by_obj: dict[int, RunResult | RunResultStreaming] = {}
-_agent_tool_run_results_by_call_id: dict[str, set[int]] = {}
+_agent_tool_run_results_by_signature: dict[
+    tuple[str, str, str, str, str | None, str | None],
+    set[int],
+] = {}
+_agent_tool_run_result_signature_by_obj: dict[
+    int,
+    tuple[str, str, str, str, str | None, str | None],
+] = {}
 
 
-def _index_agent_tool_run_result(tool_call_id: str, tool_call_obj_id: int) -> None:
-    """Track tool call objects by call ID for fallback lookup."""
-    _agent_tool_run_results_by_call_id.setdefault(tool_call_id, set()).add(tool_call_obj_id)
+def _tool_call_signature(
+    tool_call: ResponseFunctionToolCall,
+) -> tuple[str, str, str, str, str | None, str | None]:
+    """Build a stable signature for fallback lookup across tool call instances."""
+    return (
+        tool_call.call_id,
+        tool_call.name,
+        tool_call.arguments,
+        tool_call.type,
+        tool_call.id,
+        tool_call.status,
+    )
+
+
+def _index_agent_tool_run_result(
+    tool_call: ResponseFunctionToolCall, tool_call_obj_id: int
+) -> None:
+    """Track tool call objects by signature for fallback lookup."""
+    signature = _tool_call_signature(tool_call)
+    _agent_tool_run_result_signature_by_obj[tool_call_obj_id] = signature
+    _agent_tool_run_results_by_signature.setdefault(signature, set()).add(tool_call_obj_id)
 
 
-def _drop_agent_tool_run_result(tool_call_id: str | None, tool_call_obj_id: int) -> None:
+def _drop_agent_tool_run_result(tool_call_obj_id: int) -> None:
     """Remove a tool call object from the fallback index."""
-    if tool_call_id is None:
+    signature = _agent_tool_run_result_signature_by_obj.pop(tool_call_obj_id, None)
+    if signature is None:
         return
-    call_ids = _agent_tool_run_results_by_call_id.get(tool_call_id)
-    if not call_ids:
+    candidate_ids = _agent_tool_run_results_by_signature.get(signature)
+    if not candidate_ids:
         return
-    call_ids.discard(tool_call_obj_id)
-    if not call_ids:
-        _agent_tool_run_results_by_call_id.pop(tool_call_id, None)
+    candidate_ids.discard(tool_call_obj_id)
+    if not candidate_ids:
+        _agent_tool_run_results_by_signature.pop(signature, None)
 
 
 def record_agent_tool_run_result(
@@ -36,32 +62,29 @@ def record_agent_tool_run_result(
     """Store the nested agent run result by tool call identity."""
     tool_call_obj_id = id(tool_call)
     _agent_tool_run_results_by_obj[tool_call_obj_id] = run_result
-    if isinstance(tool_call.call_id, str):
-        _index_agent_tool_run_result(tool_call.call_id, tool_call_obj_id)
+    _index_agent_tool_run_result(tool_call, tool_call_obj_id)
 
 
 def consume_agent_tool_run_result(
     tool_call: ResponseFunctionToolCall,
 ) -> RunResult | RunResultStreaming | None:
-    """Return and drop the stored nested agent run result for the given tool call ID."""
+    """Return and drop the stored nested agent run result for the given tool call."""
     obj_id = id(tool_call)
     run_result = _agent_tool_run_results_by_obj.pop(obj_id, None)
     if run_result is not None:
-        _drop_agent_tool_run_result(tool_call.call_id, obj_id)
+        _drop_agent_tool_run_result(obj_id)
         return run_result
 
-    call_id = tool_call.call_id
-    if not call_id:
-        return None
-
-    candidate_ids = _agent_tool_run_results_by_call_id.get(call_id)
+    signature = _tool_call_signature(tool_call)
+    candidate_ids = _agent_tool_run_results_by_signature.get(signature)
     if not candidate_ids:
         return None
     if len(candidate_ids) != 1:
         return None
 
     candidate_id = next(iter(candidate_ids))
-    _agent_tool_run_results_by_call_id.pop(call_id, None)
+    _agent_tool_run_results_by_signature.pop(signature, None)
+    _agent_tool_run_result_signature_by_obj.pop(candidate_id, None)
     return _agent_tool_run_results_by_obj.pop(candidate_id, None)
 
 
@@ -74,11 +97,8 @@ def peek_agent_tool_run_result(
     if run_result is not None:
         return run_result
 
-    call_id = tool_call.call_id
-    if not call_id:
-        return None
-
-    candidate_ids = _agent_tool_run_results_by_call_id.get(call_id)
+    signature = _tool_call_signature(tool_call)
+    candidate_ids = _agent_tool_run_results_by_signature.get(signature)
     if not candidate_ids:
         return None
     if len(candidate_ids) != 1:
@@ -93,18 +113,17 @@ def drop_agent_tool_run_result(tool_call: ResponseFunctionToolCall) -> None:
     obj_id = id(tool_call)
     run_result = _agent_tool_run_results_by_obj.pop(obj_id, None)
     if run_result is not None:
-        _drop_agent_tool_run_result(tool_call.call_id, obj_id)
+        _drop_agent_tool_run_result(obj_id)
         return
 
-    call_id = tool_call.call_id
-    if not call_id:
-        return
-    candidate_ids = _agent_tool_run_results_by_call_id.get(call_id)
+    signature = _tool_call_signature(tool_call)
+    candidate_ids = _agent_tool_run_results_by_signature.get(signature)
     if not candidate_ids:
         return
     if len(candidate_ids) != 1:
         return
 
     candidate_id = next(iter(candidate_ids))
-    _agent_tool_run_results_by_call_id.pop(call_id, None)
+    _agent_tool_run_results_by_signature.pop(signature, None)
+    _agent_tool_run_result_signature_by_obj.pop(candidate_id, None)
     _agent_tool_run_results_by_obj.pop(candidate_id, None)