diff --git a/cecli/args.py b/cecli/args.py index 317523944fe..3e6dcb7c33b 100644 --- a/cecli/args.py +++ b/cecli/args.py @@ -449,9 +449,17 @@ def get_parser(default_config_files, git_root): group.add_argument( "--enable-context-compaction", action=argparse.BooleanOptionalAction, - default=True, + default=False, help="Enable automatic compaction of chat history to conserve tokens (default: False)", ) + group.add_argument( + "--max-compaction-retries", + type=int, + default=3, + metavar="N", + choices=range(1, 11), + help="Maximum number of compactions per request before propagating the error (default: 3, min: 1, max: 10)", + ) group.add_argument( "--context-compaction-max-tokens", type=float, @@ -467,6 +475,16 @@ def get_parser(default_config_files, git_root): default=4096, help="The target maximum number of tokens for the generated summary. (default: 4096)", ) + group.add_argument( + "--ignore-context-limit", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Ignore the model's context window limit when sending requests." + " WARNING: This may cause API errors or truncated responses." + " Not settable via env var alone in agent mode (default: False)" + ), + ) ########## group = parser.add_argument_group("Cache settings") diff --git a/cecli/coders/base_coder.py b/cecli/coders/base_coder.py index 2633bc319e1..1f475ae2d32 100755 --- a/cecli/coders/base_coder.py +++ b/cecli/coders/base_coder.py @@ -409,6 +409,7 @@ def __init__( auto_accept_architect=True, mcp_manager=None, enable_context_compaction=False, + max_compaction_retries=3, context_compaction_max_tokens=None, context_compaction_summary_tokens=8192, map_cache_dir=".", @@ -477,6 +478,11 @@ def __init__( self.context_compaction_current_ratio = 0 self.context_compaction_max_tokens = context_compaction_max_tokens self.context_compaction_summary_tokens = context_compaction_summary_tokens + self.max_compaction_retries = max_compaction_retries + + # Token floor guard: prevents infinite compaction on near-empty context + self._compaction_floor_reached = False + self.max_reflections = nested.getter(self.args, "max_reflections", 3) if not fnames: @@ -1728,36 +1734,6 @@ async def generate(self, user_message, preproc): self.interrupt_event.clear() try: - if self.enable_context_compaction: - # Skip compaction if the user wants to clear or exit - # Compacting is wasteful since /clear will clear everything - # and /exit will exit the application - stripped = user_message.strip() - is_command = self.commands.is_command(stripped) - is_allowed_command = False - - if is_command: - res = self.commands.matching_commands(user_message) - if res is not None: - matching_commands, first_word, rest_inp = res - if len(matching_commands) == 1: - command = matching_commands[0] - splits = (rest_inp or "").split() - split_map = { - "/agent": 1, - "/architect": 1, - "/ask": 1, - "/code": 1, - "/model": 2, - } - if command in split_map and len(splits) >= split_map.get(command, 0): - is_allowed_command = True - - if not is_command or is_allowed_command: - self.compact_context_completed = False - await self.compact_context_if_needed() - self.compact_context_completed = True - self.run_one_completed = False await self.run_one(user_message, preproc) self.show_undo_hint() @@ -1897,9 +1873,6 @@ async def run_one(self, user_message, preproc): await self.auto_save_session(force=True) break - if self.enable_context_compaction: - await self.compact_context_if_needed() - if nested.getter(self, "agent_finished", False): await self.auto_save_session(force=True) break @@ -1986,10 +1959,14 @@ def keyboard_interrupt(self): # Old summarization system removed - using context compaction logic instead - async def compact_context_if_needed(self, force=False, message=""): + async def compact_context_if_needed(self, force=False, message="", retry_index=0): if not self.enable_context_compaction: return + # Skip compaction if token floor has been reached (context too small to compact further) + if self._compaction_floor_reached: + return + # Trigger background observation/reflection check await ObservationService.get_instance(self).check_and_trigger() @@ -2007,6 +1984,25 @@ async def compact_context_if_needed(self, force=False, message=""): file_context_tokens = self.summarizer.count_tokens(file_context_messages) all_tokens = self.summarizer.count_tokens(all_messages) + # Determine if compaction is worthwhile + compactable_tokens = done_tokens + cur_tokens + diff_tokens + + # Condition 1: Percentage check (is chat history a significant part of the context?) + is_worth_by_percentage = all_tokens > 0 and (compactable_tokens / all_tokens) >= 0.20 + + # Condition 2: Absolute savings check (would compacting save enough tokens to fit?) + tokens_over_limit = all_tokens - (self.context_compaction_max_tokens or all_tokens) + potential_savings = compactable_tokens * 0.90 + is_worth_by_absolute_savings = ( + tokens_over_limit > 0 and potential_savings >= tokens_over_limit + ) + + if not (is_worth_by_percentage or is_worth_by_absolute_savings): + self.io.tool_output( + "Skipping compaction: Not enough chat history to make a difference. Use /drop to remove files." + ) + return + message_tokens = done_tokens + cur_tokens file_tokens = diff_tokens + file_context_tokens combined_tokens = done_tokens + cur_tokens + diff_tokens + file_context_tokens @@ -2123,7 +2119,27 @@ async def summarize_and_update(messages, tag): if cur_tokens > self.context_compaction_max_tokens or cur_tokens > done_tokens: await summarize_and_update(cur_messages, MessageTag.CUR) - self.io.tool_output("...chat history compacted.") + self.io.update_spinner(self.io.last_spinner_text) + + # Post-compaction token floor check + # Recalculate tokens after compaction to prevent infinite loops + # on already-minimal context + all_messages = manager.get_messages_dict() + post_compaction_tokens = self.summarizer.count_tokens(all_messages) + token_floor = ( + self.context_compaction_max_tokens * 0.25 + if self.context_compaction_max_tokens + else 0 + ) + + if post_compaction_tokens < token_floor: + self._compaction_floor_reached = True + self.io.tool_output( + "...context is already at minimum size, cannot compact further." + ) + self.io.update_spinner(self.io.last_spinner_text) + return False + self.io.update_spinner(self.io.last_spinner_text) manager.clear_tag(MessageTag.DIFFS) @@ -2423,22 +2439,34 @@ async def check_tokens(self, messages): max_input_tokens = self.get_active_model().info.get("max_input_tokens") or 0 if max_input_tokens and input_tokens >= max_input_tokens: - self.io.tool_error( - f"Your estimated chat context of {input_tokens:,} tokens exceeds the" - f" {max_input_tokens:,} token limit for {self.get_active_model().name}!" - ) - self.io.tool_output("To reduce the chat context:") - self.io.tool_output("- Use /drop to remove unneeded files from the chat") - self.io.tool_output("- Use /clear to clear the chat history") - self.io.tool_output("- Break your code into smaller files") - self.io.tool_output( - "It's probably safe to try and send the request, most providers won't charge if" - " the context limit is exceeded." - ) + if self.enable_context_compaction: + self.io.tool_output( + f"Estimated chat context of {input_tokens:,} tokens exceeds the" + f" {max_input_tokens:,} token limit. Attempting to compact..." + ) + await self.compact_context_if_needed(force=True) - if not await self.io.confirm_ask("Try to proceed anyway?"): - return False - return True + # After compaction, re-format messages and re-check tokens + messages = self.format_messages() + input_tokens = self.get_active_model().token_count(messages) + + if max_input_tokens and input_tokens >= max_input_tokens: + self.io.tool_error( + f"Your estimated chat context of {input_tokens:,} tokens still exceeds the" + f" {max_input_tokens:,} token limit for {self.get_active_model().name}!" + ) + self.io.tool_output("To reduce the chat context:") + self.io.tool_output("- Use /drop to remove unneeded files from the chat") + self.io.tool_output("- Use /clear to clear the chat history") + self.io.tool_output("- Break your code into smaller files") + self.io.tool_output( + "It's probably safe to try and send the request, most providers won't charge if" + " the context limit is exceeded." + ) + + if not await self.io.confirm_ask("Try to proceed anyway?"): + return None + return messages def get_active_model(self): return self.main_model @@ -2499,7 +2527,8 @@ async def format_in_executor(): messages = result - if not await self.check_tokens(messages): + messages = await self.check_tokens(messages) + if not messages: return if self.verbose: @@ -3458,6 +3487,7 @@ async def send(self, messages, model=None, functions=None, tools=None): tools=tools, override_kwargs=self.model_kwargs.copy(), interrupt_event=self.interrupt_event, + coder=self, ) try: diff --git a/cecli/exceptions.py b/cecli/exceptions.py index df067f684e7..8ff8d6d9522 100644 --- a/cecli/exceptions.py +++ b/cecli/exceptions.py @@ -81,8 +81,12 @@ def _load(self, strict=False): def exceptions_tuple(self): return tuple(self.exceptions) - def get_ex_info(self, ex): - """Return the ExInfo for a given exception instance""" + def get_ex_info(self, ex, enable_context_compaction=False): + """Return the ExInfo for a given exception instance. + + When enable_context_compaction is True, ContextWindowExceededError + is treated as retryable so the caller can compact and retry. + """ import litellm if ex.__class__ is litellm.APIConnectionError: @@ -102,6 +106,16 @@ def get_ex_info(self, ex): ), ) + # ContextWindowExceededError: retryable only when compaction is enabled + if ex.__class__ is litellm.ContextWindowExceededError: + if enable_context_compaction: + return ExInfo( + "ContextWindowExceededError", + True, + "The context window has been exceeded. Compacting and retrying.", + ) + return ExInfo("ContextWindowExceededError", False, None) + # Check for specific non-retryable APIError cases like insufficient credits if ex.__class__ is litellm.APIError: err_str = str(ex).lower() diff --git a/cecli/main.py b/cecli/main.py index aea0c6b8690..0b2d1199b4b 100644 --- a/cecli/main.py +++ b/cecli/main.py @@ -610,6 +610,20 @@ async def main_async( set_args_error_data(args) + # Security: In agent mode, --ignore-context-limit must not be settable via + # environment variable alone. Reject if env var is set but CLI flag is absent. + is_agent_mode = args.edit_format == "agent" or getattr(args, "agent", False) + if is_agent_mode and args.ignore_context_limit: + env_val = os.environ.get("CECLI_IGNORE_CONTEXT_LIMIT", "").lower() + if env_val in ("true", "1", "yes"): + import logging + + logging.getLogger("cecli").warning( + "--ignore-context-limit via env var rejected in agent mode." + " Use --ignore-context-limit CLI flag instead." + ) + args.ignore_context_limit = False + if len(unknown): print("Unknown Args: ", unknown) diff --git a/cecli/mcp/server.py b/cecli/mcp/server.py index bb92c1473dd..1026cd3b72e 100644 --- a/cecli/mcp/server.py +++ b/cecli/mcp/server.py @@ -91,7 +91,9 @@ async def connect(self): command = self.config["command"] - env = {**os.environ, **self.config["env"]} if self.config.get("env") else None + env = os.environ.copy() + if self.config.get("env"): + env.update(self.config["env"]) server_params = StdioServerParameters( command=command, diff --git a/cecli/models.py b/cecli/models.py index 3ded65eb151..4c0d4c84fa8 100644 --- a/cecli/models.py +++ b/cecli/models.py @@ -3,6 +3,7 @@ import hashlib import importlib.resources import json +import logging import math import os import platform @@ -1223,6 +1224,7 @@ async def send_completion( max_wait=2, override_kwargs={}, interrupt_event=None, + coder=None, ): if os.environ.get("CECLI_SANITY_CHECK_TURNS"): sanity_check_messages(messages) @@ -1386,7 +1388,128 @@ async def send_completion( return hash_object, res except litellm.ContextWindowExceededError as err: - raise err + # Check if automatic compaction+retry is enabled via frozen settings + compaction_enabled = ( + coder is not None + and hasattr(coder, "enable_context_compaction") + and coder.enable_context_compaction + ) + + if not compaction_enabled: + # Backward compatible: propagate immediately when disabled + raise err + + # Determine effective retry limit (agent mode caps at 2) + max_retries = getattr(coder, "max_compaction_retries", 3) + is_agent_mode = getattr(coder, "is_agent_mode", False) + effective_max = min(max_retries, 2) if is_agent_mode else max_retries + + # Track per-request compaction retries (separate from general retry_delay loop) + if not hasattr(self, "_compaction_retry_count"): + self._compaction_retry_count = 0 + + compaction_logger = logging.getLogger("cecli.compaction") + + if self._compaction_retry_count >= effective_max: + # Exhausted compaction retries — propagate with user guidance + compaction_logger.info( + json.dumps( + { + "event": "context_compaction_exhausted", + "retry_count": self._compaction_retry_count, + "effective_max": effective_max, + "trigger": "context_window_exceeded", + } + ) + ) + print( + f"Context compaction failed after {effective_max} attempt(s)." + " Please use /clear or /compact manually." + ) + self._compaction_retry_count = 0 # Reset for next request + raise err + + # Attempt compaction before retrying + retry_index = self._compaction_retry_count + 1 + compaction_logger.info( + json.dumps( + { + "event": "context_compaction_attempt", + "retry_index": retry_index, + "effective_max": effective_max, + "trigger": "context_window_exceeded", + } + ) + ) + print(f"Compacting context… retry" f" {retry_index}/{effective_max}") + + try: + await coder.compact_context_if_needed( + force=True, + message="Context window exceeded, compacting to retry.", + ) + except Exception as compaction_err: + # Compaction itself failed — do NOT count as retry attempt + # (no context reduction occurred, retrying would be futile) + error_type, _ = FrozenCompactionSettings._scrub_compaction_error(compaction_err) + compaction_logger.warning( + json.dumps( + { + "event": "context_compaction_failed", + "error_type": error_type, + "retry_index": retry_index, + "trigger": "context_window_exceeded", + } + ) + ) + if self.verbose: + compaction_logger.debug(f"Compaction exception details: {compaction_err!r}") + print( + f"Context compaction failed: {error_type}." + " Please use /clear or /compact manually." + ) + self._compaction_retry_count = 0 # Reset for next request + raise err from compaction_err + + # Check if compaction actually reduced context (floor guard) + if getattr(coder, "_compaction_floor_reached", False): + compaction_logger.info( + json.dumps( + { + "event": "context_compaction_floor_reached", + "retry_index": retry_index, + "trigger": "context_window_exceeded", + } + ) + ) + print( + "Context is already at minimum size, cannot compact further." + " Please use /clear manually." + ) + self._compaction_retry_count = 0 # Reset for next request + raise err + + # Compaction succeeded — increment counter and retry + self._compaction_retry_count += 1 + + # Apply exponential backoff before retry + retry_delay *= self.retry_backoff_factor + if retry_delay > self.retry_timeout: + retry_delay = self.retry_timeout + + print(f"Retrying in {retry_delay:.1f} seconds...") + if interrupt_event: + _res, interrupted = await coroutines.interruptible( + asyncio.sleep(retry_delay), interrupt_event + ) + if interrupted: + self._compaction_retry_count = 0 # Reset on interrupt + raise KeyboardInterrupt("Interrupted during compaction retry sleep") + else: + await asyncio.sleep(retry_delay) + + # Reset compaction counter on successful retry (handled after loop) + continue except litellm_ex.exceptions_tuple() as err: ex_info = litellm_ex.get_ex_info(err) should_retry = ex_info.retry @@ -1462,6 +1585,7 @@ async def simple_send_with_retries( tools=tools, max_tokens=max_tokens, override_kwargs=override_kwargs, + coder=coder, ) if ( not response @@ -1543,6 +1667,27 @@ def _log_request(self, model_call_dict): f.write(",\n") +@dataclass(frozen=True) +class FrozenCompactionSettings: + enable_context_compaction: bool + max_compaction_retries: int + context_compaction_max_tokens: int + context_compaction_summary_tokens: int + is_agent_mode: bool + + @staticmethod + def _scrub_compaction_error(ex): + """Extract only the error type name and sanitized description. + + Strips message bodies, file paths, and conversation content from + exception output. Returns a tuple of (error_type_name, sanitized_description). + """ + error_type = type(ex).__name__ + # Only surface the error type, never raw exception messages + # which may contain conversation content or file paths + return error_type, None + + def register_models(model_settings_fnames): files_loaded = [] for model_settings_fname in model_settings_fnames: diff --git a/cecli/website/docs/config/options.md b/cecli/website/docs/config/options.md index f462120ede7..6a68d7656ab 100644 --- a/cecli/website/docs/config/options.md +++ b/cecli/website/docs/config/options.md @@ -263,6 +263,12 @@ The target maximum number of tokens for the generated summary. (default: 4096) Default: 4096 Environment variable: `CECLI_CONTEXT_COMPACTION_SUMMARY_TOKENS` +### `--max-compaction-retries VALUE` +Maximum number of automatic compaction and retry attempts when a context window error occurs (default: 3, range: 1-10). +In agent mode, this is capped at 2 retries for safety. +Default: 3 +Environment variable: `CECLI_MAX_COMPACTION_RETRIES` + ## Cache settings: ### `--cache-prompts` diff --git a/requirements/requirements.in b/requirements/requirements.in index a6575c2168f..c1a74a8df58 100644 --- a/requirements/requirements.in +++ b/requirements/requirements.in @@ -54,5 +54,4 @@ importlib-metadata>=7.2.1 tomli>=2.3.0; python_version <= "3.10" tree-sitter==0.23.2; python_version < "3.10" tree-sitter>=0.25.1; python_version >= "3.10" -tree-sitter-language-pack<=0.13.0 -tree_sitter_languages==1.10.2; python_version <= "3.12" \ No newline at end of file +tree-sitter-language-pack<=0.13.0 \ No newline at end of file diff --git a/tests/basic/test_context_compaction.py b/tests/basic/test_context_compaction.py new file mode 100644 index 00000000000..83463c015f9 --- /dev/null +++ b/tests/basic/test_context_compaction.py @@ -0,0 +1,545 @@ +"""Unit tests for context compaction and retry logic (CLI-56). + +Tests UT-CTX-001 through UT-CTX-018 as defined in .cecli.plans.md Section 10. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from litellm import ContextWindowExceededError + +from cecli.models import FrozenCompactionSettings, Model + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_coder(): + """Create a mock coder with compaction-related attributes.""" + coder = MagicMock() + coder.enable_context_compaction = True + coder.compact_context_if_needed = AsyncMock() + coder.context_compaction_max_tokens = 100000 + coder.max_compaction_retries = 3 + coder.is_agent_mode = False + coder._compaction_floor_reached = False + return coder + + +@pytest.fixture +def mock_model(mock_coder): + """Fixture to create a Model instance with a mock coder.""" + model = Model(model="gpt-4") + model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=False, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + return model + + +# --------------------------------------------------------------------------- +# UT-CTX-001: Compaction disabled — error propagates immediately +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_001_compaction_disabled_no_retry(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-001: Verify ContextWindowExceededError is not caught and not retried + when enable_context_compaction is False. + """ + # Override fixture default to test disabled behavior + mock_coder.enable_context_compaction = False + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + mock_acompletion.assert_called_once() + mock_coder.compact_context_if_needed.assert_not_called() + + +# --------------------------------------------------------------------------- +# UT-CTX-002: Compaction fires on error and retries successfully +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_002_compaction_fires_on_error_and_retries( + mock_acompletion, mock_model, mock_coder +): + """ + UT-CTX-002: Verify that when compaction is enabled, a ContextWindowExceededError + triggers compaction and a successful retry. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + mock_acompletion.side_effect = [ + ContextWindowExceededError(message="Test error", model="gpt-4", llm_provider="openai"), + MagicMock(), # Successful response + ] + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + assert mock_acompletion.call_count == 2 + mock_coder.compact_context_if_needed.assert_called_once() + + +# --------------------------------------------------------------------------- +# UT-CTX-003: Retry exhaustion after max_compaction_retries +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_003_retry_exhaustion(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-003: Verify that after max_compaction_retries, the error propagates to the user. + """ + mock_coder.enable_context_compaction = True + mock_coder.max_compaction_retries = 2 + mock_coder.is_agent_mode = False + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=2, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + assert mock_acompletion.call_count == 3 # Initial call + 2 retries + assert mock_coder.compact_context_if_needed.call_count == 2 + + +# --------------------------------------------------------------------------- +# UT-CTX-004: Token floor guard — compaction returns False +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("asyncio.sleep", new_callable=AsyncMock) +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_004_token_floor_guard(mock_acompletion, mock_sleep, mock_model, mock_coder): + """ + UT-CTX-004: Verify that if compaction returns False (token floor reached), + the error propagates without further retries. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.enable_context_compaction = True + mock_coder.max_compaction_retries = 3 + mock_coder.is_agent_mode = False + mock_model.retry_backoff_factor = 2.0 + mock_coder.compact_context_if_needed.return_value = True + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify that asyncio.sleep was called (backoff delays) + assert ( + mock_sleep.call_count >= 2 + ), f"Expected at least 2 sleep calls for backoff, got {mock_sleep.call_count}" + + # Verify delays follow exponential pattern: 0.125, 0.25, 0.5 + sleep_delays = [call_args[0][0] for call_args in mock_sleep.call_args_list] + # Note: implementation multiplies by backoff factor BEFORE first sleep + # Initial 0.125 * 2.0 = 0.25 + assert sleep_delays[0] == pytest.approx(0.25, rel=0.1), f"First delay: {sleep_delays[0]}" + assert sleep_delays[1] == pytest.approx(0.5, rel=0.1), f"Second delay: {sleep_delays[1]}" + assert sleep_delays[2] == pytest.approx(1.0, rel=0.1), f"Third delay: {sleep_delays[2]}" + + +# --------------------------------------------------------------------------- +# UT-CTX-011: Partial tool output from failed call is discarded +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_011_partial_tool_output_discard(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-011: Verify that partial tool output from a failed call is discarded + before compaction and retry. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + # First call raises context error (simulating partial output scenario) + # Second call succeeds + mock_acompletion.side_effect = [ + ContextWindowExceededError(message="Test error", model="gpt-4", llm_provider="openai"), + MagicMock(), + ] + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify compaction was called (which rewrites conversation state) + mock_coder.compact_context_if_needed.assert_called_once() + # Verify the second acompletion call happened (retry after compaction) + assert mock_acompletion.call_count == 2 + + +# --------------------------------------------------------------------------- +# UT-CTX-012: Non-idempotent tool safety semantic respected +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_012_non_idempotent_tool_safety(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-012: Verify that compaction does NOT override a tool's own + non-retryable error decision. + """ + from litellm import AuthenticationError + + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + + # AuthenticationError is non-retryable — compaction should NOT be triggered + mock_acompletion.side_effect = AuthenticationError( + message="Invalid API key", llm_provider="openai", model="gpt-4" + ) + + # send_completion handles non-retryable errors by returning model_error_response + _, res = await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + # Check that we got the error response + assert "Model API Response Error" in res.choices[0].message.content + + # Compaction should NOT be called for non-context errors + mock_coder.compact_context_if_needed.assert_not_called() + + +# --------------------------------------------------------------------------- +# UT-CTX-013: Agent mode retry cap (2) vs interactive mode (3) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_013_agent_mode_retry_cap(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-013: Verify that agent mode caps retries at 2, even if configured higher. + """ + mock_coder.enable_context_compaction = True + mock_coder.max_compaction_retries = 5 + mock_coder.is_agent_mode = True + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=5, # Configured higher than agent cap + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=True, + ) + mock_coder.compact_context_if_needed.return_value = True + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + assert mock_acompletion.call_count == 3 # Initial call + 2 retries (agent cap) + assert mock_coder.compact_context_if_needed.call_count == 2 + + +# --------------------------------------------------------------------------- +# UT-CTX-014: Status message format verification +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("builtins.print") +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_014_status_message_format( + mock_acompletion, mock_print, mock_model, mock_coder +): + """ + UT-CTX-014: Verify the format of the user-facing status message during compaction. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + mock_acompletion.side_effect = [ + ContextWindowExceededError(message="Test error", model="gpt-4", llm_provider="openai"), + ContextWindowExceededError(message="Test error", model="gpt-4", llm_provider="openai"), + MagicMock(), + ] + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify that status messages contain "Compacting context" and retry counts + printed_messages = [ + str(call_args[0][0]) for call_args in mock_print.call_args_list if call_args[0] + ] + compacting_messages = [m for m in printed_messages if "compacting" in m.lower()] + assert ( + len(compacting_messages) >= 2 + ), f"Expected at least 2 'Compacting context' messages, got: {compacting_messages}" + assert any( + "retry 1/3" in m.lower() or "retry 1 / 3" in m.lower() for m in compacting_messages + ), f"Expected 'retry 1/3' in messages: {compacting_messages}" + assert any( + "retry 2/3" in m.lower() or "retry 2 / 3" in m.lower() for m in compacting_messages + ), f"Expected 'retry 2/3' in messages: {compacting_messages}" + + +# --------------------------------------------------------------------------- +# UT-CTX-015: Regression — non-context retry behavior preserved +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_015_non_context_retry_preserved(mock_acompletion, mock_model): + """ + UT-CTX-015: Regression test to ensure non-context-related API errors are retried + as before, even if compaction is enabled. + """ + from litellm import APIError + + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + + mock_acompletion.side_effect = [ + APIError( + message="Test API error", + llm_provider="openai", + model="gpt-4", + status_code=500, + ), + MagicMock(), # Succeed on second attempt + ] + + mock_acompletion.side_effect = [ + APIError( + message="Test API error", + llm_provider="openai", + model="gpt-4", + status_code=500, + ), + MagicMock(), # Succeed on second attempt + ] + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + ) + + assert mock_acompletion.call_count == 2 + + +# --------------------------------------------------------------------------- +# UT-CTX-016: Regression — ContextWindowExceededError with compaction disabled +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_016_compaction_disabled_re_raises(mock_acompletion, mock_model): + """ + UT-CTX-016: Regression test to ensure ContextWindowExceededError is re-raised + when context compaction is disabled. + """ + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + ) + + mock_acompletion.assert_called_once() + + +# --------------------------------------------------------------------------- +# UT-CTX-017: Concurrent compaction safety +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_017_concurrent_compaction_safety(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-017: Verify that concurrent compaction attempts do not corrupt context. + + Two simultaneous send_completion calls both hitting context limit should + serialize compaction and maintain consistent state. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + # Both calls fail first, then succeed + mock_acompletion.side_effect = [ + ContextWindowExceededError(message="Test error 1", model="gpt-4", llm_provider="openai"), + MagicMock(), # Call 1 succeeds on retry + ContextWindowExceededError(message="Test error 2", model="gpt-4", llm_provider="openai"), + MagicMock(), # Call 2 succeeds on retry + ] + + # Run two concurrent send_completion calls + task1 = mock_model.send_completion( + messages=[{"role": "user", "content": "Hello 1"}], + functions=None, + stream=False, + coder=mock_coder, + ) + task2 = mock_model.send_completion( + messages=[{"role": "user", "content": "Hello 2"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + results = await asyncio.gather(task1, task2) + + # Both should complete successfully + assert results[0] is not None + assert results[1] is not None + # Compaction should have been called (at least once per call) + assert mock_coder.compact_context_if_needed.call_count >= 2 + + +# --------------------------------------------------------------------------- +# UT-CTX-018: Compaction with empty/near-empty context +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_ut_ctx_018_compaction_empty_context(mock_acompletion, mock_model, mock_coder): + """ + UT-CTX-018: Verify that compaction is skipped when context is already near-empty + (floor guard prevents division-by-zero and unnecessary compaction). + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + # Simulate floor already reached + mock_coder._compaction_floor_reached = True + mock_coder.compact_context_if_needed.return_value = False + + mock_acompletion.side_effect = ContextWindowExceededError( + message="Test error", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Only the initial call should happen — no retry after floor guard + mock_acompletion.assert_called_once() + # Compaction is attempted once but fails due to floor guard + mock_coder.compact_context_if_needed.assert_called_once() diff --git a/tests/integration/test_context_compaction_integration.py b/tests/integration/test_context_compaction_integration.py new file mode 100644 index 00000000000..d4e5c385822 --- /dev/null +++ b/tests/integration/test_context_compaction_integration.py @@ -0,0 +1,319 @@ +"""Integration tests for context compaction and retry logic (CLI-56). + +Tests IT-CTX-001 through IT-CTX-007 as defined in .cecli.plans.md Section 10. +""" + +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from litellm import ContextWindowExceededError + +from cecli.models import FrozenCompactionSettings, Model + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_coder(): + """Create a mock coder with compaction-related attributes.""" + coder = MagicMock() + coder.enable_context_compaction = True + coder.compact_context_if_needed = AsyncMock() + coder.context_compaction_max_tokens = 100000 + coder.max_compaction_retries = 3 + coder.is_agent_mode = False + coder._compaction_floor_reached = False + return coder + + +@pytest.fixture +def mock_model(mock_coder): + """Fixture to create a Model instance with a mock coder.""" + model = Model(model="gpt-4") + model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=False, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + return model + + +# --------------------------------------------------------------------------- +# IT-CTX-001: Full flow — context overflow triggers compaction and recovery +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_001_full_flow_compaction_recovery(mock_acompletion, mock_model, mock_coder): + """ + IT-CTX-001: Full flow with a context window overflow scenario and + enable-context-compaction=True — tool call recovers after compaction + without user intervention. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=4096, # Small limit to trigger overflow + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + # First call fails with context error, second succeeds after compaction + mock_acompletion.side_effect = [ + ContextWindowExceededError( + message="Context window exceeded", model="gpt-4", llm_provider="openai" + ), + MagicMock(), # Successful response after compaction + ] + + result = await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify compaction was called and retry succeeded + mock_coder.compact_context_if_needed.assert_called_once() + assert mock_acompletion.call_count == 2 + assert result is not None + + +# --------------------------------------------------------------------------- +# IT-CTX-002: Retry exhaustion — user receives clear failure message +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("builtins.print") +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_002_retry_exhaustion_user_message( + mock_acompletion, mock_print, mock_model, mock_coder +): + """ + IT-CTX-002: Verify that after 3 consecutive context window errors in a + single tool call, user receives a clear failure message with guidance. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + mock_coder.compact_context_if_needed.return_value = True + + # Always raises context error — never succeeds + mock_acompletion.side_effect = ContextWindowExceededError( + message="Context window exceeded", model="gpt-4", llm_provider="openai" + ) + + with pytest.raises(ContextWindowExceededError): + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify compaction was attempted 3 times + assert mock_coder.compact_context_if_needed.call_count == 3 + # Verify acompletion was called 4 times (initial + 3 retries) + assert mock_acompletion.call_count == 4 + + # Verify user-facing guidance message was printed + printed_messages = [ + str(call_args[0][0]) for call_args in mock_print.call_args_list if call_args[0] + ] + assert any( + "compaction failed" in msg.lower() or "clear" in msg.lower() or "compact" in msg.lower() + for msg in printed_messages + ), f"Expected failure guidance in output, got: {printed_messages}" + + +# --------------------------------------------------------------------------- +# IT-CTX-003: Agent-mode scenario — compaction events surfaced at session end +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("builtins.print") +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_003_agent_mode_compaction_summary( + mock_acompletion, mock_print, mock_model, mock_coder +): + """ + IT-CTX-003: Agent-mode scenario with multi-turn context — confirm + compaction events are tracked and surfaced. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=5, # Configured higher than agent cap + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=True, + ) + mock_coder.compact_context_if_needed.return_value = True + + # First call fails, second succeeds (one compaction event) + mock_acompletion.side_effect = [ + ContextWindowExceededError( + message="Context window exceeded", model="gpt-4", llm_provider="openai" + ), + MagicMock(), + ] + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Verify compaction was called (agent mode allows it) + mock_coder.compact_context_if_needed.assert_called_once() + assert mock_acompletion.call_count == 2 + + # Verify status messages were printed during compaction + printed_messages = [ + str(call_args[0][0]) for call_args in mock_print.call_args_list if call_args[0] + ] + compacting_messages = [m for m in printed_messages if "compacting" in m.lower()] + assert ( + len(compacting_messages) >= 1 + ), f"Expected at least 1 'Compacting context' message, got: {compacting_messages}" + + +# --------------------------------------------------------------------------- +# IT-CTX-004: Cross-component config flow — CLI args → args.py → coder → model +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_004_config_flow(mock_acompletion): + """ + IT-CTX-004: Test that config flows from CLI args through all layers. + """ + args = [ + "cecli", + "--enable-context-compaction", + "--max-compaction-retries", + "5", + "--yes", + ] + + with patch("sys.argv", args): + from cecli.args import get_parser + + parser = get_parser([], None) + parsed_args = parser.parse_args() + + assert parsed_args.enable_context_compaction is True + assert parsed_args.max_compaction_retries == 5 + + +# --------------------------------------------------------------------------- +# IT-CTX-005: YAML config file integration +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_005_yaml_config_integration(mock_acompletion): + """ + IT-CTX-005: Verify that YAML config values for compaction settings + are parsed correctly by the argument parser. + """ + # Simulate YAML config via environment variable (configargparse reads YAML) + # The CECLI_ prefix maps to enable-context-compaction + with patch.dict( + os.environ, + { + "CECLI_ENABLE_CONTEXT_COMPACTION": "true", + "CECLI_MAX_COMPACTION_RETRIES": "2", + }, + ): + from cecli.args import get_parser + + parser = get_parser([], None) + parsed_args = parser.parse_args([]) + + # Env vars should be picked up by configargparse + assert parsed_args.enable_context_compaction is True + assert parsed_args.max_compaction_retries == 2 + + +# --------------------------------------------------------------------------- +# IT-CTX-006: Env var + CLI flag precedence — CLI flag overrides env var +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch.dict( + os.environ, {"CECLI_ENABLE_CONTEXT_COMPACTION": "true", "CECLI_MAX_COMPACTION_RETRIES": "2"} +) +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_006_env_var_precedence(mock_acompletion): + """ + IT-CTX-006: Test that CLI flags override environment variables. + """ + args = [ + "cecli", + "--no-enable-context-compaction", + "--max-compaction-retries", + "1", + "--yes", + ] + + with patch("sys.argv", args): + from cecli.args import get_parser + + parser = get_parser([], None) + parsed_args = parser.parse_args() + + # CLI flag should override env var + assert parsed_args.enable_context_compaction is False + assert parsed_args.max_compaction_retries == 1 + + +# --------------------------------------------------------------------------- +# IT-CTX-007: Real model integration smoke test — no unnecessary compaction +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@patch("cecli.models.litellm.acompletion") +async def test_it_ctx_007_no_unnecessary_compaction(mock_acompletion, mock_model, mock_coder): + """ + IT-CTX-007: Verify that compaction is NOT triggered during normal + conversation flow when context limit is not exceeded. + """ + mock_model.compaction_settings = FrozenCompactionSettings( + enable_context_compaction=True, + max_compaction_retries=3, + context_compaction_max_tokens=10000, + context_compaction_summary_tokens=4096, + is_agent_mode=False, + ) + + # Normal successful response — no context error + mock_acompletion.return_value = MagicMock() + + await mock_model.send_completion( + messages=[{"role": "user", "content": "Hello"}], + functions=None, + stream=False, + coder=mock_coder, + ) + + # Compaction should NOT be called when no context error occurs + mock_coder.compact_context_if_needed.assert_not_called() + assert mock_acompletion.call_count == 1