diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index 0c32a8f1..d42a5821 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -149,6 +149,7 @@ class TestToolEventProgress: provider.chat_with_retry = AsyncMock() loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="openai-codex/gpt-5.5") 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", @@ -165,7 +166,8 @@ class TestToolEventProgress: final = [m for m in outbound if not m.metadata.get("_progress")] assert [m.content for m in progress] == ["Hel", "lo"] - assert final[-1].content == "Hello" + assert final[-2].content == "Hello" + assert (final[-1].metadata or {}).get("_turn_end") is True provider.chat_with_retry.assert_not_awaited() @pytest.mark.asyncio @@ -214,3 +216,30 @@ class TestToolEventProgress: 'custom_tool("foo.txt")', ] assert all(item[0] != "I will inspect it." for item in progress) + + @pytest.mark.asyncio + async def test_websocket_dispatch_publishes_final_turn_end_marker(self, tmp_path: Path) -> None: + bus = MessageBus() + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="Done", tool_calls=[])) + loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model") + 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", + )) + + outbound = [] + while bus.outbound_size > 0: + outbound.append(await bus.consume_outbound()) + + assert outbound[-2].content == "Done" + assert (outbound[-2].metadata or {}).get("_turn_end") is not True + assert outbound[-1].content == "" + assert (outbound[-1].metadata or {}).get("_turn_end") is True + assert outbound[-1].chat_id == "chat1" diff --git a/tests/channels/test_websocket_channel.py b/tests/channels/test_websocket_channel.py index b5dc830b..db61fc28 100644 --- a/tests/channels/test_websocket_channel.py +++ b/tests/channels/test_websocket_channel.py @@ -287,6 +287,25 @@ async def test_send_delta_emits_delta_and_stream_end() -> None: assert second["stream_id"] == "sid" +@pytest.mark.asyncio +async def test_send_turn_end_emits_turn_end_event() -> None: + bus = MagicMock() + channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus) + mock_ws = AsyncMock() + channel._attach(mock_ws, "chat-1") + + await channel.send(OutboundMessage( + channel="websocket", + chat_id="chat-1", + content="", + metadata={"_turn_end": True}, + )) + + mock_ws.send.assert_awaited_once() + body = json.loads(mock_ws.send.await_args.args[0]) + assert body == {"event": "turn_end", "chat_id": "chat-1"} + + @pytest.mark.asyncio async def test_send_non_connection_closed_exception_is_raised() -> None: bus = MagicMock() @@ -545,6 +564,16 @@ async def test_end_to_end_server_pushes_streaming_deltas_to_client(bus: MagicMoc end = json.loads(await client.recv()) assert end["event"] == "stream_end" assert end["stream_id"] == "s1" + + await channel.send(OutboundMessage( + channel="websocket", + chat_id=chat_id, + content="", + metadata={"_turn_end": True}, + )) + + turn_end = json.loads(await client.recv()) + assert turn_end == {"event": "turn_end", "chat_id": chat_id} finally: await channel.stop() await server_task diff --git a/webui/src/tests/useNanobotStream.test.tsx b/webui/src/tests/useNanobotStream.test.tsx index f5adcf17..2c717317 100644 --- a/webui/src/tests/useNanobotStream.test.tsx +++ b/webui/src/tests/useNanobotStream.test.tsx @@ -6,6 +6,8 @@ import { useNanobotStream } from "@/hooks/useNanobotStream"; import type { InboundEvent } from "@/lib/types"; import { ClientProvider } from "@/providers/ClientProvider"; +const EMPTY_MESSAGES: import("@/lib/types").UIMessage[] = []; + function fakeClient() { const handlers = new Map void>>(); return { @@ -51,9 +53,27 @@ function wrap(client: ReturnType["client"]) { } describe("useNanobotStream", () => { + it("starts in streaming mode when history shows pending tool calls", () => { + const fake = fakeClient(); + const initialMessages = [{ + id: "m1", + role: "assistant" as const, + content: "Using tools", + createdAt: Date.now(), + }]; + const { result } = renderHook( + () => useNanobotStream("chat-p", initialMessages, true), + { + wrapper: wrap(fake.client), + }, + ); + + expect(result.current.isStreaming).toBe(true); + }); + it("collapses consecutive tool_hint frames into one trace row", () => { const fake = fakeClient(); - const { result } = renderHook(() => useNanobotStream("chat-t", []), { + const { result } = renderHook(() => useNanobotStream("chat-t", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); @@ -95,7 +115,7 @@ describe("useNanobotStream", () => { it("attaches assistant media_urls to complete messages", () => { const fake = fakeClient(); - const { result } = renderHook(() => useNanobotStream("chat-m", []), { + const { result } = renderHook(() => useNanobotStream("chat-m", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); @@ -116,7 +136,7 @@ describe("useNanobotStream", () => { it("keeps assistant buttons on complete messages", () => { const fake = fakeClient(); - const { result } = renderHook(() => useNanobotStream("chat-q", []), { + const { result } = renderHook(() => useNanobotStream("chat-q", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); @@ -136,4 +156,60 @@ describe("useNanobotStream", () => { ["Short answer", "Detailed answer"], ]); }); + + it("keeps streaming alive across stream_end and completes on turn_end", () => { + const fake = fakeClient(); + const { result } = renderHook(() => useNanobotStream("chat-s", EMPTY_MESSAGES), { + wrapper: wrap(fake.client), + }); + + act(() => { + fake.emit("chat-s", { + event: "delta", + chat_id: "chat-s", + text: "Hello", + }); + }); + + expect(result.current.isStreaming).toBe(true); + expect(result.current.messages[0]).toMatchObject({ + role: "assistant", + content: "Hello", + isStreaming: true, + }); + + act(() => { + fake.emit("chat-s", { + event: "stream_end", + chat_id: "chat-s", + }); + }); + + expect(result.current.isStreaming).toBe(true); + expect(result.current.messages[0].isStreaming).toBe(true); + + act(() => { + fake.emit("chat-s", { + event: "message", + chat_id: "chat-s", + text: "Hello world", + }); + }); + + expect(result.current.isStreaming).toBe(true); + expect(result.current.messages.at(-1)).toMatchObject({ + role: "assistant", + content: "Hello world", + }); + + act(() => { + fake.emit("chat-s", { + event: "turn_end", + chat_id: "chat-s", + }); + }); + + expect(result.current.isStreaming).toBe(false); + expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true); + }); }); diff --git a/webui/src/tests/useSessions.test.tsx b/webui/src/tests/useSessions.test.tsx index ad4f1c1a..f73e26cf 100644 --- a/webui/src/tests/useSessions.test.tsx +++ b/webui/src/tests/useSessions.test.tsx @@ -170,6 +170,53 @@ describe("useSessions", () => { ]); }); + it("flags history with trailing assistant tool calls as still pending", async () => { + vi.mocked(api.fetchSessionMessages).mockResolvedValue({ + key: "websocket:chat-pending", + created_at: "2026-04-20T10:00:00Z", + updated_at: "2026-04-20T10:05:00Z", + messages: [ + { + role: "assistant", + content: "Using 2 tools", + timestamp: "2026-04-20T10:00:01Z", + tool_calls: [{ id: "call-1" }], + }, + ], + }); + + const { result } = renderHook(() => useSessionHistory("websocket:chat-pending"), { + wrapper: wrap(fakeClient()), + }); + + await waitFor(() => expect(result.current.loading).toBe(false)); + + expect(result.current.hasPendingToolCalls).toBe(true); + }); + + it("does not flag history as pending once the assistant turn has no tool calls", async () => { + vi.mocked(api.fetchSessionMessages).mockResolvedValue({ + key: "websocket:chat-done", + created_at: "2026-04-20T10:00:00Z", + updated_at: "2026-04-20T10:05:00Z", + messages: [ + { + role: "assistant", + content: "All done", + timestamp: "2026-04-20T10:00:01Z", + }, + ], + }); + + const { result } = renderHook(() => useSessionHistory("websocket:chat-done"), { + wrapper: wrap(fakeClient()), + }); + + await waitFor(() => expect(result.current.loading).toBe(false)); + + expect(result.current.hasPendingToolCalls).toBe(false); + }); + it("keeps the session in the list when delete fails", async () => { vi.mocked(api.listSessions).mockResolvedValue([ {