Fix API stream lifecycle for tool-backed requests

tests/test_api_stream.py

@@ -10,8 +10,8 @@ import pytest
 import pytest_asyncio
 
 from nanobot.api.server import (
-    _sse_chunk,
     _SSE_DONE,
+    _sse_chunk,
     create_app,
 )
 
@@ -111,13 +111,13 @@ async def test_stream_true_returns_sse(aiohttp_client) -> None:
     assert resp.content_type == "text/event-stream"
 
     body = await resp.text()
-    lines = [l for l in body.split("\n") if l.startswith("data: ")]
+    lines = [line for line in body.split("\n") if line.startswith("data: ")]
 
     # Should have: 2 token chunks + 1 finish chunk + [DONE]
-    data_lines = [l[len("data: "):] for l in lines]
+    data_lines = [line[len("data: "):] for line in lines]
     assert data_lines[-1] == "[DONE]"
 
-    chunks = [json.loads(l) for l in data_lines[:-1]]
+    chunks = [json.loads(line) for line in data_lines[:-1]]
     assert chunks[0]["choices"][0]["delta"]["content"] == "Hello"
     assert chunks[1]["choices"][0]["delta"]["content"] == " world"
     # Last chunk before [DONE] should have finish_reason=stop
@@ -181,8 +181,12 @@ async def test_stream_sse_chunk_ids_are_consistent(aiohttp_client) -> None:
         json={"messages": [{"role": "user", "content": "go"}], "stream": True},
     )
     body = await resp.text()
-    data_lines = [l[len("data: "):] for l in body.split("\n") if l.startswith("data: ") and l != "data: [DONE]"]
-    chunks = [json.loads(l) for l in data_lines]
+    data_lines = [
+        line[len("data: "):]
+        for line in body.split("\n")
+        if line.startswith("data: ") and line != "data: [DONE]"
+    ]
+    chunks = [json.loads(line) for line in data_lines]
 
     chunk_ids = {c["id"] for c in chunks}
     assert len(chunk_ids) == 1, f"Expected single chunk id, got {chunk_ids}"
@@ -218,6 +222,85 @@ async def test_stream_passes_on_stream_callbacks(aiohttp_client) -> None:
     assert captured_kwargs.get("on_stream_end") is not None
 
 
+@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
+@pytest.mark.asyncio
+async def test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None:
+    """Intermediate stream-end callbacks should not terminate the HTTP stream."""
+    agent = MagicMock()
+
+    async def fake_process_direct(*, on_stream=None, on_stream_end=None, **kwargs):
+        assert on_stream is not None
+        assert on_stream_end is not None
+        await on_stream("planning")
+        await on_stream_end(resuming=True)
+        await asyncio.sleep(0.05)
+        await on_stream(" final")
+        await on_stream_end(resuming=False)
+        return "planning final"
+
+    agent.process_direct = fake_process_direct
+    agent._connect_mcp = AsyncMock()
+    agent.close_mcp = AsyncMock()
+
+    app = create_app(agent, model_name="m")
+    client = await aiohttp_client(app)
+
+    resp = await client.post(
+        "/v1/chat/completions",
+        json={"messages": [{"role": "user", "content": "use a tool"}], "stream": True},
+    )
+
+    assert resp.status == 200
+    body = await resp.text()
+    data_lines = [
+        line[len("data: "):] for line in body.split("\n") if line.startswith("data: ")
+    ]
+    assert data_lines[-1] == "[DONE]"
+
+    chunks = [json.loads(line) for line in data_lines[:-1]]
+    deltas = [c["choices"][0]["delta"].get("content", "") for c in chunks]
+    assert "planning" in deltas
+    assert " final" in deltas
+    assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
+
+
+@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
+@pytest.mark.asyncio
+async def test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None:
+    """stream=true should not return an empty stream when the agent returns content."""
+    agent = MagicMock()
+
+    async def fake_process_direct(*, on_stream=None, on_stream_end=None, **kwargs):
+        assert on_stream is not None
+        assert on_stream_end is not None
+        await on_stream_end(resuming=False)
+        return "plain final"
+
+    agent.process_direct = fake_process_direct
+    agent._connect_mcp = AsyncMock()
+    agent.close_mcp = AsyncMock()
+
+    app = create_app(agent, model_name="m")
+    client = await aiohttp_client(app)
+
+    resp = await client.post(
+        "/v1/chat/completions",
+        json={"messages": [{"role": "user", "content": "hi"}], "stream": True},
+    )
+
+    assert resp.status == 200
+    body = await resp.text()
+    data_lines = [
+        line[len("data: "):] for line in body.split("\n") if line.startswith("data: ")
+    ]
+    chunks = [json.loads(line) for line in data_lines[:-1]]
+    deltas = [c["choices"][0]["delta"].get("content", "") for c in chunks]
+
+    assert "plain final" in deltas
+    assert data_lines[-1] == "[DONE]"
+    assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
+
+
 @pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
 @pytest.mark.asyncio
 async def test_stream_with_session_id(aiohttp_client) -> None: