From 3a69a45f86620d4b386db0b6647ed7b338c0ff01 Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 6 Oct 2026 15:14:27 -0400 Subject: [PATCH] fix(compaction): stop the run and compact before a request would overrun the context window --- strix/core/execution.py | 51 +++++++++++++++++++++----- strix/llm/compaction.py | 43 +++++++++++++++------- tests/test_compaction.py | 19 +++++++++- tests/test_execution.py | 77 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 169 insertions(+), 21 deletions(-) diff --git a/strix/core/execution.py b/strix/core/execution.py index 91d94498..eff7e7eb 100644 --- a/strix/core/execution.py +++ b/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, diff --git a/strix/llm/compaction.py b/strix/llm/compaction.py index 47fc04df..cc76a421 100644 --- a/strix/llm/compaction.py +++ b/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) diff --git a/tests/test_compaction.py b/tests/test_compaction.py index 3b03e4b3..adddbf9c 100644 --- a/tests/test_compaction.py +++ b/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) diff --git a/tests/test_execution.py b/tests/test_execution.py index 83d3829b..cf29bd8b 100644 --- a/tests/test_execution.py +++ b/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()