fix(agent): deliver non-streamed finalization responses
This commit is contained in:
+28
-4
@@ -149,6 +149,7 @@ class TurnContext:
|
||||
all_messages: list[dict[str, Any]] = field(default_factory=list)
|
||||
stop_reason: str = ""
|
||||
had_injections: bool = False
|
||||
streamed_content: bool = False
|
||||
|
||||
input_persisted_early: bool = False
|
||||
save_skip: int = 0
|
||||
@@ -1304,6 +1305,29 @@ class AgentLoop:
|
||||
hook_factories=list(hook_factories or []),
|
||||
tools=tools,
|
||||
)
|
||||
# A streaming callback may be present even when the final text comes from a
|
||||
# non-streaming recovery. Only the last completed segment can suppress the
|
||||
# regular outbound message.
|
||||
if ctx.on_stream is not None:
|
||||
stream_callback = ctx.on_stream
|
||||
stream_end_callback = ctx.on_stream_end
|
||||
segment_streamed_content = False
|
||||
|
||||
async def _tracked_stream(delta: str) -> None:
|
||||
nonlocal segment_streamed_content
|
||||
if delta:
|
||||
segment_streamed_content = True
|
||||
await stream_callback(delta)
|
||||
|
||||
async def _tracked_stream_end(*, resuming: bool = False) -> None:
|
||||
nonlocal segment_streamed_content
|
||||
ctx.streamed_content = segment_streamed_content
|
||||
segment_streamed_content = False
|
||||
if stream_end_callback is not None:
|
||||
await stream_end_callback(resuming=resuming)
|
||||
|
||||
ctx.on_stream = _tracked_stream
|
||||
ctx.on_stream_end = _tracked_stream_end
|
||||
|
||||
while ctx.state is not TurnState.DONE:
|
||||
handler_name = f"_state_{ctx.state.name.lower()}"
|
||||
@@ -1366,7 +1390,7 @@ class AgentLoop:
|
||||
all_msgs: list[dict[str, Any]],
|
||||
stop_reason: str,
|
||||
had_injections: bool,
|
||||
on_stream: Callable[[str], Awaitable[None]] | None,
|
||||
streamed_content: bool,
|
||||
*,
|
||||
turn_latency_ms: int | None = None,
|
||||
) -> OutboundMessage | None:
|
||||
@@ -1381,7 +1405,7 @@ class AgentLoop:
|
||||
|
||||
event = None
|
||||
meta = dict(msg.metadata or {})
|
||||
if on_stream is not None and stop_reason not in {"error", "tool_error"}:
|
||||
if streamed_content and stop_reason not in {"error", "tool_error"}:
|
||||
event = StreamedResponseEvent()
|
||||
if turn_latency_ms is not None:
|
||||
meta["latency_ms"] = int(turn_latency_ms)
|
||||
@@ -1626,7 +1650,7 @@ class AgentLoop:
|
||||
ctx.outbound = ctx.delivery.background_response(
|
||||
ctx.final_content,
|
||||
stop_reason=ctx.stop_reason,
|
||||
streamed=ctx.on_stream is not None,
|
||||
streamed=ctx.streamed_content,
|
||||
latency_ms=ctx.turn_latency_ms,
|
||||
)
|
||||
return "ok"
|
||||
@@ -1636,7 +1660,7 @@ class AgentLoop:
|
||||
ctx.all_messages,
|
||||
ctx.stop_reason,
|
||||
ctx.had_injections,
|
||||
ctx.on_stream,
|
||||
ctx.streamed_content,
|
||||
turn_latency_ms=ctx.turn_latency_ms,
|
||||
)
|
||||
if ctx.ephemeral and ctx.outbound is not None:
|
||||
|
||||
@@ -534,6 +534,55 @@ class TestToolEventProgress:
|
||||
assert turn_end_msgs[0].content == ""
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streamed_finalization_is_delivered_as_regular_message(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A no-tools finalization must not be dropped after empty stream retries."""
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.supports_progress_deltas = True
|
||||
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
|
||||
provider.chat_stream_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content=None, tool_calls=[]),
|
||||
LLMResponse(content=None, tool_calls=[]),
|
||||
])
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(content="final answer", tool_calls=[]),
|
||||
)
|
||||
loop = AgentLoop(
|
||||
bus=bus,
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="openai-codex/gpt-5.5",
|
||||
)
|
||||
_attach_webui_runtime_events(loop, bus)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="say hello",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
assert not any(isinstance(message.event, StreamDeltaEvent) for message in outbound)
|
||||
assert len([
|
||||
message for message in outbound if isinstance(message.event, StreamEndEvent)
|
||||
]) == 3
|
||||
final = [message for message in outbound if message.content == "final answer"]
|
||||
assert len(final) == 1
|
||||
assert final[0].event is None
|
||||
provider.chat_stream_with_retry.assert_awaited()
|
||||
provider.chat_with_retry.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_independent_late_subagent_result_gets_complete_webui_turn(
|
||||
self,
|
||||
|
||||
@@ -512,7 +512,7 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path):
|
||||
)],
|
||||
usage={},
|
||||
)
|
||||
provider.chat_stream_with_retry = AsyncMock(side_effect=[
|
||||
responses = iter([
|
||||
tool_call_resp,
|
||||
LLMResponse(
|
||||
content="I cannot access private URLs. Please share the local file.",
|
||||
@@ -521,6 +521,13 @@ async def test_ssrf_soft_block_can_finalize_after_streamed_tool_call(tmp_path):
|
||||
),
|
||||
])
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
response = next(responses)
|
||||
await on_content_delta(response.content)
|
||||
return response
|
||||
|
||||
provider.chat_stream_with_retry = AsyncMock(side_effect=chat_stream_with_retry)
|
||||
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.tools.prepare_call = MagicMock(return_value=(None, {}, None))
|
||||
|
||||
Reference in New Issue
Block a user