diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index f15c611a58806..3e719882de04f 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -5074,6 +5074,142 @@ class TestRunConversation: assert result["final_response"] == "All done" assert result["completed"] is True + def test_engine_preflight_fires_below_threshold(self, agent): + """Sub-threshold ContextEngine.should_compress_preflight() routes to compress(). + + Regression test for #20316: when running below the threshold_tokens + cutoff, run_conversation must still consult the engine's + should_compress_preflight() hook so engines like hermes-lcm can + perform incremental maintenance (e.g. leaf-chunk compaction) + without waiting for the 75% context fill threshold. + """ + self._setup_agent(agent) + agent.compression_enabled = True + + # Build a conversation history long enough to clear the + # protect_first_n + protect_last_n + 1 guard so the preflight + # block actually executes. + protect_first = agent.context_compressor.protect_first_n + protect_last = agent.context_compressor.protect_last_n + prefill = [] + for _i in range((protect_first + protect_last + 4)): + prefill.append({"role": "user", "content": f"q{_i}"}) + prefill.append({"role": "assistant", "content": f"a{_i}"}) + + # Force the preflight estimator far below the threshold so the + # legacy ``>= threshold_tokens`` branch does NOT fire — only the + # new engine-driven elif branch should be exercised. + agent.context_compressor.threshold_tokens = 10**9 + + ok_resp = _mock_response(content="Done", finish_reason="stop") + agent.client.chat.completions.create.return_value = ok_resp + + # Engine-style hook: returns True so the elif branch should + # invoke _compress_context once for sub-threshold maintenance. + with ( + patch.object( + agent.context_compressor, + "should_compress_preflight", + return_value=True, + create=True, + ) as mock_preflight, + patch.object(agent, "_compress_context") as mock_compress, + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + mock_compress.return_value = ( + [{"role": "user", "content": "hello"}], + "compressed system prompt", + ) + result = agent.run_conversation("hello", conversation_history=prefill) + + mock_preflight.assert_called_once() + mock_compress.assert_called_once() + assert result["final_response"] == "Done" + assert result["completed"] is True + + def test_engine_preflight_skipped_when_returns_false(self, agent): + """should_compress_preflight() returning False must NOT invoke compress(). + + This guards the no-op default for the built-in ContextCompressor — + the elif branch should evaluate the hook but skip compression + when the engine reports nothing to do. + """ + self._setup_agent(agent) + agent.compression_enabled = True + + protect_first = agent.context_compressor.protect_first_n + protect_last = agent.context_compressor.protect_last_n + prefill = [] + for _i in range((protect_first + protect_last + 4)): + prefill.append({"role": "user", "content": f"q{_i}"}) + prefill.append({"role": "assistant", "content": f"a{_i}"}) + + agent.context_compressor.threshold_tokens = 10**9 + + ok_resp = _mock_response(content="Done", finish_reason="stop") + agent.client.chat.completions.create.return_value = ok_resp + + with ( + patch.object( + agent.context_compressor, + "should_compress_preflight", + return_value=False, + create=True, + ) as mock_preflight, + patch.object(agent, "_compress_context") as mock_compress, + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", conversation_history=prefill) + + mock_preflight.assert_called_once() + mock_compress.assert_not_called() + assert result["final_response"] == "Done" + assert result["completed"] is True + + def test_engine_preflight_exception_does_not_break_turn(self, agent): + """An exception in should_compress_preflight() must be swallowed. + + The plugin hook must never abort an otherwise-healthy turn — + a buggy engine raising in its preflight estimator should log + a debug message and continue without compressing. + """ + self._setup_agent(agent) + agent.compression_enabled = True + + protect_first = agent.context_compressor.protect_first_n + protect_last = agent.context_compressor.protect_last_n + prefill = [] + for _i in range((protect_first + protect_last + 4)): + prefill.append({"role": "user", "content": f"q{_i}"}) + prefill.append({"role": "assistant", "content": f"a{_i}"}) + + agent.context_compressor.threshold_tokens = 10**9 + + ok_resp = _mock_response(content="Done", finish_reason="stop") + agent.client.chat.completions.create.return_value = ok_resp + + with ( + patch.object( + agent.context_compressor, + "should_compress_preflight", + side_effect=RuntimeError("buggy engine"), + create=True, + ), + patch.object(agent, "_compress_context") as mock_compress, + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + result = agent.run_conversation("hello", conversation_history=prefill) + + mock_compress.assert_not_called() + assert result["final_response"] == "Done" + assert result["completed"] is True + def test_glm_prompt_exceeds_max_length_triggers_compression(self, agent): """GLM/Z.AI uses 'Prompt exceeds max length' for context overflow.""" self._setup_agent(agent)