diff --git a/strix/core/execution.py b/strix/core/execution.py index f7cf3759..a0280937 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -22,6 +22,7 @@ from strix.core.sessions import ( open_agent_session, strip_all_images_from_session, ) +from strix.llm.compaction import is_context_overflow, maybe_compact if TYPE_CHECKING: @@ -40,6 +41,41 @@ logger = logging.getLogger(__name__) StreamEventSink = Callable[[str, Any], None] _INPUT_REJECTION_CODES = frozenset({400, 404, 422}) +_MAX_COMPACTIONS_PER_CYCLE = 2 + + +def _run_config_model(run_config: RunConfig) -> str | None: + return run_config.model if isinstance(run_config.model, str) else None + + +def _agent_instructions(agent: Any) -> str: + instructions = getattr(agent, "instructions", None) + return instructions if isinstance(instructions, str) else "" + + +def _agent_tools_text(agent: Any) -> str: + parts: list[str] = [] + for tool in getattr(agent, "tools", []) or []: + name = getattr(tool, "name", "") + description = getattr(tool, "description", "") or "" + schema = getattr(tool, "params_json_schema", "") or "" + parts.append(f"{name} {description} {schema}") + return "\n".join(parts) + + +async def _compact_session( + agent: Any, session: Session, run_config: RunConfig, *, force: bool +) -> bool: + model = _run_config_model(run_config) + if session is None or model is None: + return False + return await maybe_compact( + session, + model=model, + instructions=_agent_instructions(agent), + tools_text=_agent_tools_text(agent), + force=force, + ) async def run_agent_loop( @@ -350,6 +386,7 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 hooks: RunHooks[dict[str, Any]] | None, ) -> RunResultBase | None: image_strips = 0 + compactions = 0 while True: try: await coordinator.mark_running(agent_id) @@ -360,6 +397,10 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 await enforce_image_budget(session, max_images) except Exception: logger.exception("image-budget enforcement failed for %s", agent_id) + try: + await _compact_session(agent, session, run_config, force=False) + except Exception: + logger.exception("proactive compaction failed for %s", agent_id) stream = Runner.run_streamed( agent, input=input_data, @@ -428,6 +469,25 @@ 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) + ): + 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 + logger.info( + "Compacted %s session after context overflow; retrying (%d)", + agent_id, + compactions, + ) + input_data = [] + continue if not interactive: raise if isinstance(exc, MaxTurnsExceeded): diff --git a/strix/core/sessions.py b/strix/core/sessions.py index 67fbba95..2bb83982 100644 --- a/strix/core/sessions.py +++ b/strix/core/sessions.py @@ -92,6 +92,39 @@ async def _rewrite_session( return True +async def replace_session_items( + session: Session, + new_items: list[Any], + *, + expected_len: int | None = None, +) -> bool: + """Overwrite the session's items, restoring the originals on failure. + + When ``expected_len`` is given, the rewrite is skipped if the session no + longer has that many items (a concurrent writer changed it), so a slow + compaction summary can't clobber newer turns. + """ + async with session_write_lock(session): + original = list(await session.get_items()) + if expected_len is not None and len(original) != expected_len: + logger.warning( + "skipping session rewrite: expected %d items, found %d", + expected_len, + len(original), + ) + return False + rebuilt = cast("list[TResponseInputItem]", new_items) + await session.clear_session() + try: + await session.add_items(rebuilt) + except Exception: + logger.exception("session rewrite failed; restoring original items") + await session.clear_session() + await session.add_items(original) + raise + return True + + async def strip_all_images_from_session(session: Session) -> bool: """Replace every image tool output with a text placeholder (rejection recovery).""" diff --git a/strix/llm/__init__.py b/strix/llm/__init__.py new file mode 100644 index 00000000..0d3f9546 --- /dev/null +++ b/strix/llm/__init__.py @@ -0,0 +1 @@ +"""LLM-facing context management: model-aware budgets and history compaction.""" diff --git a/strix/llm/compaction.py b/strix/llm/compaction.py new file mode 100644 index 00000000..da612ff2 --- /dev/null +++ b/strix/llm/compaction.py @@ -0,0 +1,354 @@ +"""Provider-agnostic conversation compaction. + +When an agent's session grows past the model's usable context window, older +turns are summarised into a single checkpoint while the most recent turns are +kept verbatim. This runs for every LiteLLM provider (not just OpenAI), keeps a +security-focused structured summary, and preserves tool-call/tool-result +pairing so the trimmed history is still valid provider input. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import litellm +from litellm.exceptions import BadRequestError, ContextWindowExceededError + +from strix.config import load_settings +from strix.core.sessions import replace_session_items, session_write_lock +from strix.llm.context_budget import context_window, count_tokens, output_limit + + +if TYPE_CHECKING: + from agents.memory import Session + + +logger = logging.getLogger(__name__) + +_CHECKPOINT_TAG = "" +_TOOL_OUTPUT_MAX_CHARS = 2_000 +_MIN_ITEMS_TO_COMPACT = 6 +_HEAD_TRUNCATED_MARKER = "\n\n[... older conversation omitted to fit the summary request ...]\n\n" + + +# Providers that don't type overflow errors (OpenRouter maps every 400 to a +# plain BadRequestError) leave only the message to go on, so we match it the way +# LiteLLM's own checker does — but with rate-limit exclusions first, so a +# throttling 429 is never mistaken for an overflow and sent into compaction. +_OVERFLOW_EXCLUSIONS = ( + "rate limit", + "too many requests", + "throttling", + "service unavailable", + "quota", +) +_OVERFLOW_MARKERS = ( + "context length", + "context window", + "context_length_exceeded", + "prompt is too long", + "input is too long", + "input length", + "maximum prompt length", + "reduce the length of the messages", + "too many tokens", + "token limit exceeded", + "request entity too large", +) + + +def is_context_overflow(exc: BaseException) -> bool: + """Whether ``exc`` is a model context-window-overflow error. + + LiteLLM types most providers' overflow as ContextWindowExceededError, but its + OpenRouter branch raises a plain BadRequestError, so for that we fall back to + matching the provider message. + """ + if isinstance(exc, ContextWindowExceededError): + return True + if isinstance(exc, BadRequestError): + msg = str(exc).lower() + if any(x in msg for x in _OVERFLOW_EXCLUSIONS): + return False + return any(x in msg for x in _OVERFLOW_MARKERS) + return False + + +_SUMMARY_INSTRUCTIONS = """\ +You are compacting the earlier part of an autonomous security-testing agent's \ +conversation so it fits the model context window. Produce a dense, factual \ +record that lets the agent continue with no loss of important state. + +This is a security engagement: dropped findings mean lost vulnerabilities. Be \ +EXHAUSTIVE, not concise. Enumerate every distinct item as its own bullet — \ +never merge, deduplicate, generalise, or omit distinct findings, credentials, \ +or dead ends, even if they seem minor or repetitive. If the source mentions \ +five vulnerabilities, list five. Copy exact values verbatim: URLs, endpoints, \ +file paths, parameters, payloads, credentials, tokens, keys, hashes, cracked \ +passwords, software versions, and error messages — never paraphrase or \ +placeholder them. Do not invent anything and do not describe this compaction \ +process. + +Return Markdown with exactly these sections: + +## Objective +The overall goal and target scope. + +## Vulnerabilities & Findings +One bullet per DISTINCT vulnerability or finding (SQLi, XSS, SSRF, auth bypass, \ +misconfig, etc.). For each: type, exact location (URL/endpoint/param/file), the \ +verbatim payload or proof, confirmation status, and impact. List them all. + +## Credentials & Secrets +One bullet per credential, secret, API key, token, hash, or cracked password, \ +copied verbatim with where it applies. Write "(none)" only if truly none. + +## System & Recon Details +Architecture, tech stack, versions, discovered endpoints/paths/params, and \ +other weak points worth keeping. + +## Work State +- Completed: what has been verified or finished. +- Active: what is in progress right now. +- Blocked: anything stuck and why. + +## Failed Attempts & Dead Ends +One bullet per approach already tried that did not work (including WAF blocks, \ +filtered inputs, non-exploitable leads) so they are not repeated. Write \ +"(none)" only if truly none. + +## Next Move +The concrete next step(s) the agent intended to take. + +## Relevant Files +Files/notes/reports created or modified and their purpose.""" + + +def _content_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + text = block.get("text") + if isinstance(text, str): + parts.append(text) + elif block.get("type") in {"input_image", "image_url", "output_image"}: + parts.append("[image]") + return "\n".join(parts) + return "" + + +def _truncate(text: str, limit: int) -> str: + return text if len(text) <= limit else f"{text[:limit]}\n[truncated]" + + +def _serialize_item(item: Any) -> 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) + 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)}" + if item_type == "reasoning": + return "" + if role or item_type == "message": + return f"[{role or 'assistant'}] {_content_text(item.get('content'))}".strip() + 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 _is_tool_call(item: Any) -> bool: + return isinstance(item, dict) and item.get("type") == "function_call" + + +def _is_tool_output(item: Any) -> bool: + return isinstance(item, dict) and item.get("type") == "function_call_output" + + +def _open_calls_at(items: list[Any]) -> list[int]: + """Prefix count of tool calls still awaiting their result at each index; + a split is only safe where this is zero.""" + balance = [0] * (len(items) + 1) + for i, item in enumerate(items): + delta = 1 if _is_tool_call(item) else -1 if _is_tool_output(item) else 0 + balance[i + 1] = max(0, balance[i] + delta) + return balance + + +def _select_split(model: str, items: list[Any], keep_tokens: int) -> int: + """Index where the kept-verbatim recent tail begins: walk newest→oldest to + ``keep_tokens``, then snap to a point with no tool call left open.""" + total = 0 + split = len(items) + for i in range(len(items) - 1, -1, -1): + total += count_tokens(model, _serialize_item(items[i])) + if total > keep_tokens: + break + split = i + open_calls = _open_calls_at(items) + while split > 0 and open_calls[split] != 0: + split -= 1 + return split + + +def _previous_summary(head: list[Any]) -> str | None: + for item in head: + if isinstance(item, dict) and item.get("role") == "user": + text = _content_text(item.get("content")) + if text.startswith(_CHECKPOINT_TAG): + return text + return None + + +def _fit_to_tokens(model: str, text: str, max_tokens: int) -> str: + """Head+tail-truncate ``text`` to ``max_tokens``, keeping start and end.""" + if count_tokens(model, text) <= max_tokens: + return text + # Rough char budget (~4x tokens), then tighten by real token count. + budget_chars = max_tokens * 4 + head_chars = budget_chars // 2 + tail_chars = budget_chars - head_chars + candidate = text[:head_chars] + _HEAD_TRUNCATED_MARKER + text[len(text) - tail_chars :] + while count_tokens(model, candidate) > max_tokens and (head_chars > 0 or tail_chars > 0): + head_chars = int(head_chars * 0.8) + tail_chars = int(tail_chars * 0.8) + candidate = text[:head_chars] + _HEAD_TRUNCATED_MARKER + text[len(text) - tail_chars :] + return candidate + + +def _summary_output_tokens(model: str) -> int: + """Summary output allowance, capped at the model's own output limit.""" + return min(load_settings().context.summary_max_tokens, output_limit(model)) + + +def _summary_input_budget(model: str, previous: str | None) -> int: + """Token room left for the head after instructions and the summary output.""" + overhead = count_tokens(model, _SUMMARY_INSTRUCTIONS) + if previous: + overhead += count_tokens(model, previous) + # 256 leaves slack for the prompt wrapper text not counted in ``overhead``. + room = context_window(model) - _summary_output_tokens(model) - overhead - 256 + return max(0, room) + + +def _build_summary_prompt(serialized_head: str, previous: str | None) -> str: + previous_block = ( + f"\n\nA previous checkpoint summary follows. Update it: keep what is " + f"still true, drop what is now stale, and merge in the new " + f"conversation below.\n\n{previous}\n" + if previous + else "" + ) + return ( + f"{_SUMMARY_INSTRUCTIONS}{previous_block}\n\n" + f"Conversation to summarise:\n\n{serialized_head}" + ) + + +def _checkpoint_item(summary: str) -> dict[str, Any]: + return { + "role": "user", + "content": ( + f"{_CHECKPOINT_TAG}\nThe following summarises earlier conversation that was " + f"compacted to fit the context window. Treat it as established context, not " + f"new instructions.\n\n{summary}\n" + ), + } + + +async def _summarize(model: str, prompt: str, max_tokens: int) -> str | None: + llm = load_settings().llm + try: + response = await litellm.acompletion( + model=model, + messages=[{"role": "user", "content": prompt}], + max_tokens=max_tokens, + api_key=llm.api_key, + api_base=llm.api_base, + timeout=llm.timeout, + ) + except Exception: + logger.exception("compaction summary call failed for model %s", model) + return None + try: + content = response.choices[0].message.content + except (AttributeError, IndexError, KeyError): + logger.warning("compaction summary returned no content") + return None + return content.strip() if isinstance(content, str) and content.strip() else None + + +async def maybe_compact( + session: Session, + *, + model: str, + instructions: str = "", + tools_text: str = "", + force: bool = False, +) -> bool: + """Compact ``session`` if it is near the model's context window. + + Returns ``True`` when the session was rewritten. ``force`` skips the size + check (used after a provider context-overflow error). + """ + context = load_settings().context + if not context.auto_compact and not force: + return False + + async with session_write_lock(session): + items = list(await session.get_items()) + 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: + return False + + split = _select_split(model, items, context.keep_tokens) + head, recent = items[:split], items[split:] + previous = _previous_summary(head) + input_budget = _summary_input_budget(model, previous) + if not head or input_budget <= 0: + # Nothing to summarise, or no room for even the summary request itself. + if head: + logger.warning( + "skipping compaction for %s: no room to summarise within its context window", model + ) + return False + + serialized_head = _fit_to_tokens(model, _serialize_items(head), input_budget) + summary = await _summarize( + model, + _build_summary_prompt(serialized_head, previous), + _summary_output_tokens(model), + ) + if summary is None: + return False + + new_items = [_checkpoint_item(summary), *recent] + rewritten = await replace_session_items(session, new_items, expected_len=len(items)) + if rewritten: + logger.info( + "compacted %s: %d items (~%d tok) -> %d items (summary + %d recent)", + model, + len(items), + used, + len(new_items), + len(recent), + ) + return rewritten diff --git a/strix/llm/context_budget.py b/strix/llm/context_budget.py new file mode 100644 index 00000000..3f153d91 --- /dev/null +++ b/strix/llm/context_budget.py @@ -0,0 +1,76 @@ +"""Model-aware token budgets, resolved from LiteLLM model metadata with a +large configurable fallback for models LiteLLM doesn't map. +""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from typing import Any + +import litellm + +from strix.config import load_settings + + +logger = logging.getLogger(__name__) + +# LiteLLM keys models without the routing prefix users type (``openai/``, +# ``litellm/``, ``ollama/`` ...). Strip a leading provider segment on lookup. +_STRIPPABLE_PREFIXES = ("openai/", "litellm/", "any-llm/", "ollama/", "ollama_chat/") + +_DEFAULT_OUTPUT_TOKENS = 8_192 + + +def _lookup_key(model: str) -> str: + for prefix in _STRIPPABLE_PREFIXES: + if model.startswith(prefix): + return model[len(prefix) :] + return model + + +def _safe_get_model_info(model: str) -> dict[str, Any] | None: + try: + return dict(litellm.get_model_info(model)) + except Exception: # noqa: BLE001 - unmapped models raise; caller falls back. + return None + + +@lru_cache(maxsize=128) +def _model_info(model: str) -> dict[str, int]: + for candidate in (model, _lookup_key(model)): + info = _safe_get_model_info(candidate) + if info is not None: + return { + "max_input_tokens": int( + info.get("max_input_tokens") or info.get("max_tokens") or 0 + ), + "max_output_tokens": int(info.get("max_output_tokens") or 0), + } + logger.debug("No LiteLLM model info for %r; using configured fallbacks", model) + return {"max_input_tokens": 0, "max_output_tokens": 0} + + +def context_window(model: str) -> int: + """Input token capacity for ``model`` (configured fallback when unmapped).""" + resolved = _model_info(model)["max_input_tokens"] + return resolved or load_settings().context.fallback_context_tokens + + +def output_limit(model: str) -> int: + """Max output tokens for ``model`` (a conservative default when unmapped).""" + return _model_info(model)["max_output_tokens"] or _DEFAULT_OUTPUT_TOKENS + + +def count_tokens(model: str, text: str) -> int: + """Token count for ``text`` under ``model``. + + Falls back to UTF-8 byte length (a guaranteed upper bound) when LiteLLM + can't count, so budget checks stay conservative. + """ + if not text: + return 0 + try: + return int(litellm.token_counter(model=_lookup_key(model), text=text)) + except Exception: # noqa: BLE001 - tokenizer may be unavailable for some models. + return len(text.encode("utf-8")) diff --git a/tests/test_compaction.py b/tests/test_compaction.py new file mode 100644 index 00000000..da9a2361 --- /dev/null +++ b/tests/test_compaction.py @@ -0,0 +1,306 @@ +"""Tests for provider-agnostic conversation compaction.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +import pytest +from litellm.exceptions import BadRequestError, ContextWindowExceededError, RateLimitError + +from strix.config import ContextSettings +from strix.llm import compaction + + +if TYPE_CHECKING: + from agents.memory.session_settings import SessionSettings + + +class FakeSession: + """Minimal in-memory Session for exercising compaction.""" + + session_id = "fake" + session_settings: SessionSettings | None = None + + def __init__(self, items: list[Any]) -> None: + self._items = list(items) + + async def get_items(self, limit: int | None = None) -> list[Any]: + return list(self._items) if limit is None else list(self._items[-limit:]) + + async def add_items(self, items: list[Any]) -> None: + self._items.extend(items) + + async def clear_session(self) -> None: + self._items = [] + + async def pop_item(self) -> Any: + return self._items.pop() if self._items else None + + +def _user(text: str) -> dict[str, Any]: + return {"role": "user", "content": text} + + +def _assistant(text: str) -> dict[str, Any]: + return {"role": "assistant", "content": text} + + +def _call(call_id: str, name: str = "exec_command") -> dict[str, Any]: + return {"type": "function_call", "call_id": call_id, "name": name, "arguments": "{}"} + + +def _output(call_id: str, text: str = "done") -> dict[str, Any]: + return {"type": "function_call_output", "call_id": call_id, "output": text} + + +def _turns(n: int) -> list[dict[str, Any]]: + items: list[dict[str, Any]] = [] + for i in range(n): + items += [ + _user(f"task {i}"), + _call(f"c{i}"), + _output(f"c{i}", f"result {i}"), + _assistant(f"ok {i}"), + ] + return items + + +def _has_orphan_tool_output(items: list[Any]) -> bool: + call_ids = {i["call_id"] for i in items if compaction._is_tool_call(i)} + return any(i["call_id"] not in call_ids for i in items if compaction._is_tool_output(i)) + + +def test_is_context_overflow_uses_litellm_typed_error() -> None: + overflow = ContextWindowExceededError( + message="context length exceeded", model="m", llm_provider="openai" + ) + assert compaction.is_context_overflow(overflow) + assert not compaction.is_context_overflow( + RateLimitError(message="slow down", model="m", llm_provider="openai") + ) + assert not compaction.is_context_overflow(RuntimeError("maximum context length is 8192")) + + +def test_is_context_overflow_matches_untyped_openrouter_400() -> None: + # OpenRouter overflows arrive as a plain BadRequestError, so match the message. + openrouter = BadRequestError( + message=( + "litellm.BadRequestError: This endpoint's maximum context length is 16385 " + "tokens. However, you requested about 75064 tokens. Please reduce the length " + "of the messages." + ), + model="openrouter/openai/gpt-3.5-turbo", + llm_provider="openrouter", + ) + assert compaction.is_context_overflow(openrouter) + + +def test_is_context_overflow_ignores_rate_limit_shaped_bad_request() -> None: + # A 400 that is really throttling must never be treated as an overflow. + throttled = BadRequestError( + message="Rate limit exceeded, please slow down", + model="openrouter/openai/gpt-4o", + llm_provider="openrouter", + ) + assert not compaction.is_context_overflow(throttled) + unrelated = BadRequestError( + message="Invalid value for 'temperature'", model="m", llm_provider="openrouter" + ) + assert not compaction.is_context_overflow(unrelated) + + +def test_select_split_never_orphans_tool_output(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t)) + items = _turns(10) + split = compaction._select_split("m", items, keep_tokens=25) + recent = items[split:] + assert recent # something is kept + assert not _has_orphan_tool_output(recent) + + +def test_select_split_handles_parallel_calls(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(compaction, "count_tokens", lambda _m, _t: 1) + # Two parallel calls then their two outputs. + items = [ + _user("start"), + _call("a"), + _call("b"), + _output("a"), + _output("b"), + _assistant("done"), + ] + # keep_tokens picks a boundary that would land between the calls/outputs. + split = compaction._select_split("m", items, keep_tokens=2) + assert not _has_orphan_tool_output(items[split:]) + + +def _patch_budget(monkeypatch: pytest.MonkeyPatch, *, keep_tokens: int, window: int) -> None: + 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) + context = ContextSettings() + context.keep_tokens = keep_tokens + context.compact_buffer_tokens = 0 + context.summary_max_tokens = 64 + context.auto_compact = True + settings = SimpleNamespace( + context=context, + llm=SimpleNamespace(api_key=None, api_base=None, timeout=1), + ) + monkeypatch.setattr(compaction, "load_settings", lambda: settings) + + +def _patch_summary(monkeypatch: pytest.MonkeyPatch, text: str) -> None: + async def fake_acompletion(**_kwargs: Any) -> Any: + message = SimpleNamespace(content=text) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + + +@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) + session = FakeSession(_turns(10)) + before = await session.get_items() + + assert await compaction.maybe_compact(session, model="m") is False + assert await session.get_items() == before + + +@pytest.mark.asyncio +async def test_maybe_compact_rewrites_and_keeps_pairs(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_budget(monkeypatch, keep_tokens=30, window=4_000) + _patch_summary(monkeypatch, "SUMMARY BODY") + session = FakeSession(_turns(12)) + + assert await compaction.maybe_compact(session, model="m", force=True) is True + + items = await session.get_items() + assert items[0]["role"] == "user" + assert items[0]["content"].startswith(compaction._CHECKPOINT_TAG) + assert "SUMMARY BODY" in items[0]["content"] + assert len(items) < len(_turns(12)) + assert not _has_orphan_tool_output(items) + + +@pytest.mark.asyncio +async def test_maybe_compact_updates_previous_summary(monkeypatch: pytest.MonkeyPatch) -> None: + # Window large enough to leave real room for the summary instructions. + _patch_budget(monkeypatch, keep_tokens=30, window=4_000) + captured: dict[str, str] = {} + + async def fake_acompletion(**kwargs: Any) -> Any: + captured["prompt"] = kwargs["messages"][0]["content"] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="NEW"))]) + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + + prior = compaction._checkpoint_item("OLD SUMMARY TEXT") + session = FakeSession([prior, *_turns(12)]) + + assert await compaction.maybe_compact(session, model="m", force=True) is True + assert "OLD SUMMARY TEXT" in captured["prompt"] + + +def test_fit_to_tokens_truncates_oversized_text(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(compaction, "count_tokens", lambda _m, t: len(t)) + text = "x" * 10_000 + + fitted = compaction._fit_to_tokens("m", text, 500) + + assert len(fitted) <= 500 + assert compaction._HEAD_TRUNCATED_MARKER in fitted + # Small text is returned untouched. + assert compaction._fit_to_tokens("m", "short", 500) == "short" + + +def test_summary_output_tokens_capped_at_model_limit(monkeypatch: pytest.MonkeyPatch) -> None: + context = ContextSettings() + monkeypatch.setattr(compaction, "load_settings", lambda: SimpleNamespace(context=context)) + + monkeypatch.setattr(compaction, "output_limit", lambda _m: 1_000) + context.summary_max_tokens = 4_096 + # Configured allowance above the model cap is clamped down to the cap. + assert compaction._summary_output_tokens("m") == 1_000 + # Below the cap, the configured value is used unchanged. + context.summary_max_tokens = 500 + assert compaction._summary_output_tokens("m") == 500 + + +@pytest.mark.asyncio +async def test_maybe_compact_bounds_summary_prompt(monkeypatch: pytest.MonkeyPatch) -> None: + # A tiny window with a huge head must not send an oversized summary request. + _patch_budget(monkeypatch, keep_tokens=30, window=4_000) + captured: dict[str, str] = {} + + async def fake_acompletion(**kwargs: Any) -> Any: + captured["prompt"] = kwargs["messages"][0]["content"] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))]) + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + big_turns = [{"role": "user", "content": "y" * 2_000} for _ in range(50)] + session = FakeSession(big_turns) + + assert await compaction.maybe_compact(session, model="m") is True + # count_tokens==len(chars); prompt must fit the model window. + assert len(captured["prompt"]) <= 4_000 + + +@pytest.mark.asyncio +async def test_summary_request_fits_when_room_is_below_old_floor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Head-input budget must shrink to the real room so the request fits. + instructions = len(compaction._SUMMARY_INSTRUCTIONS) + window = instructions + 64 + 256 + 300 # summary_max(64)+slack(256)+room(300) + _patch_budget(monkeypatch, keep_tokens=30, window=window) + captured: dict[str, str] = {} + + async def fake_acompletion(**kwargs: Any) -> Any: + captured["prompt"] = kwargs["messages"][0]["content"] + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))]) + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + session = FakeSession([{"role": "user", "content": "y" * 5_000} for _ in range(20)]) + + assert await compaction.maybe_compact(session, model="m") is True + assert len(captured["prompt"]) <= window + + +@pytest.mark.asyncio +async def test_maybe_compact_skips_when_summary_fails(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_budget(monkeypatch, keep_tokens=30, window=4_000) + + async def fake_acompletion(**_kwargs: Any) -> Any: + raise RuntimeError("boom") + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + session = FakeSession(_turns(12)) + before = await session.get_items() + + assert await compaction.maybe_compact(session, model="m", force=True) is False + assert await session.get_items() == before + + +@pytest.mark.asyncio +async def test_maybe_compact_skips_when_no_room_to_summarise( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # No room for any head -> no (doomed) summary is attempted. + _patch_budget(monkeypatch, keep_tokens=30, window=200) + called = False + + async def fake_acompletion(**_kwargs: Any) -> Any: + nonlocal called + called = True + return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="S"))]) + + monkeypatch.setattr("strix.llm.compaction.litellm.acompletion", fake_acompletion) + session = FakeSession(_turns(12)) + before = await session.get_items() + + assert await compaction.maybe_compact(session, model="m", force=True) is False + assert called is False + assert await session.get_items() == before diff --git a/tests/test_context_budget.py b/tests/test_context_budget.py new file mode 100644 index 00000000..8b34b483 --- /dev/null +++ b/tests/test_context_budget.py @@ -0,0 +1,47 @@ +"""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) + # Falls back to UTF-8 byte length (upper bound on tokens). + assert context_budget.count_tokens("weird-model", "x" * 400) == 400 + assert context_budget.count_tokens("weird-model", "😀" * 10) == 40 + + +def test_count_tokens_empty_is_zero() -> None: + assert context_budget.count_tokens("gpt-4o", "") == 0