diff --git a/agent/agent_init.py b/agent/agent_init.py index 96bfe3d873f0..43a38aa92bba 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -1540,6 +1540,11 @@ def init_agent( api_mode=agent.api_mode, abort_on_summary_failure=compression_abort_on_summary_failure, ) + _sync_output_reservation = getattr( + agent, "_sync_context_compressor_output_reservation", None + ) + if callable(_sync_output_reservation): + _sync_output_reservation() agent.compression_enabled = compression_enabled # Reject models whose context window is below the minimum required diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index f144b3f67a63..173b2055ccad 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -1611,6 +1611,11 @@ def switch_model(agent, new_model, new_provider, api_key='', base_url='', api_mo provider=agent.provider, api_mode=agent.api_mode, ) + _sync_output_reservation = getattr( + agent, "_sync_context_compressor_output_reservation", None + ) + if callable(_sync_output_reservation): + _sync_output_reservation() # ── Invalidate cached system prompt so it rebuilds next turn ── agent._cached_system_prompt = None diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 1ee1702b45e8..29f198fec274 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -1286,6 +1286,11 @@ def try_activate_fallback(agent, reason: "FailoverReason | None" = None) -> bool provider=agent.provider, api_mode=agent.api_mode, ) + _sync_output_reservation = getattr( + agent, "_sync_context_compressor_output_reservation", None + ) + if callable(_sync_output_reservation): + _sync_output_reservation() agent._buffer_status( f"🔄 Primary model failed — switching to fallback: " diff --git a/agent/context_compressor.py b/agent/context_compressor.py index 16db1bedc30f..134730cd9fb2 100644 --- a/agent/context_compressor.py +++ b/agent/context_compressor.py @@ -640,6 +640,48 @@ def on_session_end(self, session_id: str, messages: List[Dict[str, Any]]) -> Non """ self._previous_summary = None + @staticmethod + def _coerce_output_reservation_tokens(value: Any) -> int: + """Return a positive output-token reservation, or 0 when unset.""" + if isinstance(value, bool): + return 0 + try: + coerced = int(value) + except (OverflowError, TypeError, ValueError): + return 0 + return max(coerced, 0) + + def _calculate_threshold_tokens(self) -> int: + """Calculate the prompt-token threshold for proactive compression.""" + context_length = max(int(self.context_length or 0), 1) + output_reservation = self._coerce_output_reservation_tokens( + getattr(self, "output_reservation_tokens", 0) + ) + if output_reservation > 0: + input_budget = max(context_length - output_reservation, 1) + return max(int(input_budget * self.threshold_percent), 1) + return max( + int(context_length * self.threshold_percent), + MINIMUM_CONTEXT_LENGTH, + ) + + def _recalculate_token_budgets(self) -> None: + self.threshold_tokens = self._calculate_threshold_tokens() + self.tail_token_budget = int( + self.threshold_tokens * self.summary_target_ratio + ) + self.max_summary_tokens = min( + int(max(int(self.context_length or 0), 1) * 0.05), + _SUMMARY_TOKENS_CEILING, + ) + + def set_output_reservation_tokens(self, output_reservation_tokens: int) -> None: + """Update reserved response-token budget and derived thresholds.""" + self.output_reservation_tokens = self._coerce_output_reservation_tokens( + output_reservation_tokens + ) + self._recalculate_token_budgets() + def update_model( self, model: str, @@ -656,17 +698,10 @@ def update_model( self.provider = provider self.api_mode = api_mode self.context_length = context_length - self.threshold_tokens = max( - int(context_length * self.threshold_percent), - MINIMUM_CONTEXT_LENGTH, - ) # Recalculate token budgets for the new context length so the # compressor stays calibrated after a model switch (e.g. 200K → 32K). - target_tokens = int(self.threshold_tokens * self.summary_target_ratio) - self.tail_token_budget = target_tokens - self.max_summary_tokens = min( - int(context_length * 0.05), _SUMMARY_TOKENS_CEILING, - ) + + self._recalculate_token_budgets() def __init__( self, @@ -683,6 +718,7 @@ def __init__( provider: str = "", api_mode: str = "", abort_on_summary_failure: bool = False, + output_reservation_tokens: int = 0, ): self.model = model self.base_url = base_url @@ -705,30 +741,22 @@ def __init__( config_context_length=config_context_length, provider=provider, ) - # Floor: never compress below MINIMUM_CONTEXT_LENGTH tokens even if - # the percentage would suggest a lower value. This prevents premature - # compression on large-context models at 50% while keeping the % sane - # for models right at the minimum. - self.threshold_tokens = max( - int(self.context_length * threshold_percent), - MINIMUM_CONTEXT_LENGTH, + self.output_reservation_tokens = self._coerce_output_reservation_tokens( + output_reservation_tokens ) self.compression_count = 0 # Derive token budgets: ratio is relative to the threshold, not total context - target_tokens = int(self.threshold_tokens * self.summary_target_ratio) - self.tail_token_budget = target_tokens - self.max_summary_tokens = min( - int(self.context_length * 0.05), _SUMMARY_TOKENS_CEILING, - ) + self._recalculate_token_budgets() if not quiet_mode: logger.info( "Context compressor initialized: model=%s context_length=%d " - "threshold=%d (%.0f%%) target_ratio=%.0f%% tail_budget=%d " - "provider=%s base_url=%s", + "threshold=%d (%.0f%%) output_reservation=%d " + "target_ratio=%.0f%% tail_budget=%d provider=%s base_url=%s", model, self.context_length, self.threshold_tokens, - threshold_percent * 100, self.summary_target_ratio * 100, + threshold_percent * 100, self.output_reservation_tokens, + self.summary_target_ratio * 100, self.tail_token_budget, provider or "none", base_url or "none", ) diff --git a/agent/turn_context.py b/agent/turn_context.py index e94d43279abd..eb3c9f238289 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -251,6 +251,12 @@ def build_turn_context( and len(messages) > agent.context_compressor.protect_first_n + agent.context_compressor.protect_last_n + 1 ): + _sync_output_reservation = getattr( + agent, "_sync_context_compressor_output_reservation", None + ) + if callable(_sync_output_reservation): + _sync_output_reservation() + _preflight_tokens = estimate_request_tokens_rough( messages, system_prompt=active_system_prompt or "", diff --git a/run_agent.py b/run_agent.py index 0be8b1763fa4..19af692ac8fa 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1289,6 +1289,102 @@ def _requested_output_cap_from_api_kwargs(api_kwargs: Any) -> Optional[int]: return value return None + @staticmethod + def _positive_output_token_cap(value: Any) -> Optional[int]: + """Return a positive integer output cap, or None when unset.""" + if isinstance(value, bool): + return None + try: + cap = int(value) + except (OverflowError, TypeError, ValueError): + return None + return cap if cap > 0 else None + + def _resolved_output_reservation_tokens(self) -> int: + """Resolve the response-token budget the next request will reserve.""" + override_cap = self._requested_output_cap_from_api_kwargs( + getattr(self, "request_overrides", None) + ) + if override_cap is not None: + return override_cap + + ephemeral_cap = self._positive_output_token_cap( + getattr(self, "_ephemeral_max_output_tokens", None) + ) + if ephemeral_cap is not None: + return ephemeral_cap + + user_cap = self._positive_output_token_cap(getattr(self, "max_tokens", None)) + if user_cap is not None: + return user_cap + + api_mode = (getattr(self, "api_mode", "") or "").strip().lower() + model = getattr(self, "model", "") or "" + if api_mode == "anthropic_messages": + try: + from agent.anthropic_adapter import _resolve_anthropic_messages_max_tokens + + cap = _resolve_anthropic_messages_max_tokens( + None, + model, + context_length=getattr( + getattr(self, "context_compressor", None), + "context_length", + None, + ), + ) + context_length = getattr( + getattr(self, "context_compressor", None), + "context_length", + None, + ) + if context_length and cap >= context_length: + cap = max(int(context_length) - 1, 1) + return cap + except Exception: + return 0 + + try: + from providers import get_provider_profile + + profile = get_provider_profile(getattr(self, "provider", "")) + except Exception: + profile = None + if profile is not None: + profile_cap = self._positive_output_token_cap( + profile.get_max_tokens(model) + ) + if profile_cap is not None: + return profile_cap + + base_lower = (getattr(self, "_base_url_lower", None) or getattr( + self, "base_url", "" + ) or "").lower() + is_nous = "nousresearch" in base_lower + try: + is_openrouter = self._is_openrouter_url() + except Exception: + is_openrouter = "openrouter.ai" in base_lower + if (is_openrouter or is_nous) and "claude" in model.lower(): + try: + from agent.anthropic_adapter import _get_anthropic_max_output + + return _get_anthropic_max_output(model) + except Exception: + return 0 + + return 0 + + def _sync_context_compressor_output_reservation(self) -> int: + """Push the outgoing output-token reservation into the compressor.""" + compressor = getattr(self, "context_compressor", None) + setter = getattr(compressor, "set_output_reservation_tokens", None) + if not callable(setter): + return 0 + output_reservation = self._resolved_output_reservation_tokens() + setter(output_reservation) + return output_reservation + def _has_content_after_think_block(self, content: str) -> bool: """ Check if content has actual text after any reasoning/thinking blocks. diff --git a/tests/test_ctx_halving_fix.py b/tests/test_ctx_halving_fix.py index 63c965ac9658..ed58108afa37 100644 --- a/tests/test_ctx_halving_fix.py +++ b/tests/test_ctx_halving_fix.py @@ -25,7 +25,7 @@ import sys import os -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) @@ -360,3 +360,57 @@ def test_safety_margin_never_goes_below_one(self): available_out = parse_available_output_tokens_from_error(error_msg) safe_out = max(1, available_out - 64) assert safe_out == 1 + + +# --------------------------------------------------------------------------- +# Proactive compression threshold reserves the outgoing output-token budget +# --------------------------------------------------------------------------- + + +class TestOutputReservationAwareCompressionThreshold: + """Proactive compaction must leave room for max_tokens in the same window.""" + + def test_output_reservation_reduces_threshold_without_minimum_floor(self): + from agent.context_compressor import ContextCompressor + + with patch("agent.context_compressor.get_model_context_length", return_value=131_072): + compressor = ContextCompressor( + model="local/gemma", + quiet_mode=True, + threshold_percent=0.50, + output_reservation_tokens=65_536, + ) + + assert compressor.context_length == 131_072 + assert compressor.output_reservation_tokens == 65_536 + assert compressor.threshold_tokens == 32_768 + assert compressor.should_compress(43_089) is True + + def test_custom_provider_default_max_tokens_updates_compressor_threshold(self): + from agent.context_compressor import ContextCompressor + from run_agent import AIAgent + + agent = object.__new__(AIAgent) + agent.model = "local/gemma" + agent.provider = "custom" + agent.base_url = "http://localhost:8000/v1" + agent._base_url_lower = agent.base_url.lower() + agent.api_mode = "chat_completions" + agent.max_tokens = None + agent.request_overrides = {} + agent._ephemeral_max_output_tokens = None + + with patch("agent.context_compressor.get_model_context_length", return_value=131_072): + agent.context_compressor = ContextCompressor( + model=agent.model, + provider=agent.provider, + base_url=agent.base_url, + quiet_mode=True, + threshold_percent=0.50, + ) + + output_reservation = agent._sync_context_compressor_output_reservation() + + assert output_reservation == 65_536 + assert agent.context_compressor.output_reservation_tokens == 65_536 + assert agent.context_compressor.threshold_tokens == 32_768