fix(compaction): rerun without the budget check when compaction can't fit the session

strix/core/execution.py

@@ -127,6 +127,19 @@ async def _compact_session(
     )
 
 
+async def _session_needs_compaction(agent: Any, session: Session, run_config: RunConfig) -> bool:
+    model = _run_config_model(run_config)
+    if model is None:
+        return False
+    return await asyncio.to_thread(
+        needs_compaction,
+        model,
+        _agent_instructions(agent),
+        _agent_tools_text(agent),
+        list(await session.get_items()),
+    )
+
+
 def _with_compaction_check(run_config: RunConfig) -> RunConfig:
     """Stop the run before a model request that would overrun the context window so
     the cycle can compact and resume: the SDK builds each request from its in-memory
@@ -759,6 +772,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
     image_strips = 0
     compactions = 0
     model_retries = 0
+    check_context = session is not None
     request_log.set_retry_attempt(0)
     while True:
         stream: Any = None
@@ -781,7 +795,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
             stream = Runner.run_streamed(
                 agent,
                 input=input_data,
-                run_config=run_config if session is None else _with_compaction_check(run_config),
+                run_config=_with_compaction_check(run_config) if check_context else run_config,
                 context=context,
                 max_turns=max_turns,
                 session=session,
@@ -858,19 +872,38 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                     )
                     input_data = []
                     continue
-            budget_hit = isinstance(exc, CompactionNeededError)
-            overflowed = compactions < _MAX_COMPACTIONS_PER_CYCLE and is_context_overflow(exc)
-            if session is not None and (budget_hit or overflowed):
+            if session is not None and isinstance(exc, CompactionNeededError):
+                # Rerun without the check if compaction can't bring the session under
+                # the budget; a real overflow then falls to the capped recovery below.
+                try:
+                    compacted = await _compact_session(agent, session, run_config, force=True)
+                    check_context = compacted and not await _session_needs_compaction(
+                        agent, session, run_config
+                    )
+                except Exception:
+                    logger.exception("context budget compaction failed for %s", agent_id)
+                    check_context = False
+                logger.info(
+                    "%s hit the context budget; %s",
+                    agent_id,
+                    "compacted and resuming"
+                    if check_context
+                    else "rerunning without the budget check",
+                )
+                input_data = []
+                continue
+            if (
+                compactions < _MAX_COMPACTIONS_PER_CYCLE
+                and session is not None
+                and is_context_overflow(exc)
+            ):
                 try:
                     compacted = await _compact_session(agent, session, run_config, force=True)
                 except Exception:
                     logger.exception("overflow compaction recovery failed for %s", agent_id)
                     compacted = False
                 if compacted:
-                    # Each check_context restart makes at least one model call, so only
-                    # provider overflows need a cap.
-                    if not budget_hit:
-                        compactions += 1
+                    compactions += 1
                     logger.info(
                         "Compacted %s session after context overflow; retrying (%d)",
                         agent_id,

tests/test_execution.py

@@ -1584,6 +1584,7 @@ async def test_run_cycle_compacts_and_resumes_on_every_budget_stop(
 
     monkeypatch.setattr("strix.core.execution.Runner.run_streamed", _run_streamed)
     monkeypatch.setattr(execution, "_compact_session", _compact)
+    monkeypatch.setattr(execution, "needs_compaction", lambda *_args: False)
     coordinator = AgentCoordinator()
     await coordinator.register("root", "strix", parent_id=None)
     session = SQLiteSession("root", tmp_path / "agents.db")
@@ -1608,3 +1609,51 @@ async def test_run_cycle_compacts_and_resumes_on_every_budget_stop(
     assert all(run["input"] == [] for run in runs[1:])
     assert all(run["run_config"].call_model_input_filter is not None for run in runs)
     session.close()
+
+
+@pytest.mark.parametrize(
+    ("compacted", "still_over"),
+    [(False, False), (True, True)],
+    ids=["compaction_failed", "still_over_budget"],
+)
+@pytest.mark.asyncio
+async def test_run_cycle_reruns_unchecked_when_compaction_cannot_fit(
+    tmp_path: Any, monkeypatch: pytest.MonkeyPatch, compacted: bool, still_over: bool
+) -> None:
+    streams = [_BudgetStream(CompactionNeededError("m")), _BudgetStream(None)]
+    runs: list[dict[str, Any]] = []
+
+    def _run_streamed(*_args: Any, **kwargs: Any) -> _BudgetStream:
+        runs.append(kwargs)
+        return streams[len(runs) - 1]
+
+    async def _compact(*_args: Any, **_kwargs: Any) -> bool:
+        return compacted
+
+    monkeypatch.setattr("strix.core.execution.Runner.run_streamed", _run_streamed)
+    monkeypatch.setattr(execution, "_compact_session", _compact)
+    monkeypatch.setattr(execution, "needs_compaction", lambda *_args: still_over)
+    coordinator = AgentCoordinator()
+    await coordinator.register("root", "strix", parent_id=None)
+    session = SQLiteSession("root", tmp_path / "agents.db")
+
+    result = await execution._run_cycle(
+        MagicMock(),
+        coordinator,
+        "root",
+        input_data="task",
+        run_config=RunConfig(model="m"),
+        context={},
+        max_turns=5,
+        session=session,
+        interactive=True,
+        event_sink=None,
+        hooks=None,
+    )
+
+    assert cast("Any", result) is streams[-1]
+    assert len(runs) == 2
+    assert runs[0]["run_config"].call_model_input_filter is not None
+    assert runs[1]["run_config"].call_model_input_filter is None
+    assert runs[1]["input"] == []
+    session.close()