diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 6c7f8590a5f24..0603594c63487 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -3072,8 +3072,8 @@ def run_conversation( if finish_reason == "length": if getattr(response, "id", "") == PARTIAL_STREAM_STUB_ID: agent._vprint( - f"{agent.log_prefix}⚠️ Stream interrupted by network error " - f"(finish_reason='length' on partial-stream-stub)", + f"{agent.log_prefix}⚠️ Response truncated — stream " + f"ended before completion", force=True, ) else: @@ -3200,6 +3200,11 @@ def run_conversation( # gets a coherent continuation point. if truncated_response_parts: messages = agent._get_messages_up_to_last_assistant(messages) + # Unmark survivors: their text left the stitched partial. + for _frag in messages: + if isinstance(_frag, dict): + _frag.pop("_length_continuation_fragment", None) + _frag.pop("_length_continuation_nudge", None) agent._session_messages = messages length_continue_retries = 0 truncated_response_parts = [] @@ -3236,6 +3241,8 @@ def run_conversation( ) if not _is_empty_partial_stub: interim_msg = agent._build_assistant_message(assistant_message, finish_reason) + # Marked so the ceiling exit can drop the fragment trail. + interim_msg["_length_continuation_fragment"] = True messages.append(interim_msg) if assistant_message.content: truncated_response_parts.append(assistant_message.content) @@ -3274,6 +3281,7 @@ def run_conversation( continue_msg = { "role": "user", "content": _continue_content, + "_length_continuation_nudge": True, } messages.append(continue_msg) agent._session_messages = messages @@ -3281,6 +3289,37 @@ def run_conversation( break partial_response = agent._strip_think_blocks("".join(truncated_response_parts)).strip() + if partial_response: + agent._vprint( + f"{agent.log_prefix}⚠️ Response still truncated " + f"after 4 continuation attempts — keeping the " + f"partial response received so far.", + force=True, + ) + # Unanswered continue nudges made every later turn re-truncate. + _turn_start = ( + current_turn_user_idx + 1 + if isinstance(current_turn_user_idx, int) + and current_turn_user_idx >= 0 + else 0 + ) + messages[_turn_start:] = [ + m for m in messages[_turn_start:] + if not ( + isinstance(m, dict) + and ( + m.get("_length_continuation_fragment") + or m.get("_length_continuation_nudge") + ) + ) + ] + if partial_response: + messages.append({ + "role": "assistant", + "content": partial_response, + "finish_reason": "length", + }) + agent._session_messages = messages agent._cleanup_task_resources(effective_task_id) agent._persist_session(messages, conversation_history) return { @@ -7160,6 +7199,11 @@ def run_conversation( final_response = "".join(truncated_response_parts) + final_response truncated_response_parts = [] length_continue_retries = 0 + # The continuation recovered, so the fragments stay in the transcript. + for _frag in messages: + if isinstance(_frag, dict): + _frag.pop("_length_continuation_fragment", None) + _frag.pop("_length_continuation_nudge", None) final_response = agent._strip_think_blocks(final_response).strip() diff --git a/tests/run_agent/test_continuation_ceiling_wedge.py b/tests/run_agent/test_continuation_ceiling_wedge.py new file mode 100644 index 0000000000000..bba451b1e38aa --- /dev/null +++ b/tests/run_agent/test_continuation_ceiling_wedge.py @@ -0,0 +1,168 @@ +"""Regression tests for the post-ceiling session wedge. + +A turn that exhausts all 4 length-continuation attempts must leave the +session usable: the next user message issues a fresh upstream request, +inherits no continuation counter, and the partial text that WAS received +is surfaced instead of dropped. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from hermes_constants import PARTIAL_STREAM_STUB_ID, FINISH_REASON_LENGTH + + +@pytest.fixture() +def loop_agent(): + from run_agent import AIAgent + with ( + patch("run_agent.get_tool_definitions", return_value=[]), + patch("run_agent.check_toolset_requirements", return_value={}), + patch("run_agent.OpenAI"), + ): + a = AIAgent( + api_key="test-key-1234567890", + base_url="https://openrouter.ai/api/v1", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + a.client = MagicMock() + a._cached_system_prompt = "You are helpful." + a._use_prompt_caching = False + a.compression_enabled = False + a.save_trajectories = False + return a + + +def _stub(content): + from tests.run_agent.test_run_agent import _mock_assistant_msg + return SimpleNamespace( + id=PARTIAL_STREAM_STUB_ID, + model="test/model", + choices=[SimpleNamespace( + index=0, + message=_mock_assistant_msg(content=content), + finish_reason=FINISH_REASON_LENGTH, + )], + usage=None, + ) + + +def _run(agent, message, history=None): + with ( + patch.object(agent, "_persist_session"), + patch.object(agent, "_save_trajectory"), + patch.object(agent, "_cleanup_task_resources"), + ): + return agent.run_conversation(message, conversation_history=history) + + +class TestContinuationCeilingWedge: + def _exhaust_ceiling(self, agent): + agent.client.chat.completions.create.side_effect = [ + _stub("part one "), _stub("part two "), + _stub("part three "), _stub("part four."), + ] + return _run(agent, "write me a long report") + + def test_partial_text_surfaced_at_ceiling(self, loop_agent): + result = self._exhaust_ceiling(loop_agent) + assert result["completed"] is False + assert result["partial"] is True + assert "part one" in (result["final_response"] or "") + assert "part four" in (result["final_response"] or "") + + def test_new_user_message_issues_fresh_request(self, loop_agent): + """Core regression: after the ceiling, a new user turn must reach + the provider instead of replaying wedge state.""" + from tests.run_agent.test_run_agent import _mock_response + + result1 = self._exhaust_ceiling(loop_agent) + assert "truncated after 4 continuation attempts" in (result1.get("error") or "") + calls_after_turn1 = loop_agent.client.chat.completions.create.call_count + assert calls_after_turn1 == 4 + + loop_agent.client.chat.completions.create.side_effect = [ + _mock_response(content="Hello! How can I help?", finish_reason="stop"), + ] + result2 = _run(loop_agent, "hi", history=result1["messages"]) + + assert loop_agent.client.chat.completions.create.call_count == calls_after_turn1 + 1, ( + "A new user message after the continuation ceiling must issue " + "exactly one fresh upstream request." + ) + assert result2["completed"] is True + assert result2["final_response"] == "Hello! How can I help?" + assert not result2.get("error") + + def test_ceiling_replaces_scaffolding_with_settled_turn(self, loop_agent): + """The persisted tail must not keep the continuation scaffolding. + Unanswered "continue" nudges make every later turn resume the + truncated response and re-exhaust the same ceiling.""" + result = self._exhaust_ceiling(loop_agent) + msgs = result["messages"] + + nudges = [ + m for m in msgs + if m.get("role") == "user" + and "Continue exactly where you left off" in (m.get("content") or "") + ] + assert nudges == [], ( + "Continuation nudges must not survive the ceiling exit — they " + "steer every subsequent turn back into the truncated response." + ) + + assistants = [m for m in msgs if m.get("role") == "assistant"] + assert len(assistants) == 1, ( + "The fragment trail must collapse into one settled assistant turn." + ) + assert msgs[-1]["role"] == "assistant" + content = msgs[-1]["content"] or "" + for part in ("part one", "part two", "part three", "part four"): + assert part in content, "Stitched partial must keep every fragment." + + def test_ceiling_not_labeled_network_error(self, loop_agent): + """A finish_reason='length' stub is a truncation, not a network + error — the user-facing message must not blame the network.""" + printed = [] + original = loop_agent._vprint + + def _capture(text, **kwargs): + printed.append(str(text)) + return original(text, **kwargs) + + with patch.object(loop_agent, "_vprint", side_effect=_capture): + self._exhaust_ceiling(loop_agent) + + network_lines = [line for line in printed if "network error" in line.lower()] + assert network_lines == [], ( + "Truncation must not be reported as a network error: " + f"{network_lines!r}" + ) + assert any("truncated" in line.lower() for line in printed), ( + "The user-facing message must name the truncation." + ) + + def test_new_turn_does_not_inherit_continuation_counter(self, loop_agent): + """A single truncation on the turn AFTER the ceiling must get its + own full 4-attempt budget, not the exhausted counter.""" + from tests.run_agent.test_run_agent import _mock_response + + result1 = self._exhaust_ceiling(loop_agent) + loop_agent.client.chat.completions.create.side_effect = [ + _stub("second turn partial "), + _mock_response(content="and the rest.", finish_reason="stop"), + ] + result2 = _run(loop_agent, "try again", history=result1["messages"]) + + assert result2["completed"] is True, ( + "One truncation on a fresh turn must continue (1/4), not fail " + "with an inherited exhausted counter." + ) + assert "second turn partial" in result2["final_response"] + assert "and the rest." in result2["final_response"]