From 94adabd67e02aa4a77e0afacba0cd019dd684760 Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 6 Oct 2026 15:23:39 -0400 Subject: [PATCH] fix(compaction): rerun without the budget check when compaction can't fit the session --- strix/core/execution.py | 49 ++++++++++++++++++++++++++++++++++------- tests/test_execution.py | 49 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 8 deletions(-) diff --git a/strix/core/execution.py b/strix/core/execution.py index eff7e7eb..27d1a071 100644 --- a/strix/core/execution.py +++ b/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, diff --git a/tests/test_execution.py b/tests/test_execution.py index cf29bd8b..0e4bbf5c 100644 --- a/tests/test_execution.py +++ b/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()