test: add output-cap retry with compression disabled + fix request-pressure test

This commit is contained in:
dmabry 2026-07-13 11:22:55 -05:00 committed by kshitij
parent 57f1814832
commit 7ef9345a55
2 changed files with 225 additions and 76 deletions

View file

@ -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

View file

@ -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