From f7e3af49bd5cbd23353db70d645bbc6e858ded8d Mon Sep 17 00:00:00 2001
From: Ahmed Allam <49919286+0xallam@users.noreply.github.com>
Date: Tue, 9 Jun 2026 02:48:30 -0700
Subject: [PATCH] Strip all images from session on vision-rejection, not just
 the latest (#553)

---
 strix/core/execution.py          |  6 ++--
 strix/core/sessions.py           | 53 ++++++++++++++++++++------------
 strix/tools/view_image/README.md |  2 +-
 3 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/strix/core/execution.py b/strix/core/execution.py
index 1b561e9..0e9f914 100644
--- a/strix/core/execution.py
+++ b/strix/core/execution.py
@@ -16,7 +16,7 @@ from docker import errors as docker_errors  # type: ignore[import-untyped, unuse
 from openai import APIError
 
 from strix.core.inputs import child_initial_input
-from strix.core.sessions import open_agent_session, strip_latest_image_from_session
+from strix.core.sessions import open_agent_session, strip_all_images_from_session
 
 
 if TYPE_CHECKING:
@@ -384,14 +384,14 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                 and getattr(exc, "status_code", None) in _INPUT_REJECTION_CODES
             ):
                 try:
-                    stripped = await strip_latest_image_from_session(session)
+                    stripped = await strip_all_images_from_session(session)
                 except Exception:
                     logger.exception("image-strip recovery failed for %s", agent_id)
                     stripped = False
                 if stripped:
                     image_strips += 1
                     logger.info(
-                        "Stripped latest image from %s session after rejection; retrying (%d)",
+                        "Stripped images from %s session after rejection; retrying (%d)",
                         agent_id,
                         image_strips,
                     )
diff --git a/strix/core/sessions.py b/strix/core/sessions.py
index 28bdf71..8e4f414 100644
--- a/strix/core/sessions.py
+++ b/strix/core/sessions.py
@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, cast
+import contextlib
+from typing import TYPE_CHECKING, Any, cast
 
 from agents.memory import SQLiteSession
 
@@ -22,29 +23,43 @@ def open_agent_session(agent_id: str, path: Path) -> SQLiteSession:
 _IMAGE_REJECTED_TEXT = "[image rejected by the model]"
 
 
-async def strip_latest_image_from_session(session: Session) -> bool:
+async def strip_all_images_from_session(session: Session) -> bool:
     items = await session.get_items()
     if not items:
         return False
-    latest = items[-1]
-    if not isinstance(latest, dict) or latest.get("type") != "function_call_output":
-        return False
-    output = latest.get("output")
-    if not isinstance(output, list):
-        return False
-    if not any(isinstance(b, dict) and b.get("type") == "input_image" for b in output):
-        return False
-    await session.pop_item()
-    await session.add_items(
-        cast(
-            "list[TResponseInputItem]",
-            [
+
+    rebuilt: list[Any] = []
+    changed = False
+    for item in items:
+        item_dict = cast("dict[str, Any]", item) if isinstance(item, dict) else None
+        if (
+            item_dict is not None
+            and item_dict.get("type") == "function_call_output"
+            and isinstance(item_dict.get("output"), list)
+            and any(
+                isinstance(b, dict) and b.get("type") == "input_image" for b in item_dict["output"]
+            )
+        ):
+            rebuilt.append(
                 {
                     "type": "function_call_output",
-                    "call_id": latest.get("call_id"),
+                    "call_id": item_dict.get("call_id"),
                     "output": [{"type": "input_text", "text": _IMAGE_REJECTED_TEXT}],
                 },
-            ],
-        ),
-    )
+            )
+            changed = True
+        else:
+            rebuilt.append(item)
+
+    if not changed:
+        return False
+
+    rebuilt_items = cast("list[TResponseInputItem]", rebuilt)
+    await session.clear_session()
+    try:
+        await session.add_items(rebuilt_items)
+    except Exception:
+        with contextlib.suppress(Exception):
+            await session.add_items(rebuilt_items)
+        raise
     return True
diff --git a/strix/tools/view_image/README.md b/strix/tools/view_image/README.md
index 3827206..cd4c622 100644
--- a/strix/tools/view_image/README.md
+++ b/strix/tools/view_image/README.md
@@ -13,5 +13,5 @@ returns it as an image content block for vision-capable models.
   `containers/docker-entrypoint.sh`).
 - **Skill:** screenshot workflow lives in `strix/skills/tooling/agent_browser.md`.
 - **Recovery:** vision-not-supported model rejections are auto-recovered
-  via `strix.core.sessions.strip_latest_image_from_session`, invoked from
+  via `strix.core.sessions.strip_all_images_from_session`, invoked from
   `strix/core/execution.py`.