From fa52c344a62dec2582e924ed4e46ada1f76405aa Mon Sep 17 00:00:00 2001
From: oyasumi <oyasumi@kantilabs.xyz>
Date: Wed, 19 Aug 2026 21:32:46 +0000
Subject: [PATCH] fix(llm): preserve images during prompt retries

---
 strix/core/execution.py                 | 42 +++++++++++++------------
 tests/test_execution_transient_retry.py | 39 ++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 21 deletions(-)

diff --git a/strix/core/execution.py b/strix/core/execution.py
index 31551d29..bb2c10f4 100644
--- a/strix/core/execution.py
+++ b/strix/core/execution.py
@@ -729,8 +729,29 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
             await coordinator.trigger_budget_stop()
             raise
         except Exception as exc:
+            prompt_policy_rejection = _is_openrouter_prompt_policy_rejection(exc)
             if (
-                image_strips < 3
+                prompt_policy_rejection
+                and prompt_policy_retries < _MAX_OPENROUTER_PROMPT_POLICY_RETRIES
+            ):
+                prompt_policy_retries += 1
+                delay = _transient_model_retry_delay(prompt_policy_retries)
+                logger.warning(
+                    "intermittent OpenRouter prompt-policy rejection for %s; replaying "
+                    "unchanged turn (attempt %d/%d, backoff %.1fs): %r",
+                    agent_id,
+                    prompt_policy_retries,
+                    _MAX_OPENROUTER_PROMPT_POLICY_RETRIES,
+                    delay,
+                    exc,
+                )
+                await asyncio.sleep(delay)
+                if session is not None:
+                    input_data = []
+                continue
+            if (
+                not prompt_policy_rejection
+                and image_strips < 3
                 and session is not None
                 and getattr(exc, "status_code", None) in _INPUT_REJECTION_CODES
             ):
@@ -767,25 +788,6 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                     )
                     input_data = []
                     continue
-            if (
-                prompt_policy_retries < _MAX_OPENROUTER_PROMPT_POLICY_RETRIES
-                and _is_openrouter_prompt_policy_rejection(exc)
-            ):
-                prompt_policy_retries += 1
-                delay = _transient_model_retry_delay(prompt_policy_retries)
-                logger.warning(
-                    "intermittent OpenRouter prompt-policy rejection for %s; replaying "
-                    "unchanged turn (attempt %d/%d, backoff %.1fs): %r",
-                    agent_id,
-                    prompt_policy_retries,
-                    _MAX_OPENROUTER_PROMPT_POLICY_RETRIES,
-                    delay,
-                    exc,
-                )
-                await asyncio.sleep(delay)
-                if session is not None:
-                    input_data = []
-                continue
             if model_retries < _MAX_TRANSIENT_MODEL_RETRIES and _is_transient_model_error(exc):
                 model_retries += 1
                 delay = _transient_model_retry_delay(model_retries)
diff --git a/tests/test_execution_transient_retry.py b/tests/test_execution_transient_retry.py
index 44ccdd2c..b94e8742 100644
--- a/tests/test_execution_transient_retry.py
+++ b/tests/test_execution_transient_retry.py
@@ -124,6 +124,8 @@ def _patch_fast_backoff(monkeypatch: pytest.MonkeyPatch) -> None:
 async def _run_once(
     monkeypatch: pytest.MonkeyPatch,
     streams: list[_FakeStream],
+    *,
+    session: Any = None,
 ) -> Any:
     _patch_fast_backoff(monkeypatch)
     calls = {"n": 0}
@@ -146,7 +148,7 @@ async def _run_once(
         run_config=cast("RunConfig", object()),
         context={},
         max_turns=5,
-        session=None,
+        session=session,
         interactive=False,
         event_sink=None,
         hooks=None,
@@ -193,6 +195,41 @@ async def test_run_cycle_gives_up_after_openrouter_prompt_policy_retry_limit(
         await _run_once(monkeypatch, streams)
 
 
+@pytest.mark.asyncio
+async def test_openrouter_prompt_policy_retry_preserves_session_images(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class _Session:
+        async def get_items(self) -> list[Any]:
+            return []
+
+    image_strips = 0
+
+    async def _strip_images(_session: Any) -> bool:
+        nonlocal image_strips
+        image_strips += 1
+        return True
+
+    async def _no_compaction(*_args: Any, **_kwargs: Any) -> bool:
+        return False
+
+    async def _no_salvage(*_args: Any, **_kwargs: Any) -> None:
+        return None
+
+    monkeypatch.setattr(execution, "strip_all_images_from_session", _strip_images)
+    monkeypatch.setattr(execution, "_compact_session", _no_compaction)
+    monkeypatch.setattr(execution, "_salvage_stream_to_session", _no_salvage)
+    streams = [
+        _FakeStream(exc=_openrouter_prompt_policy_rejection())
+        for _ in range(execution._MAX_OPENROUTER_PROMPT_POLICY_RETRIES + 1)
+    ]
+
+    with pytest.raises(BadRequestError):
+        await _run_once(monkeypatch, streams, session=_Session())
+
+    assert image_strips == 0
+
+
 @pytest.mark.asyncio
 async def test_run_cycle_gives_up_after_max_retries(
     monkeypatch: pytest.MonkeyPatch,