From b19039f9d09b7425958f0673d0f1877b0064c63f Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 23 Jul 2026 16:37:20 +0800 Subject: [PATCH] fix(agent): preserve length-recovered output --- nanobot/agent/runner.py | 32 +++++++++++++- tests/agent/test_runner_core.py | 64 +++++++++++++++++++++++++++ tests/agent/test_runner_hooks.py | 49 ++++++++++++++++++++ tests/agent/test_runner_injections.py | 36 +++++++++++++++ 4 files changed, 180 insertions(+), 1 deletion(-) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 9071a34a..83344763 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -60,6 +60,18 @@ _MAX_LENGTH_RECOVERIES = 3 _MAX_INJECTIONS_PER_TURN = 3 _MAX_INJECTION_CYCLES = 5 + +def _restore_outer_whitespace(content: str, original: str | None) -> str: + """Restore boundary whitespace stripped while cleaning one recovered segment.""" + if not original: + return content + leading_size = len(original) - len(original.lstrip()) + trailing_size = len(original) - len(original.rstrip()) + leading = original[:leading_size] + trailing = original[-trailing_size:] if trailing_size else "" + return f"{leading}{content}{trailing}" + + @dataclass(slots=True) class AgentRunSpec: """Configuration for a single agent execution.""" @@ -381,6 +393,9 @@ class AgentRunner: workspace_violation_counts: dict[str, int] = {} empty_content_retries = 0 length_recovery_count = 0 + # Segments from one uninterrupted length-recovery chain. Tool work or + # injected user input starts a new logical answer and clears the chain. + length_recovery_parts: list[str] = [] had_injections = False injection_cycles = 0 compacted_tool_call_ids: set[str] = set() @@ -418,6 +433,7 @@ class AgentRunner: context.response = response context.tool_calls = list(response.tool_calls) + original_content = response.content reasoning_text, cleaned_content = extract_reasoning( response.reasoning_content, response.thinking_blocks, @@ -519,6 +535,7 @@ class AgentRunner: ) empty_content_retries = 0 length_recovery_count = 0 + length_recovery_parts.clear() # Checkpoint 1: drain injections after tools, before next LLM call _drained, injection_cycles = await self._try_drain_injections( spec, messages, None, injection_cycles, @@ -567,11 +584,15 @@ class AgentRunner: context.response = response context.usage = dict(raw_usage) context.tool_calls = list(response.tool_calls) + original_content = response.content clean = hook.finalize_content(context, response.content) if response.finish_reason == "length" and not is_blank_text(clean): length_recovery_count += 1 if length_recovery_count <= _MAX_LENGTH_RECOVERIES: + length_recovery_parts.append( + _restore_outer_whitespace(clean, original_content) + ) logger.info( "Output truncated on turn {} for {} ({}/{}); continuing", iteration, @@ -614,6 +635,7 @@ class AgentRunner: await hook.on_stream_end(context, resuming=should_continue) if should_continue: + length_recovery_parts.clear() await hook.after_iteration(context) continue @@ -635,6 +657,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_parts.clear() continue break if is_blank_text(clean): @@ -652,6 +675,7 @@ class AgentRunner: ) if should_continue: had_injections = True + length_recovery_parts.clear() continue break @@ -671,7 +695,13 @@ class AgentRunner: "pending_tool_calls": [], }, ) - final_content = clean + if length_recovery_parts: + final_content = ( + "".join(length_recovery_parts) + + _restore_outer_whitespace(clean, original_content) + ).strip() + else: + final_content = clean context.final_content = final_content context.stop_reason = stop_reason await hook.after_iteration(context) diff --git a/tests/agent/test_runner_core.py b/tests/agent/test_runner_core.py index 39b74725..584df875 100644 --- a/tests/agent/test_runner_core.py +++ b/tests/agent/test_runner_core.py @@ -450,6 +450,70 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry(): assert result.stop_reason == "empty_final_response" +@pytest.mark.asyncio +async def test_runner_length_recovery_returns_all_segments(): + """Recovered output segments are returned together instead of only the tail.""" + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first ", finish_reason="length"), + LLMResponse(content="second ", finish_reason="length"), + LLMResponse(content="third", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "first second third" + assert [ + message["content"] + for message in result.messages + if message.get("role") == "assistant" + ] == ["first", "second", "third"] + assert provider.chat_with_retry.await_count == 3 + + +@pytest.mark.asyncio +async def test_runner_length_recovery_does_not_leak_across_tool_calls(): + """A recovered prefix belongs only to its contiguous response chain.""" + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="working", finish_reason="length"), + LLMResponse( + content=None, + tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})], + finish_reason="tool_calls", + ), + LLMResponse(content="final answer", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + tools.execute = AsyncMock(return_value="file content") + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "inspect a file"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "final answer" + assert result.tools_used == ["read_file"] + + @pytest.mark.asyncio async def test_runner_empty_response_does_not_break_tool_chain(): """An empty intermediate response must not kill an ongoing tool chain. diff --git a/tests/agent/test_runner_hooks.py b/tests/agent/test_runner_hooks.py index f441d285..0bc8a6fe 100644 --- a/tests/agent/test_runner_hooks.py +++ b/tests/agent/test_runner_hooks.py @@ -143,6 +143,55 @@ async def test_runner_streaming_hook_receives_deltas_and_end_signal(): provider.chat_with_retry.assert_not_awaited() +@pytest.mark.asyncio +async def test_runner_length_recovery_streams_segments_once_and_returns_all_content(): + from nanobot.agent.hook import AgentHook, AgentHookContext + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + streamed: list[str] = [] + endings: list[bool] = [] + responses = iter([ + LLMResponse(content="first ", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + ]) + + async def chat_stream_with_retry(*, on_content_delta, **kwargs): + response = next(responses) + await on_content_delta(response.content or "") + return response + + provider.chat_stream_with_retry = chat_stream_with_retry + provider.chat_with_retry = AsyncMock() + tools = MagicMock() + tools.get_definitions.return_value = [] + + class StreamingHook(AgentHook): + def wants_streaming(self) -> bool: + return True + + async def on_stream(self, context: AgentHookContext, delta: str) -> None: + streamed.append(delta) + + async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: + endings.append(resuming) + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=3, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + hook=StreamingHook(), + )) + + assert result.final_content == "first second" + assert streamed == ["first ", "second"] + assert endings == [True, False] + provider.chat_with_retry.assert_not_awaited() + + @pytest.mark.asyncio async def test_runner_passes_cached_tokens_to_hook_context(): """Hook context.usage should contain cached_tokens.""" diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index 9b96b4a8..cd7c76a4 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -352,6 +352,42 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream(): assert stream_end_calls[-1] is False +@pytest.mark.asyncio +async def test_injected_followup_starts_new_length_recovery_chain(): + """Recovered content from the prior answer must not prefix a follow-up reply.""" + from nanobot.agent.runner import AgentRunner + from nanobot.bus.events import InboundMessage + + provider = MagicMock() + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first", finish_reason="length"), + LLMResponse(content="second", finish_reason="stop"), + LLMResponse(content="follow-up answer", finish_reason="stop"), + ]) + tools = MagicMock() + tools.get_definitions.return_value = [] + + injection_queue = asyncio.Queue() + inject_cb = _make_injection_callback(injection_queue) + await injection_queue.put( + InboundMessage(channel="cli", sender_id="u", chat_id="c", content="follow-up question") + ) + + runner = AgentRunner() + result = await runner.run(make_run_spec(provider, + initial_messages=[{"role": "user", "content": "give a long answer"}], + tools=tools, + model="test-model", + max_iterations=5, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + injection_callback=inject_cb, + )) + + assert result.had_injections is True + assert result.final_content == "follow-up answer" + assert provider.chat_with_retry.await_count == 3 + + @pytest.mark.asyncio async def test_checkpoint2_preserves_final_response_in_history_before_followup(): """A follow-up injected after a final answer must still see that answer in history."""