fix(agent): preserve length-recovered output
This commit is contained in:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user