Files
strix/tests/test_context_budget.py
T
Ahmed Allam a73300f840 feat(context): model-aware conversation compaction for long scans
Restore cumulative history compaction (removed in the SDK migration) so a
long scan no longer replays an ever-growing transcript until it overflows
the model's context window and the run fails.

- strix/llm/context_budget.py: resolve the model's real input/output token
  limits from LiteLLM metadata (128k gpt-4o, 272k gpt-5, 1M claude-sonnet-4,
  131k deepseek), with a large configurable fallback for unmapped models and
  a chars/4 token-count fallback.
- strix/llm/compaction.py: provider-agnostic compaction via litellm. Keeps a
  security-focused structured summary (objective, findings, credentials,
  payloads, URLs/paths, work state, dead ends, next move), keeps the most
  recent turns by token budget, and snaps the summary boundary so no tool
  call is separated from its result. Both triggers: proactive before each
  run and reactive compact-and-retry on a real context-overflow error.
- Wire both triggers into the agent run loop next to the existing image
  recovery; add replace_session_items() with restore-on-failure.

Env-tunable via STRIX_CONTEXT_* (auto-compact on by default).
2026-07-25 23:46:27 +00:00

47 lines
1.4 KiB
Python

"""Tests for model-aware token budgets."""
from __future__ import annotations
from typing import TYPE_CHECKING
from strix.config import load_settings
from strix.llm import context_budget
if TYPE_CHECKING:
import pytest
def test_context_window_known_model() -> None:
# gpt-4o is mapped by LiteLLM at 128k input tokens.
assert context_budget.context_window("gpt-4o") == 128_000
def test_context_window_strips_provider_prefix() -> None:
assert context_budget.context_window("openai/gpt-4o") == 128_000
def test_context_window_unmapped_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
context_budget._model_info.cache_clear()
def _raise(_model: str) -> dict[str, int]:
raise ValueError("This model isn't mapped yet.")
monkeypatch.setattr("strix.llm.context_budget.litellm.get_model_info", _raise)
expected = load_settings().context.fallback_context_tokens
assert context_budget.context_window("totally-made-up-model") == expected
context_budget._model_info.cache_clear()
def test_count_tokens_fallback_on_error(monkeypatch: pytest.MonkeyPatch) -> None:
def _raise(**_kwargs: object) -> int:
raise RuntimeError("no tokenizer")
monkeypatch.setattr("strix.llm.context_budget.litellm.token_counter", _raise)
text = "x" * 400
assert context_budget.count_tokens("weird-model", text) == 100
def test_count_tokens_empty_is_zero() -> None:
assert context_budget.count_tokens("gpt-4o", "") == 0