diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index 396657ea4..7c4ad7cb4 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -5355,6 +5355,231 @@ class TestRunConversation: "_record_task_failure should not be called outside kanban mode" ) + # ── Output-cap retry: safe_out uses provider available_out + request estimate ── + + def test_output_cap_retry_uses_provider_available_out(self, agent): + """run_conversation retries an output-cap error with max_tokens <= + available_out - 64, and does NOT halve context_length or trigger + compression. + """ + self._setup_agent(agent) + agent.api_mode = "chat_completions" + agent.provider = "openrouter" + agent.model = "some/model" + agent.max_tokens = 65_536 + agent.compression_enabled = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.should_compress = MagicMock(return_value=False) + + error_msg = ( + "max_tokens: 65536 > context_window: 200000 " + "- input_tokens: 199000 = available_tokens: 1000" + ) + exc = Exception(error_msg) + exc.status_code = 400 + exc.code = 400 + + ok_resp = _mock_response(content="done", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [exc, ok_resp] + + mock_compress = MagicMock() + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch.object(agent.context_compressor, "update_model"), + patch.object(agent, "_compress_context", mock_compress), + ): + result = agent.run_conversation("hello") + + second_call = agent.client.chat.completions.create.call_args_list[1].kwargs + assert result["completed"] is True + assert second_call["max_tokens"] <= 936 + assert agent.context_compressor.context_length == 200_000 + mock_compress.assert_not_called() + + def test_output_cap_retry_with_large_api_only_content(self, agent): + """When a large system prompt makes api_messages huge while persisted + messages stay tiny, the retry cap must still respect provider + available_tokens — not blow up to the full context window. + """ + self._setup_agent(agent) + agent.api_mode = "chat_completions" + agent.provider = "openrouter" + agent.model = "some/model" + agent.max_tokens = 65_536 + agent.compression_enabled = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.should_compress = MagicMock(return_value=False) + + # Huge API-only system prompt; persisted messages are tiny. + agent._cached_system_prompt = "S" * 796_000 + + error_msg = ( + "max_tokens: 65536 > context_window: 200000 " + "- input_tokens: 199000 = available_tokens: 1000" + ) + exc = Exception(error_msg) + exc.status_code = 400 + exc.code = 400 + + ok_resp = _mock_response(content="done", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [exc, ok_resp] + + mock_compress = MagicMock() + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch.object(agent.context_compressor, "update_model"), + patch.object(agent, "_compress_context", mock_compress), + ): + result = agent.run_conversation("hello") + + second_call = agent.client.chat.completions.create.call_args_list[1].kwargs + assert result["completed"] is True + # The current branch (messages-only estimate) would send max_tokens + # near 199927 — this test fails on it. + assert second_call["max_tokens"] <= 936 + assert agent.context_compressor.context_length == 200_000 + mock_compress.assert_not_called() + + def test_output_cap_retry_request_pressure_lower_bound(self, agent): + """When the provider reports a large available_tokens but local request + pressure leaves less room, the retry cap is the smaller of the two. + """ + self._setup_agent(agent) + agent.api_mode = "chat_completions" + agent.provider = "openrouter" + agent.model = "some/model" + agent.max_tokens = 65_536 + agent.compression_enabled = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.should_compress = MagicMock(return_value=False) + + # A large API-only system prompt so the local estimate is the binding + # constraint, not the provider's available_tokens. + agent._cached_system_prompt = "S" * 796_000 + + error_msg = ( + "max_tokens: 65536 > context_window: 200000 " + "- input_tokens: 190000 = available_tokens: 50000" + ) + exc = Exception(error_msg) + exc.status_code = 400 + exc.code = 400 + + ok_resp = _mock_response(content="done", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [exc, ok_resp] + + mock_compress = MagicMock() + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch.object(agent.context_compressor, "update_model"), + patch.object(agent, "_compress_context", mock_compress), + ): + result = agent.run_conversation("hello") + + first_call = agent.client.chat.completions.create.call_args_list[0].kwargs + second_call = agent.client.chat.completions.create.call_args_list[1].kwargs + assert result["completed"] is True + + # Verify the local estimate is actually the lower bound. + from agent.model_metadata import estimate_request_tokens_rough + estimated_request = estimate_request_tokens_rough( + first_call["messages"], tools=agent.tools or None, + ) + local_available = 200_000 - estimated_request + expected_cap = max(1, min(50_000, local_available) - 64) + assert local_available < 50_000 + assert second_call["max_tokens"] == expected_cap + assert agent.context_compressor.context_length == 200_000 + mock_compress.assert_not_called() + + def test_output_cap_retry_safety_floor_at_one(self, agent): + """When provider available_tokens is 1, the retry cap is floored at 1.""" + self._setup_agent(agent) + agent.api_mode = "chat_completions" + agent.provider = "openrouter" + agent.model = "some/model" + agent.max_tokens = 65_536 + agent.compression_enabled = True + agent.context_compressor.context_length = 200_000 + agent.context_compressor.should_compress = MagicMock(return_value=False) + + error_msg = ( + "max_tokens: 65536 > context_window: 200000 " + "- input_tokens: 199999 = available_tokens: 1" + ) + exc = Exception(error_msg) + exc.status_code = 400 + exc.code = 400 + + ok_resp = _mock_response(content="done", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [exc, ok_resp] + + mock_compress = MagicMock() + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch.object(agent.context_compressor, "update_model"), + patch.object(agent, "_compress_context", mock_compress), + ): + result = agent.run_conversation("hello") + + second_call = agent.client.chat.completions.create.call_args_list[1].kwargs + assert result["completed"] is True + assert second_call["max_tokens"] == 1 + assert agent.context_compressor.context_length == 200_000 + mock_compress.assert_not_called() + + def test_output_cap_retry_with_compression_disabled(self, agent): + """Output-cap retry must still work when compression.enabled is false. + The recovery is a max_tokens-only retry — it does not require compression, + so the compression-disabled guard must not block it. + """ + self._setup_agent(agent) + agent.api_mode = "chat_completions" + agent.provider = "openrouter" + agent.model = "some/model" + agent.max_tokens = 65_536 + agent.compression_enabled = False + agent.context_compressor.context_length = 200_000 + agent.context_compressor.should_compress = MagicMock(return_value=False) + + error_msg = ( + "max_tokens: 65536 > context_window: 200000 " + "- input_tokens: 199000 = available_tokens: 1000" + ) + exc = Exception(error_msg) + exc.status_code = 400 + exc.code = 400 + + ok_resp = _mock_response(content="done", finish_reason="stop") + agent.client.chat.completions.create.side_effect = [exc, ok_resp] + + mock_compress = MagicMock() + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + patch.object(agent.context_compressor, "update_model"), + patch.object(agent, "_compress_context", mock_compress), + ): + result = agent.run_conversation("hello") + + # Two API calls: the failed one and the retried one. + assert len(agent.client.chat.completions.create.call_args_list) == 2 + second_call = agent.client.chat.completions.create.call_args_list[1].kwargs + assert result["completed"] is True + assert result.get("compaction_disabled") is None + assert second_call["max_tokens"] <= 936 + assert agent.context_compressor.context_length == 200_000 + mock_compress.assert_not_called() + class TestHookPayloadSanitizesSimpleNamespace: """Regression: ``_hook_jsonable`` referenced ``SimpleNamespace`` without diff --git a/tests/test_ctx_halving_fix.py b/tests/test_ctx_halving_fix.py index 9482d160a..83b2d0e1d 100644 --- a/tests/test_ctx_halving_fix.py +++ b/tests/test_ctx_halving_fix.py @@ -362,80 +362,4 @@ class TestContextNotHalvedOnOutputCapError: assert safe_out == 1 -# --------------------------------------------------------------------------- -# Regression: safe_out must track growing input, not stale error budget -# --------------------------------------------------------------------------- -class TestSafeOutTracksGrowingInput: - """Regression test for #55546: safe_out derived from stale available_tokens - reuses a budget computed from the *previous* request. Between retries the - agent appends tool results and error text, so the real input token count - grows. The fix computes safe_out from the *current* message token estimate - so the cap stays valid on every retry. - """ - - def _make_agent(self, context_length=262_144): - from run_agent import AIAgent - from agent.context_compressor import ContextCompressor - - agent = object.__new__(AIAgent) - agent.api_mode = "chat_completions" - agent.model = "qwen3.6-27b" - agent.base_url = "http://localhost:1234/v1" - agent.tools = [] - agent.max_tokens = 65_536 - agent.reasoning_config = None - agent._ephemeral_max_output_tokens = None - - compressor = MagicMock(spec=ContextCompressor) - compressor.context_length = context_length - agent.context_compressor = compressor - - agent._prepare_messages_for_api = MagicMock( - return_value=[{"role": "user", "content": "hi"}] - ) - agent._vprint = MagicMock() - agent.request_overrides = {} - return agent - - def test_safe_out_uses_current_input_not_stale_available(self): - """safe_out is computed from the current message token estimate, not - the stale available_tokens from the error message. - """ - from agent.model_metadata import estimate_messages_tokens_rough - - agent = self._make_agent(context_length=262_144) - old_ctx = agent.context_compressor.context_length - - # Simulate a conversation that's near the ceiling. - messages = [{"role": "user", "content": "x" * 800_000}] - _current_input = estimate_messages_tokens_rough(messages) - - # The fix: derive safe_out from the current input estimate. - safe_out = max(1, old_ctx - _current_input - 64) - agent._ephemeral_max_output_tokens = safe_out - - # Verify: safe_out is based on current input, not a stale error value. - assert agent._ephemeral_max_output_tokens == safe_out - assert agent.context_compressor.context_length == old_ctx - - def test_safe_out_tracks_growing_input(self): - """When messages grow between retries, safe_out shrinks accordingly.""" - from agent.model_metadata import estimate_messages_tokens_rough - - agent = self._make_agent(context_length=262_144) - old_ctx = agent.context_compressor.context_length - - # Initial messages. - messages = [{"role": "user", "content": "x" * 800_000}] - input_1 = estimate_messages_tokens_rough(messages) - safe_out_1 = max(1, old_ctx - input_1 - 64) - - # Simulate the agent appending a tool result (input grows). - messages.append({"role": "assistant", "content": "tool result" * 200}) - input_2 = estimate_messages_tokens_rough(messages) - safe_out_2 = max(1, old_ctx - input_2 - 64) - - # safe_out_2 must be <= safe_out_1 (shrinks as input grows). - assert safe_out_2 <= safe_out_1 - assert safe_out_2 >= 1