fix(gateway): recover from truncated responses

This commit is contained in:
Dutch Dim 2026-06-17 22:30:44 +02:00 committed by kshitij
parent 9cf47fef54
commit 154c382d65
4 changed files with 59 additions and 16 deletions

View file

@ -1790,7 +1790,7 @@ def run_conversation(
if assistant_message.content:
truncated_response_parts.append(assistant_message.content)
if length_continue_retries < 3:
if length_continue_retries < 4:
_is_partial_stream_stub = (
getattr(response, "id", "") == PARTIAL_STREAM_STUB_ID
)
@ -1804,18 +1804,18 @@ def run_conversation(
f"{agent.log_prefix}↻ Stream interrupted mid "
f"tool-call ({_tool_list}) — requesting "
f"chunked retry "
f"({length_continue_retries}/3)..."
f"({length_continue_retries}/4)..."
)
elif _is_partial_stream_stub:
agent._vprint(
f"{agent.log_prefix}↻ Stream interrupted — "
f"requesting continuation "
f"({length_continue_retries}/3)..."
f"({length_continue_retries}/4)..."
)
else:
agent._vprint(
f"{agent.log_prefix}↻ Requesting continuation "
f"({length_continue_retries}/3)..."
f"({length_continue_retries}/4)..."
)
_continue_content = _get_continuation_prompt(
@ -1839,7 +1839,7 @@ def run_conversation(
"api_calls": api_call_count,
"completed": False,
"partial": True,
"error": "Response remained truncated after 3 continuation attempts",
"error": "Response remained truncated after 4 continuation attempts",
}
if agent.api_mode in {"chat_completions", "bedrock_converse", "anthropic_messages"}:
@ -1848,7 +1848,7 @@ def run_conversation(
_is_stub_stall = (
getattr(response, "id", "") == PARTIAL_STREAM_STUB_ID
)
if truncated_tool_call_retries < 3:
if truncated_tool_call_retries < 4:
truncated_tool_call_retries += 1
if _is_stub_stall:
# The stream broke mid tool-call (network /
@ -1856,13 +1856,13 @@ def run_conversation(
# cap — say so instead of "max output tokens".
agent._buffer_vprint(
f"⚠️ Stream interrupted mid tool-call — "
f"retrying ({truncated_tool_call_retries}/3)..."
f"retrying ({truncated_tool_call_retries}/4)..."
)
else:
agent._buffer_vprint(
f"⚠️ Truncated tool call detected — "
f"retrying API call "
f"({truncated_tool_call_retries}/3)..."
f"({truncated_tool_call_retries}/4)..."
)
# Boost max_tokens on each retry so the model has
# more room to complete the tool-call JSON. A
@ -1870,7 +1870,7 @@ def run_conversation(
# a genuine output-cap truncation does, and the
# boost is harmless for the stall case.
_tc_boost_base = agent.max_tokens if agent.max_tokens else 4096
_tc_boost = _tc_boost_base * (truncated_tool_call_retries + 1)
_tc_boost = _tc_boost_base * (2 ** truncated_tool_call_retries)
_tc_requested_cap = agent._requested_output_cap_from_api_kwargs(api_kwargs)
if _tc_requested_cap is not None:
_tc_boost = max(_tc_boost, _tc_requested_cap)
@ -1883,7 +1883,7 @@ def run_conversation(
agent._flush_status_buffer()
if _is_stub_stall:
agent._vprint(
f"{agent.log_prefix}⚠️ Stream kept dropping mid tool-call after 3 retries — the action was not executed.",
f"{agent.log_prefix}⚠️ Stream kept dropping mid tool-call after 4 retries — the action was not executed.",
force=True,
)
else:
@ -4016,13 +4016,14 @@ def run_conversation(
if _retry.restart_with_length_continuation:
# Progressively boost the output token budget on each retry.
# Retry 1 → 2× base, retry 2 → 3× base, capped at 32 768.
# Retry 1 → 2× base, retry 2 → 4× base, retry 3 → 8× base,
# retry 4 → 16× base, then cap at 32 768.
# Applies to all providers via _ephemeral_max_output_tokens.
# If the original request already used a larger provider/model
# default budget, keep that floor so continuation retries do
# not accidentally downshift to a much smaller cap.
_boost_base = agent.max_tokens if agent.max_tokens else 4096
_boost = _boost_base * (length_continue_retries + 1)
_boost = _boost_base * (2 ** length_continue_retries)
_requested_cap = agent._requested_output_cap_from_api_kwargs(api_kwargs)
if _requested_cap is not None:
_boost = max(_boost, _requested_cap)

View file

@ -17611,9 +17611,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
if not final_response:
error_msg = f"⚠️ {result['error']}" if result.get("error") else ""
final_response = _normalize_empty_agent_response(
result, final_response or "", history_len=len(agent_history),
)
final_response = _sanitize_gateway_final_response(source.platform, final_response)
if not final_response:
final_response = f"⚠️ {result['error']}" if result.get("error") else ""
return {
"final_response": error_msg,
"final_response": final_response,
"messages": result.get("messages", []),
"api_calls": result.get("api_calls", 0),
"failed": result.get("failed", False),

View file

@ -106,6 +106,29 @@ class InterruptedAgent:
return {"final_response": "interrupted", "messages": [], "api_calls": 1}
class PartialTruncationAgent:
"""Returns an incomplete turn with no visible assistant text."""
def __init__(self, **kwargs):
self.tool_progress_callback = kwargs.get("tool_progress_callback")
self.tools = []
self._interrupt_requested = False
@property
def is_interrupted(self) -> bool:
return self._interrupt_requested
def run_conversation(self, message, conversation_history=None, task_id=None):
return {
"final_response": None,
"messages": [],
"api_calls": 2,
"completed": False,
"partial": True,
"error": "Response truncated due to output length limit",
}
def _make_runner(adapter):
gateway_run = importlib.import_module("gateway.run")
GatewayRunner = gateway_run.GatewayRunner
@ -181,6 +204,20 @@ async def test_baseline_non_interrupted_agent_renders_progress(monkeypatch, tmp_
)
@pytest.mark.asyncio
async def test_partial_empty_agent_response_is_normalized(monkeypatch, tmp_path):
"""Messaging gateways should not echo raw truncation errors as final text."""
adapter, result = await _run_once(
monkeypatch, tmp_path, PartialTruncationAgent, "sess-partial-empty"
)
assert result["final_response"].startswith("⚠️ Processing stopped:")
assert "Response truncated due to output length limit" in result["final_response"]
assert result["final_response"] != "⚠️ Response truncated due to output length limit"
assert result["partial"] is True
assert adapter.sent == []
@pytest.mark.asyncio
async def test_progress_suppressed_when_agent_is_interrupted(monkeypatch, tmp_path):
"""Post-interrupt tool.started events must not render as bubbles.

View file

@ -5077,8 +5077,8 @@ class TestRunConversation:
result = agent.run_conversation("hello")
# Without think tags, the agent should attempt continuation retries
# (up to 3), not immediately fire thinking-exhaustion.
assert result["api_calls"] == 3
# (up to 4), not immediately fire thinking-exhaustion.
assert result["api_calls"] == 4
assert result["completed"] is False
def test_length_with_tool_calls_returns_partial_without_executing_tools(self, agent):