fix(compaction): stop the run and compact before a request would overrun the context window

strix/core/execution.py

@@ -7,6 +7,7 @@ import contextlib
 import logging
 import uuid
 from collections.abc import Callable
+from dataclasses import replace
 from functools import cache
 from typing import TYPE_CHECKING, Any, cast
 
@@ -34,7 +35,12 @@ from strix.core.sessions import (
     strip_all_images_from_session,
 )
 from strix.llm import request_log
-from strix.llm.compaction import is_context_overflow, maybe_compact
+from strix.llm.compaction import (
+    CompactionNeededError,
+    is_context_overflow,
+    maybe_compact,
+    needs_compaction,
+)
 
 
 if TYPE_CHECKING:
@@ -44,6 +50,7 @@ if TYPE_CHECKING:
     from agents.lifecycle import RunHooks
     from agents.memory import Session, SQLiteSession
     from agents.result import RunResultBase
+    from agents.run_config import CallModelData, ModelInputData
 
     from strix.core.agents import AgentCoordinator, Status
 
@@ -120,6 +127,33 @@ async def _compact_session(
     )
 
 
+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
+    history, so compacting the session mid-run would not shrink it."""
+    model = _run_config_model(run_config)
+    if model is None:
+        return run_config
+    first_call = True
+
+    async def _check(data: CallModelData[Any]) -> ModelInputData:
+        nonlocal first_call
+        # Skip the first call: the cycle compacts right before the run, and the run's
+        # new input is only persisted after this filter, so stopping here would lose it.
+        if not first_call and await asyncio.to_thread(
+            needs_compaction,
+            model,
+            data.model_data.instructions or "",
+            _agent_tools_text(data.agent),
+            data.model_data.input,
+        ):
+            raise CompactionNeededError(model)
+        first_call = False
+        return data.model_data
+
+    return replace(run_config, call_model_input_filter=_check)
+
+
 _MAX_TRANSIENT_MODEL_RETRIES = 5
 _TRANSIENT_MODEL_RETRY_BASE_DELAY_S = 2.0
 _TRANSIENT_MODEL_RETRY_MAX_DELAY_S = 90.0
@@ -747,7 +781,7 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
             stream = Runner.run_streamed(
                 agent,
                 input=input_data,
-                run_config=run_config,
+                run_config=run_config if session is None else _with_compaction_check(run_config),
                 context=context,
                 max_turns=max_turns,
                 session=session,
@@ -824,18 +858,19 @@ async def _run_cycle(  # noqa: PLR0912, PLR0915
                     )
                     input_data = []
                     continue
-            if (
-                compactions < _MAX_COMPACTIONS_PER_CYCLE
-                and session is not None
-                and is_context_overflow(exc)
-            ):
+            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):
                 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:
-                    compactions += 1
+                    # Each check_context restart makes at least one model call, so only
+                    # provider overflows need a cap.
+                    if not budget_hit:
+                        compactions += 1
                     logger.info(
                         "Compacted %s session after context overflow; retrying (%d)",
                         agent_id,

strix/llm/compaction.py

@@ -164,22 +164,22 @@ def _content_text(content: Any) -> str:
     return ""
 
 
-def _truncate(text: str, limit: int) -> str:
-    return text if len(text) <= limit else f"{text[:limit]}\n[truncated]"
+def _truncate(text: str, limit: int | None) -> str:
+    return text if limit is None or len(text) <= limit else f"{text[:limit]}\n[truncated]"
 
 
-def _serialize_item(item: Any) -> str:
+def _serialize_item(item: Any, max_chars: int | None = _TOOL_OUTPUT_MAX_CHARS) -> str:
     if not isinstance(item, dict):
         return str(item)
     item_type = item.get("type")
     role = item.get("role")
     if item_type == "function_call":
-        args = _truncate(str(item.get("arguments", "")), _TOOL_OUTPUT_MAX_CHARS)
+        args = _truncate(str(item.get("arguments", "")), max_chars)
         return f"[tool_call {item.get('name', '?')}] {args}"
     if item_type == "function_call_output":
         output = item.get("output")
         text = output if isinstance(output, str) else _content_text(output)
-        return f"[tool_result] {_truncate(text, _TOOL_OUTPUT_MAX_CHARS)}"
+        return f"[tool_result] {_truncate(text, max_chars)}"
     if item_type == "reasoning":
         return ""
     if role or item_type == "message":
@@ -187,8 +187,30 @@ def _serialize_item(item: Any) -> str:
     return ""
 
 
-def _serialize_items(items: list[Any]) -> str:
-    return "\n".join(s for s in (_serialize_item(item) for item in items) if s)
+def _serialize_items(items: list[Any], max_chars: int | None = _TOOL_OUTPUT_MAX_CHARS) -> str:
+    return "\n".join(s for s in (_serialize_item(item, max_chars) for item in items) if s)
+
+
+class CompactionNeededError(Exception):
+    """The next model request would overrun the model's usable context window."""
+
+
+def _context_budget(model: str) -> int:
+    context = load_settings().context
+    reserve = max(context.compact_buffer_tokens, output_limit(model))
+    return max(context.keep_tokens, context_window(model) - reserve)
+
+
+def _request_tokens(model: str, instructions: str, tools_text: str, items: list[Any]) -> int:
+    text = "\n".join((instructions, tools_text, _serialize_items(items, max_chars=None)))
+    return count_tokens(model, text)
+
+
+def needs_compaction(model: str, instructions: str, tools_text: str, items: list[Any]) -> bool:
+    """Whether a request carrying ``items`` would overrun ``model``'s context budget."""
+    if not load_settings().context.auto_compact:
+        return False
+    return _request_tokens(model, instructions, tools_text, items) > _context_budget(model)
 
 
 def _is_tool_call(item: Any) -> bool:
