fix: #2664 drop orphan hosted shell calls before multi-turn replay (#2665)

Author: seratch

src/agents/run_internal/items.py

@@ -33,6 +33,7 @@
     "copy_input_items",
     "drop_orphan_function_calls",
     "ensure_input_item_format",
+    "prepare_model_input_items",
     "run_item_to_input_item",
     "run_items_to_input_items",
     "normalize_input_items_for_api",
@@ -86,7 +87,11 @@ def run_items_to_input_items(
     return converted
 
 
-def drop_orphan_function_calls(items: list[TResponseInputItem]) -> list[TResponseInputItem]:
+def drop_orphan_function_calls(
+    items: list[TResponseInputItem],
+    *,
+    pruning_indexes: set[int] | None = None,
+) -> list[TResponseInputItem]:
     """
     Remove tool call items that do not have corresponding outputs so resumptions or retries do not
     replay stale tool calls.
@@ -108,6 +113,9 @@ def drop_orphan_function_calls(items: list[TResponseInputItem]) -> list[TRespons
         if output_type is None:
             filtered.append(entry)
             continue
+        if pruning_indexes is not None and index not in pruning_indexes:
+            filtered.append(entry)
+            continue
         call_id = entry.get("call_id")
         if isinstance(call_id, str) and call_id in completed_call_ids.get(output_type, set()):
             filtered.append(entry)
@@ -145,6 +153,20 @@ def normalize_input_items_for_api(items: list[TResponseInputItem]) -> list[TResp
     return normalized
 
 
+def prepare_model_input_items(
+    caller_items: Sequence[TResponseInputItem],
+    generated_items: Sequence[TResponseInputItem] = (),
+) -> list[TResponseInputItem]:
+    """Normalize model input while pruning orphans only from runner-generated history."""
+    normalized_caller_items = normalize_input_items_for_api(list(caller_items))
+    if not generated_items:
+        return normalized_caller_items
+
+    normalized_generated_items = normalize_input_items_for_api(list(generated_items))
+    filtered_generated_items = drop_orphan_function_calls(normalized_generated_items)
+    return normalized_caller_items + filtered_generated_items
+
+
 def normalize_resumed_input(
     raw_input: str | list[TResponseInputItem],
 ) -> str | list[TResponseInputItem]:

src/agents/run_internal/oai_conversation.py

@@ -23,6 +23,7 @@
     drop_orphan_function_calls,
     fingerprint_input_item,
     normalize_input_items_for_api,
+    prepare_model_input_items,
     run_item_to_input_item,
 )
 
@@ -153,8 +154,7 @@ def hydrate_from_state(
 
         normalized_input = original_input
         if isinstance(original_input, list):
-            normalized = normalize_input_items_for_api(original_input)
-            normalized_input = drop_orphan_function_calls(normalized)
+            normalized_input = prepare_model_input_items(original_input)
 
         for item in ItemHelpers.input_to_new_input_list(normalized_input):
             if item is None:
@@ -404,13 +404,17 @@ def prepare_input(
         generated_items: list[RunItem],
     ) -> list[TResponseInputItem]:
         """Assemble the next model input while skipping duplicates and approvals."""
-        input_items: list[TResponseInputItem] = []
+        prepared_initial_items: list[TResponseInputItem] = []
+        prepared_generated_items: list[TResponseInputItem] = []
+        generated_item_sources: dict[int, TResponseInputItem] = {}
 
         if not self.sent_initial_input:
             initial_items = ItemHelpers.input_to_new_input_list(original_input)
-            input_items.extend(initial_items)
-            for item in initial_items:
-                self._register_prepared_item_source(item)
+            prepared_initial_items = normalize_input_items_for_api(initial_items)
+            for prepared_item, source_item in zip(
+                prepared_initial_items, initial_items, strict=False
+            ):
+                self._register_prepared_item_source(prepared_item, source_item)
             filtered_initials = []
             for item in initial_items:
                 if item is None or isinstance(item, (str, bytes)):
@@ -419,9 +423,11 @@ def prepare_input(
             self.remaining_initial_input = filtered_initials or None
             self.sent_initial_input = True
         elif self.remaining_initial_input:
-            input_items.extend(self.remaining_initial_input)
-            for item in self.remaining_initial_input:
-                self._register_prepared_item_source(item)
+            prepared_initial_items = normalize_input_items_for_api(self.remaining_initial_input)
+            for prepared_item, source_item in zip(
+                prepared_initial_items, self.remaining_initial_input, strict=False
+            ):
+                self._register_prepared_item_source(prepared_item, source_item)
 
         for item in generated_items:  # type: ignore[assignment]
             run_item: RunItem = cast(RunItem, item)
@@ -474,13 +480,23 @@ def prepare_input(
             ):
                 continue
 
-            input_items.append(converted_input_item)
-            self._register_prepared_item_source(
-                converted_input_item,
-                cast(TResponseInputItem, raw_item),
-            )
+            prepared_generated_items.append(converted_input_item)
+            generated_item_sources[id(converted_input_item)] = cast(TResponseInputItem, raw_item)
 
-        return input_items
+        normalized_generated_items = normalize_input_items_for_api(prepared_generated_items)
+        normalized_generated_sources = {
+            id(normalized_item): generated_item_sources[id(source_item)]
+            for normalized_item, source_item in zip(
+                normalized_generated_items, prepared_generated_items, strict=False
+            )
+        }
+        filtered_generated_items = drop_orphan_function_calls(normalized_generated_items)
+        for item in filtered_generated_items:
+            prepared_source_item = normalized_generated_sources.get(id(item))
+            if prepared_source_item is not None:
+                self._register_prepared_item_source(item, prepared_source_item)
+
+        return prepared_initial_items + filtered_generated_items
 
     def _register_prepared_item_source(
         self, prepared_item: TResponseInputItem, source_item: TResponseInputItem | None = None

src/agents/run_internal/run_loop.py

@@ -69,10 +69,7 @@
 from ..usage import Usage
 from ..util import _coro, _error_tracing
 from .agent_runner_helpers import apply_resumed_conversation_settings
-from .approvals import (
-    append_input_items_excluding_approvals,
-    approvals_from_step,
-)
+from .approvals import approvals_from_step
 from .error_handlers import (
     build_run_error_data,
     create_message_output_item,
@@ -93,8 +90,9 @@
     copy_input_items,
     deduplicate_input_items_preferring_latest,
     ensure_input_item_format,
-    normalize_input_items_for_api,
     normalize_resumed_input,
+    prepare_model_input_items,
+    run_items_to_input_items,
 )
 from .model_retry import (
     apply_retry_attempt_usage,
@@ -244,6 +242,16 @@ async def _should_persist_stream_items(
     return should_skip_session_save is False
 
 
+def _prepare_turn_input_items(
+    caller_input: str | list[TResponseInputItem],
+    generated_items: list[RunItem],
+    reasoning_item_id_policy: ReasoningItemIdPolicy | None,
+) -> list[TResponseInputItem]:
+    caller_items = ItemHelpers.input_to_new_input_list(caller_input)
+    continuation_items = run_items_to_input_items(generated_items, reasoning_item_id_policy)
+    return prepare_model_input_items(caller_items, continuation_items)
+
+
 def _complete_stream_interruption(
     streamed_result: RunResultStreaming,
     *,
@@ -1164,16 +1172,12 @@ def _tool_search_fingerprint(raw_item: Any) -> str:
             else 0,
         )
     else:
-        input = ItemHelpers.input_to_new_input_list(streamed_result.input)
-        append_input_items_excluding_approvals(
-            input,
+        input = _prepare_turn_input_items(
+            streamed_result.input,
             streamed_result._model_input_items,
             reasoning_item_id_policy,
         )
 
-    if isinstance(input, list):
-        input = normalize_input_items_for_api(input)
-
     filtered = await maybe_filter_model_input(
         agent=agent,
         run_config=run_config,
@@ -1512,23 +1516,7 @@ async def run_single_turn(
     if server_conversation_tracker is not None:
         input = server_conversation_tracker.prepare_input(original_input, generated_items)
     else:
-        input = ItemHelpers.input_to_new_input_list(original_input)
-        if isinstance(input, list):
-            append_input_items_excluding_approvals(
-                input,
-                generated_items,
-                reasoning_item_id_policy,
-            )
-        else:
-            input = ItemHelpers.input_to_new_input_list(input)
-            append_input_items_excluding_approvals(
-                input,
-                generated_items,
-                reasoning_item_id_policy,
-            )
-
-    if isinstance(input, list):
-        input = normalize_input_items_for_api(input)
+        input = _prepare_turn_input_items(original_input, generated_items, reasoning_item_id_policy)
 
     new_response = await get_new_response(
         agent,

src/agents/run_internal/session_persistence.py

@@ -91,13 +91,17 @@ async def prepare_input_with_session(
         ensure_input_item_format(item) for item in ItemHelpers.input_to_new_input_list(input)
     ]
 
+    prune_history_indexes: set[int] = set()
+
     if session_input_callback is None or not include_history_in_prepared_input:
         prepared_items_raw: list[TResponseInputItem] = (
             converted_history + new_input_list
             if include_history_in_prepared_input
             else list(new_input_list)
         )
         appended_items = list(new_input_list)
+        if include_history_in_prepared_input:
+            prune_history_indexes = set(range(len(converted_history)))
     else:
         if not callable(session_input_callback):
             raise UserError(
@@ -121,17 +125,19 @@ async def prepare_input_with_session(
         new_counts = _build_frequency_map(new_items_for_callback)
 
         appended: list[Any] = []
-        for item in combined:
+        for combined_index, item in enumerate(combined):
             key = _session_item_key(item)
             if _consume_reference(new_refs, key, item):
                 new_counts[key] = max(new_counts.get(key, 0) - 1, 0)
                 appended.append(item)
                 continue
             if _consume_reference(history_refs, key, item):
                 history_counts[key] = max(history_counts.get(key, 0) - 1, 0)
+                prune_history_indexes.add(combined_index)
                 continue
             if history_counts.get(key, 0) > 0:
                 history_counts[key] = history_counts.get(key, 0) - 1
+                prune_history_indexes.add(combined_index)
                 continue
             if new_counts.get(key, 0) > 0:
                 new_counts[key] = max(new_counts.get(key, 0) - 1, 0)
@@ -151,7 +157,10 @@ async def prepare_input_with_session(
     # Normalize exactly as the runtime does elsewhere so the prepared model input and the
     # persisted session items are derived from the same item shape and dedupe rules.
     prepared_as_inputs = [ensure_input_item_format(item) for item in prepared_items_raw]
-    filtered = drop_orphan_function_calls(prepared_as_inputs)
+    filtered = drop_orphan_function_calls(
+        prepared_as_inputs,
+        pruning_indexes=prune_history_indexes,
+    )
     normalized = normalize_input_items_for_api(filtered)
     deduplicated = deduplicate_input_items_preferring_latest(normalized)