refactor

Author: seratch

src/agents/run_internal/items.py

@@ -46,19 +46,6 @@ def drop_orphan_function_calls(items: list[TResponseInputItem]) -> list[TRespons
     not replay stale tool calls.
     """
 
-    def _completed_call_ids(payload: list[TResponseInputItem]) -> set[str]:
-        completed: set[str] = set()
-        for entry in payload:
-            if not isinstance(entry, dict):
-                continue
-            item_type = entry.get("type")
-            if item_type not in ("function_call_output", "function_call_result"):
-                continue
-            call_id = entry.get("call_id") or entry.get("callId")
-            if call_id and isinstance(call_id, str):
-                completed.add(call_id)
-        return completed
-
     completed_call_ids = _completed_call_ids(items)
 
     filtered: list[TResponseInputItem] = []
@@ -77,19 +64,7 @@ def _completed_call_ids(payload: list[TResponseInputItem]) -> set[str]:
 
 def ensure_input_item_format(item: TResponseInputItem) -> TResponseInputItem:
     """Ensure a single item is normalized for model input (function_call_output, snake_case)."""
-
-    def _coerce_dict(value: TResponseInputItem) -> dict[str, Any] | None:
-        """Convert dataclass/Pydantic items into plain dicts when possible."""
-        if isinstance(value, dict):
-            return dict(value)
-        if hasattr(value, "model_dump"):
-            try:
-                return cast(dict[str, Any], value.model_dump(exclude_unset=True))
-            except Exception:
-                return None
-        return None
-
-    coerced = _coerce_dict(item)
+    coerced = _coerce_to_dict(item)
     if coerced is None:
         return item
 
@@ -100,17 +75,6 @@ def _coerce_dict(value: TResponseInputItem) -> dict[str, Any] | None:
 def normalize_input_items_for_api(items: list[TResponseInputItem]) -> list[TResponseInputItem]:
     """Normalize input items for API submission and strip provider data for downstream services."""
 
-    def _coerce_to_dict(value: TResponseInputItem) -> dict[str, Any] | None:
-        """Convert model items to dicts so fields can be renamed and sanitized."""
-        if isinstance(value, dict):
-            return dict(value)
-        if hasattr(value, "model_dump"):
-            try:
-                return cast(dict[str, Any], value.model_dump(exclude_unset=True))
-            except Exception:
-                return None
-        return None
-
     normalized: list[TResponseInputItem] = []
     for item in items:
         coerced = _coerce_to_dict(item)
@@ -223,17 +187,33 @@ def extract_mcp_request_id_from_run(mcp_run: Any) -> str | None:
     return candidate if isinstance(candidate, str) else None
 
 
-__all__ = [
-    "REJECTION_MESSAGE",
-    "copy_input_items",
-    "drop_orphan_function_calls",
-    "ensure_input_item_format",
-    "normalize_input_items_for_api",
-    "fingerprint_input_item",
-    "deduplicate_input_items",
-    "function_rejection_item",
-    "shell_rejection_item",
-    "apply_patch_rejection_item",
-    "extract_mcp_request_id",
-    "extract_mcp_request_id_from_run",
-]
+# --------------------------
+# Private helpers
+# --------------------------
+
+
+def _completed_call_ids(payload: list[TResponseInputItem]) -> set[str]:
+    """Return the call ids that already have outputs."""
+    completed: set[str] = set()
+    for entry in payload:
+        if not isinstance(entry, dict):
+            continue
+        item_type = entry.get("type")
+        if item_type not in ("function_call_output", "function_call_result"):
+            continue
+        call_id = entry.get("call_id") or entry.get("callId")
+        if call_id and isinstance(call_id, str):
+            completed.add(call_id)
+    return completed
+
+
+def _coerce_to_dict(value: TResponseInputItem) -> dict[str, Any] | None:
+    """Convert model items to dicts so fields can be renamed and sanitized."""
+    if isinstance(value, dict):
+        return dict(value)
+    if hasattr(value, "model_dump"):
+        try:
+            return cast(dict[str, Any], value.model_dump(exclude_unset=True))
+        except Exception:
+            return None
+    return None

src/agents/run_internal/session_persistence.py

@@ -85,57 +85,19 @@ async def prepare_input_with_session(
         if not isinstance(combined, list):
             raise UserError("Session input callback must return a list of input items.")
 
-        def session_item_key(item: Any) -> str:
-            try:
-                if hasattr(item, "model_dump"):
-                    payload = item.model_dump(exclude_unset=True)
-                elif isinstance(item, dict):
-                    payload = item
-                else:
-                    payload = ensure_input_item_format(item)
-                return json.dumps(payload, sort_keys=True, default=str)
-            except Exception:
-                return repr(item)
-
-        def build_reference_map(items: Sequence[Any]) -> dict[str, list[Any]]:
-            refs: dict[str, list[Any]] = {}
-            for item in items:
-                key = session_item_key(item)
-                refs.setdefault(key, []).append(item)
-            return refs
-
-        def consume_reference(ref_map: dict[str, list[Any]], key: str, candidate: Any) -> bool:
-            candidates = ref_map.get(key)
-            if not candidates:
-                return False
-            for idx, existing in enumerate(candidates):
-                if existing is candidate:
-                    candidates.pop(idx)
-                    if not candidates:
-                        ref_map.pop(key, None)
-                    return True
-            return False
-
-        def build_frequency_map(items: Sequence[Any]) -> dict[str, int]:
-            freq: dict[str, int] = {}
-            for item in items:
-                key = session_item_key(item)
-                freq[key] = freq.get(key, 0) + 1
-            return freq
-
-        history_refs = build_reference_map(history_for_callback)
-        new_refs = build_reference_map(new_items_for_callback)
-        history_counts = build_frequency_map(history_for_callback)
-        new_counts = build_frequency_map(new_items_for_callback)
+        history_refs = _build_reference_map(history_for_callback)
+        new_refs = _build_reference_map(new_items_for_callback)
+        history_counts = _build_frequency_map(history_for_callback)
+        new_counts = _build_frequency_map(new_items_for_callback)
 
         appended: list[Any] = []
         for item in combined:
-            key = session_item_key(item)
-            if consume_reference(new_refs, key, item):
+            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):
+            if _consume_reference(history_refs, key, item):
                 history_counts[key] = max(history_counts.get(key, 0) - 1, 0)
                 continue
             if history_counts.get(key, 0) > 0:
@@ -440,3 +402,54 @@ async def wait_for_session_cleanup(
     logger.debug(
         "Session cleanup verification exhausted attempts; targets may still linger temporarily"
     )
+
+
+# --------------------------
+# Private helpers
+# --------------------------
+
+
+def _session_item_key(item: Any) -> str:
+    """Return a stable representation of a session item for comparison."""
+    try:
+        if hasattr(item, "model_dump"):
+            payload = item.model_dump(exclude_unset=True)
+        elif isinstance(item, dict):
+            payload = item
+        else:
+            payload = ensure_input_item_format(item)
+        return json.dumps(payload, sort_keys=True, default=str)
+    except Exception:
+        return repr(item)
+
+
+def _build_reference_map(items: Sequence[Any]) -> dict[str, list[Any]]:
+    """Map serialized keys to the concrete session items used to build them."""
+    refs: dict[str, list[Any]] = {}
+    for item in items:
+        key = _session_item_key(item)
+        refs.setdefault(key, []).append(item)
+    return refs
+
+
+def _consume_reference(ref_map: dict[str, list[Any]], key: str, candidate: Any) -> bool:
+    """Remove a specific candidate from a reference map when it is consumed."""
+    candidates = ref_map.get(key)
+    if not candidates:
+        return False
+    for idx, existing in enumerate(candidates):
+        if existing is candidate:
+            candidates.pop(idx)
+            if not candidates:
+                ref_map.pop(key, None)
+            return True
+    return False
+
+
+def _build_frequency_map(items: Sequence[Any]) -> dict[str, int]:
+    """Count how many times each serialized key appears in a collection."""
+    freq: dict[str, int] = {}
+    for item in items:
+        key = _session_item_key(item)
+        freq[key] = freq.get(key, 0) + 1
+    return freq

src/agents/run_internal/tool_actions.py

@@ -90,11 +90,6 @@ async def execute(
         """Run a computer action, capturing a screenshot and notifying hooks."""
         computer = await resolve_computer(tool=action.computer_tool, run_context=context_wrapper)
         agent_hooks = agent.hooks
-        output_func = (
-            cls._get_screenshot_async(computer, action.tool_call)
-            if hasattr(computer, "screenshot_async")
-            else cls._get_screenshot_sync(computer, action.tool_call)
-        )
         await asyncio.gather(
             hooks.on_tool_start(context_wrapper, agent, action.computer_tool),
             (
@@ -104,7 +99,7 @@ async def execute(
             ),
         )
 
-        output = await output_func
+        output = await cls._execute_action_and_capture(computer, action.tool_call)
 
         await asyncio.gather(
             hooks.on_tool_end(context_wrapper, agent, action.computer_tool, output),
@@ -131,62 +126,40 @@ async def execute(
         )
 
     @classmethod
-    async def _get_screenshot_sync(
-        cls,
-        computer: Any,
-        tool_call: ResponseComputerToolCall,
+    async def _execute_action_and_capture(
+        cls, computer: Any, tool_call: ResponseComputerToolCall
     ) -> str:
-        """Execute the computer action for sync drivers and return the screenshot."""
-        action = tool_call.action
-        if isinstance(action, ActionClick):
-            computer.click(action.x, action.y, action.button)
-        elif isinstance(action, ActionDoubleClick):
-            computer.double_click(action.x, action.y)
-        elif isinstance(action, ActionDrag):
-            computer.drag([(p.x, p.y) for p in action.path])
-        elif isinstance(action, ActionKeypress):
-            computer.keypress(action.keys)
-        elif isinstance(action, ActionMove):
-            computer.move(action.x, action.y)
-        elif isinstance(action, ActionScreenshot):
-            computer.screenshot()
-        elif isinstance(action, ActionScroll):
-            computer.scroll(action.x, action.y, action.scroll_x, action.scroll_y)
-        elif isinstance(action, ActionType):
-            computer.type(action.text)
-        elif isinstance(action, ActionWait):
-            computer.wait()
+        """Execute the computer action (sync or async drivers) and return the screenshot."""
 
-        return cast(str, computer.screenshot())
+        async def maybe_call(method_name: str, *args: Any) -> Any:
+            method = getattr(computer, method_name, None)
+            if method is None or not callable(method):
+                raise ModelBehaviorError(f"Computer driver missing method {method_name}")
+            result = method(*args)
+            return await result if inspect.isawaitable(result) else result
 
-    @classmethod
-    async def _get_screenshot_async(
-        cls,
-        computer: Any,
-        tool_call: ResponseComputerToolCall,
-    ) -> str:
-        """Execute the computer action for async drivers and return the screenshot."""
         action = tool_call.action
         if isinstance(action, ActionClick):
-            await computer.click(action.x, action.y, action.button)
+            await maybe_call("click", action.x, action.y, action.button)
         elif isinstance(action, ActionDoubleClick):
-            await computer.double_click(action.x, action.y)
+            await maybe_call("double_click", action.x, action.y)
         elif isinstance(action, ActionDrag):
-            await computer.drag([(p.x, p.y) for p in action.path])
+            await maybe_call("drag", [(p.x, p.y) for p in action.path])
         elif isinstance(action, ActionKeypress):
-            await computer.keypress(action.keys)
+            await maybe_call("keypress", action.keys)
         elif isinstance(action, ActionMove):
-            await computer.move(action.x, action.y)
+            await maybe_call("move", action.x, action.y)
         elif isinstance(action, ActionScreenshot):
-            await computer.screenshot()
+            await maybe_call("screenshot")
         elif isinstance(action, ActionScroll):
-            await computer.scroll(action.x, action.y, action.scroll_x, action.scroll_y)
+            await maybe_call("scroll", action.x, action.y, action.scroll_x, action.scroll_y)
         elif isinstance(action, ActionType):
-            await computer.type(action.text)
+            await maybe_call("type", action.text)
         elif isinstance(action, ActionWait):
-            await computer.wait()
+            await maybe_call("wait")
 
-        return cast(str, await computer.screenshot())
+        screenshot_result = await maybe_call("screenshot")
+        return cast(str, screenshot_result)
 
 
 class LocalShellAction: