diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py index ffaca66..e63b1e5 100644 --- a/amplifier_module_context_simple/__init__.py +++ b/amplifier_module_context_simple/__init__.py @@ -55,6 +55,9 @@ async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = config = config or {} context = SimpleContextManager( max_tokens=config.get("max_tokens", 200_000), + # Track whether max_tokens was *explicitly* set. When it is, we treat it + # as a hard ceiling on the provider-derived budget (see _calculate_budget). + max_tokens_explicit="max_tokens" in config, compact_threshold=config.get("compact_threshold", 0.92), target_usage=config.get("target_usage", 0.50), protected_recent=config.get("protected_recent", 0.30), @@ -111,6 +114,7 @@ class SimpleContextManager: def __init__( self, max_tokens: int = 200_000, + max_tokens_explicit: bool = False, compact_threshold: float = 0.92, target_usage: float = 0.50, protected_recent: float = 0.30, @@ -127,7 +131,14 @@ def __init__( Initialize the context manager. Args: - max_tokens: Maximum context size in tokens + max_tokens: Maximum context size in tokens. When max_tokens_explicit + is True this also acts as a hard ceiling on any provider-derived + budget (see _calculate_budget). + max_tokens_explicit: True when max_tokens was explicitly configured by + the user (as opposed to the default). When True, the effective + budget is capped at max_tokens even if the provider advertises a + larger context window. Defaults to False to preserve the historical + "provider window wins" behavior for callers that do not set it. compact_threshold: Trigger compaction at this usage ratio (0.0-1.0) target_usage: Compact down to this usage ratio (0.0-1.0) protected_recent: Always protect last N% of messages (0.0-1.0) @@ -144,6 +155,7 @@ def __init__( """ self.messages: list[dict[str, Any]] = [] self.max_tokens = max_tokens + self.max_tokens_explicit = max_tokens_explicit self.compact_threshold = compact_threshold self.target_usage = target_usage self.protected_recent = protected_recent @@ -1154,6 +1166,27 @@ def _format_affected_items(self, level: int, stats: dict[str, Any]) -> str: "- If context is critical, consider asking user to clarify their current goal" ) + def _apply_max_tokens_cap(self, budget: int) -> int: + """Cap a provider-derived budget at the explicitly configured max_tokens. + + When max_tokens was set explicitly (max_tokens_explicit=True), it acts as a + hard ceiling: the effective budget is min(provider_budget, max_tokens). This + lets an operator bound context growth even on providers that advertise a very + large context window (e.g. a 1M-token model), where the provider-derived + budget would otherwise be hundreds of thousands of tokens and compaction + would effectively never trigger. + + When max_tokens was not set explicitly, the provider budget is returned + unchanged, preserving the historical "provider window wins" behavior. + """ + if self.max_tokens_explicit and budget > self.max_tokens: + logger.info( + f"Capping provider-derived budget {budget:,} to explicitly " + f"configured max_tokens {self.max_tokens:,}" + ) + return self.max_tokens + return budget + def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> int: """Calculate effective token budget from provider or fallback to config. @@ -1163,6 +1196,12 @@ def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> i 3. Provider defaults (legacy: some providers may put limits here) 4. Configured max_tokens fallback + Provider-derived budgets (cases 2 and 3) are additionally capped at the + configured max_tokens when max_tokens was set explicitly (see + _apply_max_tokens_cap). This lets an operator bound context growth even + when the provider advertises a much larger window (e.g. a 1M-token model), + which is otherwise impossible because max_tokens is only a fallback. + Note: We reserve only 50% of max_output_tokens since most responses are much smaller than the maximum. This prevents over-conservative budgets that would trigger compaction too early. @@ -1193,7 +1232,7 @@ def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> i f"(context={context_window:,}, reserved_output={reserved_output:,} " f"[{output_reserve_fraction:.0%} of {max_output:,}])" ) - return budget + return self._apply_max_tokens_cap(budget) # Check provider info defaults (legacy approach) info = provider.get_info() @@ -1209,7 +1248,7 @@ def _calculate_budget(self, token_budget: int | None, provider: Any | None) -> i f"(context={context_window:,}, reserved_output={reserved_output:,} " f"[{output_reserve_fraction:.0%} of {max_output_tokens:,}])" ) - return budget + return self._apply_max_tokens_cap(budget) else: logger.debug( f"Provider defaults missing context_window ({context_window}) " diff --git a/tests/test_max_tokens_cap.py b/tests/test_max_tokens_cap.py new file mode 100644 index 0000000..5071bfa --- /dev/null +++ b/tests/test_max_tokens_cap.py @@ -0,0 +1,104 @@ +"""Tests for honoring an explicitly-configured max_tokens as a budget ceiling. + +Regression coverage for the bug where max_tokens was silently ignored whenever a +provider advertised a context window. On large-context models (e.g. a 1M-token +provider) the provider-derived budget was hundreds of thousands of tokens, so the +compaction threshold was never reached and context grew unbounded -- even when an +operator had explicitly set max_tokens to bound it. + +The fix: when max_tokens is set explicitly (max_tokens_explicit=True), it acts as a +hard ceiling -- budget = min(provider_budget, max_tokens). When it is NOT explicit, +the provider window wins (historical behavior preserved). +""" + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest + +from amplifier_module_context_simple import SimpleContextManager, mount + +# Provider-derived budget for a 1M window with the default 0.5 output reserve: +# 1_000_000 - int(128_000 * 0.5) - 4096 (safety margin) == 931_904 +PROVIDER_BUDGET_1M = 1_000_000 - int(128_000 * 0.5) - 4096 + + +@dataclass +class _ModelInfo: + context_window: int + max_output_tokens: int + + +class _ModelInfoProvider: + """Provider exposing get_model_info() -- the modern budget path.""" + + def __init__(self, context_window: int, max_output_tokens: int): + self._info = _ModelInfo(context_window, max_output_tokens) + + def get_model_info(self): + return self._info + + +class _DefaultsProvider: + """Provider exposing only get_info().defaults -- the legacy budget path.""" + + def __init__(self, context_window: int, max_output_tokens: int): + self._defaults = { + "context_window": context_window, + "max_output_tokens": max_output_tokens, + } + + def get_info(self): + return SimpleNamespace(defaults=self._defaults) + + +class _FakeCoordinator: + def __init__(self): + self.hooks = None + self.mounted: dict[str, object] = {} + + async def mount(self, name, obj): + self.mounted[name] = obj + + +def test_explicit_max_tokens_caps_provider_budget(): + ctx = SimpleContextManager(max_tokens=300_000, max_tokens_explicit=True) + provider = _ModelInfoProvider(context_window=1_000_000, max_output_tokens=128_000) + assert ctx._calculate_budget(None, provider) == 300_000 + + +def test_non_explicit_max_tokens_uses_full_provider_budget(): + # Default (not explicit): provider window wins -- historical behavior preserved. + ctx = SimpleContextManager() + provider = _ModelInfoProvider(context_window=1_000_000, max_output_tokens=128_000) + assert ctx._calculate_budget(None, provider) == PROVIDER_BUDGET_1M + + +def test_explicit_cap_never_inflates_above_provider_budget(): + # The cap only lowers the budget; it never raises it above the provider value. + ctx = SimpleContextManager(max_tokens=5_000_000, max_tokens_explicit=True) + provider = _ModelInfoProvider(context_window=1_000_000, max_output_tokens=128_000) + assert ctx._calculate_budget(None, provider) == PROVIDER_BUDGET_1M + + +def test_cap_applies_to_legacy_defaults_path(): + ctx = SimpleContextManager(max_tokens=250_000, max_tokens_explicit=True) + provider = _DefaultsProvider(context_window=1_000_000, max_output_tokens=128_000) + assert ctx._calculate_budget(None, provider) == 250_000 + + +@pytest.mark.asyncio +async def test_mount_marks_explicit_only_when_configured(): + coord = _FakeCoordinator() + await mount(coord, {"max_tokens": 123_456}) + ctx = coord.mounted["context"] + assert ctx.max_tokens_explicit is True + assert ctx.max_tokens == 123_456 + + +@pytest.mark.asyncio +async def test_mount_no_explicit_flag_when_max_tokens_absent(): + coord = _FakeCoordinator() + await mount(coord, {}) + ctx = coord.mounted["context"] + assert ctx.max_tokens_explicit is False