@@ -369,11 +391,8 @@ async def maybe_compact(
     if len(items) < _MIN_ITEMS_TO_COMPACT:
         return False
 
-    window = context_window(model)
-    reserve = max(context.compact_buffer_tokens, output_limit(model))
-    budget = max(context.keep_tokens, window - reserve)
-    used = count_tokens(model, "\n".join((instructions, tools_text, _serialize_items(items))))
-    if not force and used <= budget:
+    used = _request_tokens(model, instructions, tools_text, items)
+    if not force and used <= _context_budget(model):
         return False
 
     split = _select_split(model, items, context.keep_tokens)

tests/test_compaction.py

@@ -138,7 +138,9 @@ def test_select_split_handles_parallel_calls(monkeypatch: pytest.MonkeyPatch) ->
     assert not _has_orphan_tool_output(items[split:])
 
 
-def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window: int) -> None:
+def _patch_budget(
+    monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window: int
+) -> ContextSettings:
     monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t))
     monkeypatch.setattr(compaction, "context_window", lambda _m: window)
     monkeypatch.setattr(compaction, "output_limit", lambda _m: 0)
@@ -152,6 +154,7 @@ def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window:
         llm=SimpleNamespace(api_key=None, api_base=None, timeout=1, extra_headers=None),
     )
     monkeypatch.setattr(compaction, "load_settings", lambda: settings)
+    return context
 
 
 def _model_response(text: str) -> Any:
@@ -180,6 +183,20 @@ def _patch_summary(
     monkeypatch.setattr(compaction, "StrixProvider", FakeProvider)
 
 
+def test_needs_compaction_counts_full_tool_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
+    _patch_budget(monkeypatch, keep_tokens=50, window=10_000)
+    items = [_call("c1"), _output("c1", "x" * 20_000)]
+
+    assert compaction.needs_compaction("m", "", "", items) is True
+    assert compaction.needs_compaction("m", "", "", _turns(2)) is False
+
+
+def test_needs_compaction_respects_auto_compact_off(monkeypatch: pytest.MonkeyPatch) -> None:
+    _patch_budget(monkeypatch, keep_tokens=50, window=10_000).auto_compact = False
+
+    assert compaction.needs_compaction("m", "", "", [_output("c1", "x" * 20_000)]) is False
+
+
 @pytest.mark.asyncio
 async def test_maybe_compact_noop_when_within_budget(monkeypatch: pytest.MonkeyPatch) -> None:
     _patch_budget(monkeypatch, keep_tokens=50, window=1_000_000)

tests/test_execution.py

@@ -9,9 +9,11 @@ from typing import Any, cast
 from unittest.mock import MagicMock
 
 import pytest
+from agents import RunConfig
 from agents.exceptions import MaxTurnsExceeded
 from agents.items import MessageOutputItem
 from agents.memory import SQLiteSession
+from agents.run_config import CallModelData, ModelInputData
 from agents.tool_context import ToolContext
 from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal, ResponseOutputText
 
@@ -22,6 +24,7 @@ from strix.core.execution import (
     notify_parent_on_terminal,
 )
 from strix.core.sessions import seed_initial_input
+from strix.llm.compaction import CompactionNeededError
 from strix.tools.agents_graph.tools import agent_finish, stop_agent
 from strix.tools.finish.tool import finish_scan
 
@@ -1531,3 +1534,77 @@ async def test_autonomous_nudge_does_not_offer_the_user() -> None:
     )
 
     assert "wait_for_user" not in items[0]["content"]
+
+
+@pytest.mark.asyncio
+async def test_compaction_check_stops_run_after_first_call_when_over_budget(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(execution, "needs_compaction", lambda *_args: True)
+    guard: Any = execution._with_compaction_check(RunConfig(model="m")).call_model_input_filter
+    data = CallModelData(
+        model_data=ModelInputData(input=[{"role": "user", "content": "hi"}], instructions="sys"),
+        agent=MagicMock(tools=[]),
+        context=None,
+    )
+
+    assert await guard(data) is data.model_data
+    with pytest.raises(CompactionNeededError):
+        await guard(data)
+
+
+class _BudgetStream:
+    def __init__(self, exc: BaseException | None) -> None:
+        self.run_loop_exception = exc
+        self.new_items: list[Any] = []
+
+    async def stream_events(self) -> Any:
+        for event in _NO_STREAM_EVENTS:
+            yield event
+
+
+@pytest.mark.asyncio
+async def test_run_cycle_compacts_and_resumes_on_every_budget_stop(
+    tmp_path: Any, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    stops = execution._MAX_COMPACTIONS_PER_CYCLE + 1
+    streams = [_BudgetStream(CompactionNeededError("m")) for _ in range(stops)]
+    streams.append(_BudgetStream(None))
+    runs: list[dict[str, Any]] = []
+
+    def _run_streamed(*_args: Any, **kwargs: Any) -> _BudgetStream:
+        runs.append(kwargs)
+        return streams[len(runs) - 1]
+
+    forced: list[bool] = []
+
+    async def _compact(*_args: Any, force: bool) -> bool:
+        forced.append(force)
+        return True
+
+    monkeypatch.setattr("strix.core.execution.Runner.run_streamed", _run_streamed)
+    monkeypatch.setattr(execution, "_compact_session", _compact)
+    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) == stops + 1
+    assert forced.count(True) == stops
+    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()