rework how streams work

Author: paul-nechifor

dimos/agents/mcp/mcp_client.py

@@ -27,7 +27,6 @@
 from langgraph.graph.state import CompiledStateGraph
 from reactivex.disposable import Disposable
 
-from dimos.agents.mcp.tool_stream import ToolStreamEvent
 from dimos.agents.system_prompt import SYSTEM_PROMPT
 from dimos.agents.utils import pretty_print_langchain_message
 from dimos.core.core import rpc
@@ -62,8 +61,6 @@ class McpClient(Module[McpClientConfig]):
     _stop_event: Event
     _http_client: httpx.Client
     _seq_ids: SequentialIds
-    _sse_thread: Thread | None
-    _sse_client: httpx.Client | None
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
@@ -80,8 +77,6 @@ def __init__(self, **kwargs: Any) -> None:
         self._stop_event = Event()
         self._http_client = httpx.Client(timeout=120.0)
         self._seq_ids = SequentialIds()
-        self._sse_thread = None
-        self._sse_client = None
 
     def __reduce__(self) -> Any:
         return (self.__class__, (), {})
@@ -105,6 +100,59 @@ def _mcp_request(self, method: str, params: dict[str, Any] | None = None) -> dic
         result: dict[str, Any] = data.get("result")
         return result
 
+    def _mcp_tool_call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        """Execute a tool call, handling both JSON and SSE streaming responses."""
+        body: dict[str, Any] = {
+            "jsonrpc": "2.0",
+            "id": self._seq_ids.next(),
+            "method": "tools/call",
+            "params": {"name": name, "arguments": arguments},
+        }
+
+        with self._http_client.stream(
+            "POST",
+            self.config.mcp_server_url,
+            json=body,
+            headers={"Accept": "application/json, text/event-stream"},
+        ) as resp:
+            resp.raise_for_status()
+            content_type = resp.headers.get("content-type", "")
+
+            if "text/event-stream" in content_type:
+                return self._consume_sse_tool_response(resp, name)
+
+            data = json.loads(resp.read())
+            if "error" in data:
+                raise RuntimeError(f"MCP error {data['error']['code']}: {data['error']['message']}")
+            result: dict[str, Any] = data.get("result", {})
+            return result
+
+    def _consume_sse_tool_response(
+        self, response: httpx.Response, tool_name: str
+    ) -> dict[str, Any]:
+        """Parse an SSE tool response, injecting notifications as HumanMessages."""
+        result: dict[str, Any] | None = None
+        for line in response.iter_lines():
+            if not line.startswith("data: "):
+                continue
+            try:
+                data = json.loads(line[6:])
+            except json.JSONDecodeError:
+                continue
+
+            if data.get("method") == "notifications/message":
+                text = data.get("params", {}).get("data", "")
+                if text:
+                    self._message_queue.put(
+                        HumanMessage(content=f"[Tool stream update from '{tool_name}']: {text}")
+                    )
+            elif "result" in data:
+                result = data["result"]
+
+        if result is None:
+            return {"content": [{"type": "text", "text": "Stream ended without result."}]}
+        return result
+
     def _fetch_tools(self, timeout: float = 60.0, interval: float = 1.0) -> list[StructuredTool]:
         result = self._try_fetch_tools(timeout=timeout, interval=interval)
         if result is None:
@@ -144,7 +192,7 @@ def _mcp_tool_to_langchain(self, mcp_tool: dict[str, Any]) -> StructuredTool:
         input_schema = mcp_tool.get("inputSchema", {"type": "object", "properties": {}})
 
         def call_tool(**kwargs: Any) -> str:
-            result = self._mcp_request("tools/call", {"name": name, "arguments": kwargs})
+            result = self._mcp_tool_call(name, kwargs)
             content = result.get("content", [])
             parts = [c.get("text", "") for c in content if c.get("type") == "text"]
             text = "\n".join(parts)
@@ -194,18 +242,11 @@ def on_system_modules(self, _modules: list[RPCClient]) -> None:
             )
             self._thread.start()
 
-        self._start_sse_listener()
-
     @rpc
     def stop(self) -> None:
         self._stop_event.set()
-        with self._lock:
-            if self._sse_client is not None:
-                self._sse_client.close()
         if self._thread.is_alive():
             self._thread.join(timeout=2.0)
-        if self._sse_thread is not None and self._sse_thread.is_alive():
-            self._sse_thread.join(timeout=2.0)
         self._http_client.close()
         super().stop()
 
