Merge pull request 'feat: stream all mcp logs' (#505) from listen-log into main

Reviewed-on: https://git.biggo.com/Funmula/dive-mcp-host/pulls/505

Author: LouisFunmula

dive_mcp_host/host/errors.py

@@ -115,4 +115,5 @@ class LogBufferNotFoundError(MCPHostError):
 
     def __init__(self, name: str) -> None:
         """Initialize the error."""
+        self.mcp_name = name
         super().__init__(f"Log buffer {name} not found")

dive_mcp_host/host/tools/log.py

@@ -42,7 +42,7 @@
 import asyncio
 import sys
 from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager
 from datetime import UTC, datetime
 from enum import StrEnum
 from logging import INFO, getLogger
@@ -373,30 +373,53 @@ async def register_buffer(self, buffer: LogBuffer) -> AsyncGenerator[None, None]
     @asynccontextmanager
     async def listen_log(
         self,
-        name: str,
+        names: list[str],
         listener: Callable[[LogMsg], Coroutine[None, None, None]],
     ) -> AsyncGenerator[None, None]:
-        """Listen to log updates from a specific MCP server.
+        """Listen to log updates from MCP servers.
 
         The listener is a context manager,
         user can decide how long it will listen to the buffer.
 
-        Only buffers that are registered to the log manager
-        can be listened to. If the buffer is not registered,
-        `LogBufferNotFoundError` will be raised.
+        Args:
+            names: The names of the MCP servers to listen to.
+            listener: The callback function to call when a log is received.
+
+        Raises:
+            LogBufferNotFoundError: If a specific name is provided but not found.
 
         Example:
             ```python
             async def listener(log: LogMsg) -> None:
                 print(log)
 
 
-            async with log_manager.listen_log(buffer.name, listener):
+            async with log_manager.listen_log(["mcp_server"], listener):
                 await asyncio.sleep(10)
             ```
         """
-        buffer = self._buffers.get(name)
-        if buffer is None:
-            raise LogBufferNotFoundError(name)
-        async with buffer.add_listener(listener):
-            yield
+        if not names:
+            raise LogBufferNotFoundError("no name provided")
+
+        async with AsyncExitStack() as stack:
+            found_buffer: list[LogBuffer] = []
+            not_found: list[str] = []
+
+            for name in names:
+                if buffer := self._buffers.get(name):
+                    found_buffer.append(self._buffers[name])
+                else:
+                    not_found.append(name)
+
+            if not_found:
+                raise LogBufferNotFoundError(",".join(not_found))
+
+            for buffer in found_buffer:
+                await stack.enter_async_context(buffer.add_listener(listener))
+
+            try:
+                yield
+            except Exception:
+                raise
+            finally:
+                pass

dive_mcp_host/httpd/routers/tools.py

@@ -2,7 +2,7 @@
 
 from fastapi import APIRouter, Depends
 from fastapi.responses import StreamingResponse
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 
 from dive_mcp_host.host.tools.model_types import ClientState
 from dive_mcp_host.httpd.conf.mcp_servers import Config
@@ -125,21 +125,28 @@ async def list_tools(  # noqa: PLR0912, C901
     return ToolsResult(success=True, message=None, tools=list(result.values()))
 
 
-@tools.get("/logs/stream")
+class LogsStreamBody(BaseModel):
+    """Body for logs stream API."""
+
+    names: list[str]
+    stream_until: ClientState | None = None
+    stop_on_notfound: bool = True
+    max_retries: int = 10
+
+
+@tools.post("/logs/stream")
 async def stream_server_logs(
-    server_name: str,
-    stream_until: ClientState | None = None,
-    stop_on_notfound: bool = True,
-    max_retries: int = 10,
+    body: LogsStreamBody,
     app: DiveHostAPI = Depends(get_app),
 ) -> StreamingResponse:
-    """Stream logs from a specific MCP server.
+    """Stream logs from MCP servers.
 
     Args:
-        server_name (str): The name of the MCP server to stream logs from.
-        stream_until (ClientState | None): stream until client state is reached.
-        stop_on_notfound (bool): If True, stop streaming if the server is not found.
-        max_retries (int): The maximum number of retries to stream logs.
+        body (LogsStreamBody):
+            - names: MCP servers to listen for logs
+            - stream_until: Stream until mcp server state matches the provided state
+            - stop_on_notfound: Stop streaming if mcp server is not found
+            - max_retries: Retry N times to listen for logs
         app (DiveHostAPI): The DiveHostAPI instance.
 
     Returns:
@@ -155,11 +162,12 @@ async def process() -> None:
             processor = LogStreamHandler(
                 stream=stream,
                 log_manager=log_manager,
-                stream_until=stream_until,
-                stop_on_notfound=stop_on_notfound,
-                max_retries=max_retries,
+                stream_until=body.stream_until,
+                stop_on_notfound=body.stop_on_notfound,
+                max_retries=body.max_retries,
+                server_names=body.names,
             )
-            await processor.stream_logs(server_name)
+            await processor.stream_logs()
 
     stream.add_task(process)
     return response

dive_mcp_host/httpd/routers/utils.py

@@ -814,13 +814,14 @@ async def _get_history_user_input(
 class LogStreamHandler:
     """Handles streaming of logs."""
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         stream: EventStreamContextManager,
         log_manager: LogManager,
         stream_until: ClientState | None = None,
         stop_on_notfound: bool = True,
         max_retries: int = 10,
+        server_names: list[str] | None = None,
     ) -> None:
         """Initialize the log processor."""
         self._stream = stream
@@ -836,26 +837,34 @@ def __init__(
         if stream_until:
             self._stream_until.add(stream_until)
 
+        self._servers_reached_target: set[str] = set()
+        self._server_names: set[str] = set(server_names) if server_names else set()
+
     async def _log_listener(self, msg: LogMsg) -> None:
         await self._stream.write(msg.model_dump_json())
         if msg.client_state in self._stream_until:
-            self._end_event.set()
+            self._servers_reached_target.add(msg.mcp_server_name)
+            if self._server_names == self._servers_reached_target:
+                self._end_event.set()
 
-    async def stream_logs(self, server_name: str) -> None:
-        """Stream logs from specific MCP server.
+    async def stream_logs(self) -> None:
+        """Stream logs from MCP servers.
 
         Keep the connection open until client disconnects or
         client state is reached.
 
+        Streams logs from the given server names.
+
         If self._stop_on_notfound is False, it will keep retrying until
         the log buffer is found or max retries is reached.
         """
         while self._max_retries > 0:
             self._max_retries -= 1
+            self._servers_reached_target = set()
 
             try:
                 async with self._log_manager.listen_log(
-                    name=server_name,
+                    names=list(self._server_names),
                     listener=self._log_listener,
                 ):
                     with suppress(asyncio.CancelledError):
@@ -864,28 +873,30 @@ async def stream_logs(self, server_name: str) -> None:
             except LogBufferNotFoundError as e:
                 logger.warning(
                     "Log buffer not found for server %s, retries left: %d",
-                    server_name,
+                    e.mcp_name,
                     self._max_retries,
                 )
 
                 msg = LogMsg(
                     event=LogEvent.STREAMING_ERROR,
                     body=f"Error streaming logs: {e}",
-                    mcp_server_name=server_name,
+                    mcp_server_name=e.mcp_name,
                 )
                 await self._stream.write(msg.model_dump_json())
 
                 if self._stop_on_notfound or self._max_retries == 0:
                     break
 
-                await asyncio.sleep(1)
+                await asyncio.sleep(0.5)
 
             except Exception as e:
-                logger.exception("Error in log streaming for server %s", server_name)
+                logger.exception(
+                    "Error in log streaming for servers %s", self._server_names
+                )
                 msg = LogMsg(
                     event=LogEvent.STREAMING_ERROR,
                     body=f"Error streaming logs: {e}",
-                    mcp_server_name=server_name,
+                    mcp_server_name="unknown",
                 )
                 await self._stream.write(msg.model_dump_json())
                 break

tests/httpd/routers/test_tools.py

@@ -377,9 +377,7 @@ def test_tools_cache_after_update(test_client):
 def test_stream_logs_notfound(test_client: tuple[TestClient, DiveHostAPI]):
     """Test stream_logs function with not found server."""
     client, _ = test_client
-    response = client.get(
-        "/api/tools/logs/stream", params={"server_name": "missing_server"}
-    )
+    response = client.post("/api/tools/logs/stream", json={"names": ["missing_server"]})
     for line in response.iter_lines():
         content = line.removeprefix("data: ")
         if content in ("[DONE]", ""):
@@ -397,7 +395,7 @@ def test_stream_logs_notfound_wait(test_client: tuple[TestClient, DiveHostAPI]):
 
     def update_tools():
         time.sleep(2)
-        _ = client.post(
+        response = client.post(
             "/api/config/mcpserver",
             json={
                 "mcpServers": {
@@ -418,12 +416,12 @@ def update_tools():
 
     with ThreadPoolExecutor(1) as executor:
         executor.submit(update_tools)
-        response = client.get(
+        response = client.post(
             "/api/tools/logs/stream",
-            params={
-                "server_name": "missing_server",
+            json={
+                "names": ["missing_server"],
                 "stop_on_notfound": False,
-                "max_retries": 5,
+                "max_retries": 10,
                 "stream_until": "running",
             },
         )
@@ -451,7 +449,7 @@ def test_stream_logs_name_with_slash(test_client: tuple[TestClient, DiveHostAPI]
     client, _ = test_client
 
     def update_tools():
-        _ = client.post(
+        response = client.post(
             "/api/config/mcpserver",
             json={
                 "mcpServers": {
@@ -472,10 +470,10 @@ def update_tools():
 
     with ThreadPoolExecutor(1) as executor:
         executor.submit(update_tools)
-        response = client.get(
+        response = client.post(
             "/api/tools/logs/stream",
-            params={
-                "server_name": "name/with/slash",
+            json={
+                "names": ["name/with/slash"],
                 "stop_on_notfound": False,
                 "max_retries": 5,
                 "stream_until": "running",
@@ -498,3 +496,76 @@ def update_tools():
 
         assert responses[-1].event == LogEvent.STATUS_CHANGE
         assert responses[-1].client_state == ClientState.RUNNING
+
+
+def test_stream_multiple_logs(test_client: tuple[TestClient, DiveHostAPI]):
+    """Test streaming multiple server logs."""
+    client, _ = test_client
+
+    def setup_multiple_servers():
+        time.sleep(1)
+        response = client.post(
+            "/api/config/mcpserver",
+            json={
+                "mcpServers": {
+                    "echo": {
+                        "transport": "stdio",
+                        "enabled": True,
+                        "command": "python",
+                        "args": [
+                            "-m",
+                            "dive_mcp_host.host.tools.echo",
+                            "--transport=stdio",
+                        ],
+                    },
+                    "server_two": {
+                        "transport": "stdio",
+                        "enabled": True,
+                        "command": "python",
+                        "args": [
+                            "-m",
+                            "dive_mcp_host.host.tools.echo",
+                            "--transport=stdio",
+                        ],
+                    },
+                }
+            },
+        )
+        assert response.status_code == status.HTTP_200_OK
+
+    with ThreadPoolExecutor(1) as executer:
+        executer.submit(setup_multiple_servers)
+        response = client.post(
+            "/api/tools/logs/stream",
+            json={
+                "names": ["echo", "server_two"],
+                "stream_until": "running",
+                "stop_on_notfound": False,
+                "max_retries": 5,
+            },
+        )
+        responses: list[LogMsg] = []
+        server_names: set[str] = set()
+        servers_reached_running: set[str] = set()
+
+        for line in response.iter_lines():
+            content = line.removeprefix("data: ")
+            if content in ("[DONE]", ""):
+                continue
+
+            data = LogMsg.model_validate_json(content)
+            responses.append(data)
+            server_names.add(data.mcp_server_name)
+
+            if data.client_state == ClientState.RUNNING:
+                servers_reached_running.add(data.mcp_server_name)
+
+        assert len(responses) > 0, "Should receive logs"
+        assert len(server_names) >= 2, "Should receive logs from multiple servers"
+        assert "echo" in server_names or "server_two" in server_names
+
+        running_states = [r for r in responses if r.client_state == ClientState.RUNNING]
+        assert len(running_states) >= 2, "Should have at least 2 servers reach RUNNING"
+        assert len(servers_reached_running) >= 2, (
+            "At least 2 different servers should reach RUNNING state"
+        )

tests/test_log.py

@@ -205,7 +205,7 @@ async def test_listener(msg: LogMsg) -> None:
             await log_buffer.push_log(test_log)
 
             # Listen to logs
-            async with log_manager.listen_log("test_listen", test_listener):
+            async with log_manager.listen_log(["test_listen"], test_listener):
                 # Should receive the existing log
                 assert len(captured_logs) == 1
                 assert captured_logs[0].body == "before listener"
@@ -240,7 +240,7 @@ async def test_listener(msg: LogMsg) -> None:
 
         # Try to listen to a buffer that doesn't exist
         with pytest.raises(LogBufferNotFoundError):
-            async with log_manager.listen_log("nonexistent", test_listener):
+            async with log_manager.listen_log(["nonexistent"], test_listener):
                 pass
     finally:
         import shutil