test: add output-cap retry with compression disabled + fix request-pressure test
This commit is contained in:
parent
57f1814832
commit
7ef9345a55
2 changed files with 225 additions and 76 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue