fix(agent): reset recovery chains across injections

Reset both the recovered segments and retry budget whenever injected input starts a new logical answer. Cover fatal tool-error boundaries and rename the prompt test module so pytest can collect the full suite.
This commit is contained in:
chengyongru
2026-07-27 01:39:46 +08:00
committed by Xubin Ren
parent 154cbc1974
commit 1d2ed6e4d2
3 changed files with 25 additions and 8 deletions
+5
View File
@@ -520,6 +520,8 @@ class AgentRunner:
) )
if should_continue: if should_continue:
had_injections = True had_injections = True
length_recovery_count = 0
length_recovery_parts.clear()
continue continue
break break
await self._emit_checkpoint( await self._emit_checkpoint(
@@ -635,6 +637,7 @@ class AgentRunner:
await hook.on_stream_end(context, resuming=should_continue) await hook.on_stream_end(context, resuming=should_continue)
if should_continue: if should_continue:
length_recovery_count = 0
length_recovery_parts.clear() length_recovery_parts.clear()
await hook.after_iteration(context) await hook.after_iteration(context)
continue continue
@@ -657,6 +660,7 @@ class AgentRunner:
) )
if should_continue: if should_continue:
had_injections = True had_injections = True
length_recovery_count = 0
length_recovery_parts.clear() length_recovery_parts.clear()
continue continue
break break
@@ -675,6 +679,7 @@ class AgentRunner:
) )
if should_continue: if should_continue:
had_injections = True had_injections = True
length_recovery_count = 0
length_recovery_parts.clear() length_recovery_parts.clear()
continue continue
break break
+18 -8
View File
@@ -354,15 +354,18 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_injected_followup_starts_new_length_recovery_chain(): async def test_injected_followup_starts_new_length_recovery_chain():
"""Recovered content from the prior answer must not prefix a follow-up reply.""" """A follow-up gets a fresh recovery budget and no content from the prior answer."""
from nanobot.agent.runner import AgentRunner from nanobot.agent.runner import AgentRunner
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
provider = MagicMock() provider = MagicMock()
provider.chat_with_retry = AsyncMock(side_effect=[ provider.chat_with_retry = AsyncMock(side_effect=[
LLMResponse(content="first", finish_reason="length"), LLMResponse(content="first-1 ", finish_reason="length"),
LLMResponse(content="second", finish_reason="stop"), LLMResponse(content="first-2 ", finish_reason="length"),
LLMResponse(content="follow-up answer", finish_reason="stop"), LLMResponse(content="first-3 ", finish_reason="length"),
LLMResponse(content="first-final", finish_reason="stop"),
LLMResponse(content="follow-up ", finish_reason="length"),
LLMResponse(content="answer", finish_reason="stop"),
]) ])
tools = MagicMock() tools = MagicMock()
tools.get_definitions.return_value = [] tools.get_definitions.return_value = []
@@ -378,14 +381,14 @@ async def test_injected_followup_starts_new_length_recovery_chain():
initial_messages=[{"role": "user", "content": "give a long answer"}], initial_messages=[{"role": "user", "content": "give a long answer"}],
tools=tools, tools=tools,
model="test-model", model="test-model",
max_iterations=5, max_iterations=8,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
injection_callback=inject_cb, injection_callback=inject_cb,
)) ))
assert result.had_injections is True assert result.had_injections is True
assert result.final_content == "follow-up answer" assert result.final_content == "follow-up answer"
assert provider.chat_with_retry.await_count == 3 assert provider.chat_with_retry.await_count == 6
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -1348,7 +1351,7 @@ async def test_dispatch_republishes_leftover_queue_messages(tmp_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_drain_injections_on_fatal_tool_error(): async def test_drain_injections_on_fatal_tool_error():
"""Pending injections should be drained even when a fatal tool error occurs.""" """A fatal tool error must not leak recovered content into an injected follow-up."""
from nanobot.agent.runner import AgentRunner from nanobot.agent.runner import AgentRunner
from nanobot.bus.events import InboundMessage from nanobot.bus.events import InboundMessage
@@ -1358,12 +1361,18 @@ async def test_drain_injections_on_fatal_tool_error():
async def chat_with_retry(*, messages, **kwargs): async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1 call_count["n"] += 1
if call_count["n"] == 1: if call_count["n"] == 1:
return LLMResponse(
content="stale prefix ",
finish_reason="length",
usage={},
)
if call_count["n"] == 2:
return LLMResponse( return LLMResponse(
content="", content="",
tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})], tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})],
usage={}, usage={},
) )
# Second call: respond normally to the injected follow-up # Third call: respond normally to the injected follow-up.
return LLMResponse(content="reply to follow-up", tool_calls=[], usage={}) return LLMResponse(content="reply to follow-up", tool_calls=[], usage={})
provider.chat_with_retry = chat_with_retry provider.chat_with_retry = chat_with_retry
@@ -1391,6 +1400,7 @@ async def test_drain_injections_on_fatal_tool_error():
assert result.had_injections is True assert result.had_injections is True
assert result.final_content == "reply to follow-up" assert result.final_content == "reply to follow-up"
assert call_count["n"] == 3
# The injection should be in the messages history # The injection should be in the messages history
injected = [ injected = [
m for m in result.messages m for m in result.messages
@@ -1,3 +1,5 @@
"""Tests for length-recovery prompt construction."""
from nanobot.utils.runtime import build_length_recovery_message from nanobot.utils.runtime import build_length_recovery_message