Merge branch 'feat/streaming-llm' into staging

Author: kelvin1804

openhands/agenthub/codeact_agent/codeact_agent.py

@@ -302,9 +302,7 @@ async def _handle_streaming_response(self, streaming_response):
                         {
                             'message': {
                                 'role': 'assistant',
-                                'content': accumulated_content
-                                if accumulated_content
-                                else None,
+                                'content': None,
                                 'tool_calls': formatted_tool_calls,
                             },
                             'index': 0,

openhands/mcp/utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import re
 from typing import Any, Optional
@@ -55,32 +56,66 @@ async def create_mcp_clients(
     mnemonic: Optional[str] = None,
 ) -> list[MCPClient]:
     mcp_clients: list[MCPClient] = []
-    # Initialize SSE connections
+
+    # Create connection tasks for parallel execution
+    connection_tasks = []
+    clients = []
+
     for name, mcp_config in dict_mcp_config.items():
         logger.info(
             f'Initializing MCP {name} agent for {mcp_config.url} with {mcp_config.mode} connection...'
         )
         # check if the name in a search engine config
         if f'search_engine_{name}' in dict_mcp_config:
             continue
+
         client = MCPClient(name=name)
-        try:
-            await client.connect_sse(mcp_config.url, sid, mnemonic)
-            # Only add the client to the list after a successful connection
-            mcp_clients.append(client)
-            logger.info(f'Connected to MCP server {mcp_config.url} via SSE')
-        except Exception as e:
-            logger.error(f'Failed to connect to {mcp_config.url}: {str(e)}')
+        clients.append((client, name, mcp_config))
+
+        # Create connection task
+        async def connect_client(client, mcp_config):
             try:
-                await client.disconnect()
-            except Exception as disconnect_error:
-                logger.error(
-                    f'Error during disconnect after failed connection: {str(disconnect_error)}'
-                )
+                await client.connect_sse(mcp_config.url, sid, mnemonic)
+                return client, None  # Success: return client and no error
+            except Exception as e:
+                return None, e  # Failure: return no client and the error
+
+        connection_tasks.append(connect_client(client, mcp_config))
+
+    if connection_tasks:
+        results = await asyncio.gather(*connection_tasks, return_exceptions=True)
+
+        for i, result in enumerate(results):
+            client, name, mcp_config = clients[i]
+
+            if isinstance(result, Exception):
+                logger.error(f'Failed to connect to {mcp_config.url}: {str(result)}')
+                await _safe_disconnect(client)
+            elif isinstance(result, tuple) and len(result) == 2:
+                client_result, error = result
+                if error is None:  # Success case
+                    mcp_clients.append(client_result)
+                    logger.info(f'Connected to MCP server {mcp_config.url} via SSE')
+                else:  # Error case
+                    logger.error(f'Failed to connect to {mcp_config.url}: {str(error)}')
+                    await _safe_disconnect(client)
+            else:
+                logger.error(f'Unexpected result format for {mcp_config.url}: {result}')
+                await _safe_disconnect(client)
 
     return mcp_clients
 
 
+async def _safe_disconnect(client: MCPClient):
+    """Safely disconnect a client with error handling."""
+    try:
+        await client.disconnect()
+    except Exception as disconnect_error:
+        logger.error(
+            f'Error during disconnect after failed connection: {str(disconnect_error)}'
+        )
+
+
 async def fetch_mcp_tools_from_config(
     dict_mcp_config: dict[str, MCPConfig],
     sid: Optional[str] = None,

openhands/memory/conversation_memory.py

@@ -24,6 +24,7 @@
     A2ASendTaskAction,
 )
 from openhands.events.action.mcp import McpAction
+from openhands.events.action.message import StreamingMessageAction
 from openhands.events.event import Event, RecallType
 from openhands.events.observation import (
     AgentCondensationObservation,
@@ -328,6 +329,13 @@ def _process_action(
                     content=content,
                 )
             ]
+        elif isinstance(action, StreamingMessageAction):
+            return [
+                Message(
+                    role='assistant',
+                    content=[TextContent(text=action.content)],
+                )
+            ]
         elif isinstance(action, CmdRunAction) and action.source == 'user':
             content = [
                 TextContent(text=f'User executed the command:\n{action.command}')