@@ -254,7 +295,7 @@ def dispatch_continuation(
                     tool_args[key] = continuation_context[context_key]
 
         try:
-            result = self._mcp_request("tools/call", {"name": tool_name, "arguments": tool_args})
+            result = self._mcp_tool_call(tool_name, tool_args)
             content = result.get("content", [])
             parts = [c.get("text", "") for c in content if c.get("type") == "text"]
             text = "\n".join(parts)
@@ -302,53 +343,6 @@ def _process_message(
         if self._message_queue.empty():
             self.agent_idle.publish(True)
 
-    def _start_sse_listener(self) -> None:
-        """Connect to the MCP server SSE endpoint to receive tool stream updates."""
-        self._sse_thread = Thread(target=self._sse_loop, name="McpClient-SSE", daemon=True)
-        self._sse_thread.start()
-
-    def _sse_loop(self) -> None:
-        base_url = self.config.mcp_server_url.rsplit("/mcp", 1)[0]
-        sse_url = f"{base_url}/mcp/streams"
-
-        while not self._stop_event.is_set():
-            try:
-                self._sse_connect(sse_url)
-            except Exception:
-                if not self._stop_event.is_set():
-                    # Try reconnecting after a short delay
-                    time.sleep(1.0)
-
-    def _sse_connect(self, sse_url: str) -> None:
-        client = httpx.Client(timeout=None)
-        with self._lock:
-            self._sse_client = client
-        try:
-            with client.stream("GET", sse_url) as response:
-                self._sse_consume(response)
-        finally:
-            with self._lock:
-                self._sse_client = None
-            client.close()
-
-    def _sse_consume(self, response: httpx.Response) -> None:
-        for line in response.iter_lines():
-            if self._stop_event.is_set():
-                return
-            if not line.startswith("data: "):
-                continue
-            try:
-                data = json.loads(line[6:])
-            except json.JSONDecodeError:
-                continue
-            event = ToolStreamEvent(**data)
-            if event.type == "update":
-                self._message_queue.put(
-                    HumanMessage(
-                        content=f"[Tool stream update from '{event.tool_name}']: {event.text}"
-                    )
-                )
-
 
 def _append_image_to_history(
     mcp_client: McpClient, func_name: str, uuid_: str, result: Any

dimos/agents/mcp/mcp_server.py

@@ -16,6 +16,7 @@
 import asyncio
 from collections.abc import AsyncGenerator
 import concurrent.futures
+from dataclasses import dataclass
 import json
 import os
 import time
@@ -30,6 +31,7 @@
 import uvicorn
 
 from dimos.agents.annotation import skill
+from dimos.agents.mcp.tool_stream import ToolStreamEvent
 from dimos.core.core import rpc
 from dimos.core.module import Module
 from dimos.core.rpc_client import RpcCall, RPCClient
@@ -66,12 +68,20 @@ def _jsonrpc_error(req_id: Any, code: int, message: str) -> dict[str, Any]:
     return {"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}}
 
 
+@dataclass(frozen=True, slots=True)
+class _StreamingToolResult:
+    """Marker returned when a tool starts a background stream."""
+
+    req_id: Any
+    tool_name: str
+
+
 def _handle_initialize(req_id: Any) -> dict[str, Any]:
     return _jsonrpc_result(
         req_id,
         {
             "protocolVersion": "2025-11-25",
-            "capabilities": {"tools": {}},
+            "capabilities": {"tools": {}, "logging": {}},
             "serverInfo": {"name": "dimensional", "version": "1.0.0"},
         },
     )
@@ -94,7 +104,7 @@ def _handle_tools_list(req_id: Any, skills: list[SkillInfo]) -> dict[str, Any]:
 
 async def _handle_tools_call(
     req_id: Any, params: dict[str, Any], rpc_calls: dict[str, Any]
-) -> dict[str, Any]:
+) -> dict[str, Any] | _StreamingToolResult:
     name = params.get("name", "")
     args: dict[str, Any] = params.get("arguments") or {}
 
@@ -115,8 +125,8 @@ async def _handle_tools_call(
     duration = f"{time.monotonic() - t0:.3f}s"
 
     if result is None:
-        logger.info("MCP tool done (async)", tool=name, duration=duration)
-        return _jsonrpc_result_text(req_id, "It has started. You will be updated later.")
+        logger.info("MCP tool streaming", tool=name, duration=duration)
+        return _StreamingToolResult(req_id=req_id, tool_name=name)
 
     response = str(result)[:200]
     if hasattr(result, "agent_encode"):
@@ -131,7 +141,7 @@ async def handle_request(
     request: dict[str, Any],
     skills: list[SkillInfo],
     rpc_calls: dict[str, Any],
-) -> dict[str, Any] | None:
+) -> dict[str, Any] | _StreamingToolResult | None:
     """Handle a single MCP JSON-RPC request.
 
     Returns None for JSON-RPC notifications (no ``id``), which must not
@@ -165,27 +175,101 @@ async def mcp_endpoint(request: Request) -> Response:
             {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": "Parse error"}},
             status_code=400,
         )
+
+    # Pre-register a queue for tool stream events when the client accepts SSE.
+    accept = request.headers.get("accept", "")
+    is_tool_call = body.get("method") == "tools/call"
+    client_accepts_sse = "text/event-stream" in accept
+
+    stream_queue: asyncio.Queue[dict[str, Any]] | None = None
+    if is_tool_call and client_accepts_sse:
+        stream_queue = asyncio.Queue()
+        app.state.sse_queues.append(stream_queue)
+
     result = await handle_request(body, request.app.state.skills, request.app.state.rpc_calls)
+
+    # Streaming tool: return SSE response if client supports it.
+    if isinstance(result, _StreamingToolResult):
+        if stream_queue is not None:
+            return _streaming_tool_response(result, stream_queue)
+        # Client doesn't support SSE — fall back to immediate JSON.
+        return JSONResponse(
+            _jsonrpc_result_text(result.req_id, "It has started. You will be updated later.")
+        )
+
+    # Non-streaming: remove the pre-registered queue if any.
+    if stream_queue is not None:
+        try:
+            app.state.sse_queues.remove(stream_queue)
+        except ValueError:
+            pass
+
     if result is None:
         return Response(status_code=204)
     return JSONResponse(result)
 
 
-@app.get("/mcp/streams")
-async def streams_sse_endpoint() -> StreamingResponse:
-    """Server-Sent Events endpoint for tool stream updates.
+_STREAM_TIMEOUT = 300.0  # seconds
+
+
+def _sse_event(data: dict[str, Any]) -> str:
+    """Format a JSON-RPC message as an SSE ``event: message`` frame."""
+    return f"event: message\ndata: {json.dumps(data)}\n\n"
+
+
+def _streaming_tool_response(
+    streaming: _StreamingToolResult,
+    queue: asyncio.Queue[dict[str, Any]],
+) -> StreamingResponse:
+    """Build an SSE response that forwards ToolStream events as MCP log notifications.
 
-    Clients subscribe here to receive real-time updates from long-running
-    skills that use ``ToolStream``.
+    The response streams ``notifications/message`` for each update and ends
+    with the JSON-RPC result carrying the accumulated text.
     """
-    queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
-    app.state.sse_queues.append(queue)
 
     async def event_generator() -> AsyncGenerator[str, None]:
+        stream_id: str | None = None
+        collected: list[str] = []
         try:
             while True:
-                data = await queue.get()
-                yield f"data: {json.dumps(data)}\n\n"
+                try:
+                    data = await asyncio.wait_for(queue.get(), timeout=_STREAM_TIMEOUT)
+                except asyncio.TimeoutError:
+                    text = "\n".join(collected) if collected else "No updates received."
+                    yield _sse_event(_jsonrpc_result_text(streaming.req_id, text))
+                    return
+
+                try:
+                    event = ToolStreamEvent(**data)
+                except (TypeError, KeyError):
+                    continue
+
+                # Filter: match by tool_name, lock onto the first stream_id seen.
+                if event.tool_name != streaming.tool_name:
+                    continue
+                if stream_id is None:
+                    stream_id = event.stream_id
+                elif event.stream_id != stream_id:
+                    continue
+
+                if event.type == "update" and event.text:
+                    collected.append(event.text)
+                    yield _sse_event(
+                        {
+                            "jsonrpc": "2.0",
+                            "method": "notifications/message",
+                            "params": {
+                                "level": "info",
+                                "logger": event.tool_name,
+                                "data": event.text,
+                            },
+                        }
+                    )
+
+                elif event.type == "close":
+                    text = "\n".join(collected) if collected else "Stream completed."
+                    yield _sse_event(_jsonrpc_result_text(streaming.req_id, text))
+                    return
         except asyncio.CancelledError:
             pass
         finally: