Merge branch 'development' of ssh://git.biggo.com:222/Funmula/dive-mcp-host into development

Author: LouisFunmula

.gitignore

@@ -45,3 +45,4 @@ cache/
 upload/
 cli_config.json
 logs/
+oap_config.json

dive_mcp_host/host/conf/__init__.py

@@ -58,6 +58,7 @@ class ServerConfig(BaseModel):
         ProxyUrl | None,
         BeforeValidator(_rewrite_socks),
     ] = None
+    initial_timeout: float = 10
 
     @field_serializer("headers", when_used="json")
     def dump_headers(self, v: dict[str, SecretStr] | None) -> dict[str, str] | None:

dive_mcp_host/host/tools/hack/stdio_server.py

@@ -4,19 +4,22 @@
 import subprocess
 import sys
 from collections.abc import AsyncGenerator
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 from pathlib import Path
 
 import anyio
 import anyio.abc
 import anyio.lowlevel
+from anyio.abc import Process
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from anyio.streams.text import TextReceiveStream
 from mcp import types
 from mcp.client.stdio import StdioServerParameters, get_default_environment
-from mcp.client.stdio.win32 import (
+from mcp.os.posix.utilities import terminate_posix_process_tree
+from mcp.os.win32.utilities import (
+    FallbackProcess,
     get_windows_executable_command,
-    terminate_windows_process,
+    terminate_windows_process_tree,
 )
 from mcp.shared.message import SessionMessage
 
@@ -43,6 +46,8 @@
     else ["HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER"]
 )
 
+PROCESS_TERMINATION_TIMEOUT = 2.0
+
 
 @asynccontextmanager
 async def stdio_client(  # noqa: C901, PLR0915
@@ -156,12 +161,29 @@ async def stdin_writer() -> None:
             logger.error("Error, closing process %s: %s", process.pid, exc)
             raise
         finally:
-            # Clean up process to prevent any dangling orphaned processes
-            logger.info("Terminated process %s", process.pid)
-            # Some process never terminates, so we need to kill it.
-            await terminate_windows_process(process)
-            status = await process.wait()
-            logger.info("Process %s exited with status %s", process.pid, status)
+            # MCP spec: stdio shutdown sequence
+            # 1. Close input stream to server
+            # 2. Wait for server to exit, or send SIGTERM if it doesn't exit in time
+            # 3. Send SIGKILL if still not exited
+            if process.stdin:
+                with suppress(Exception):
+                    await process.stdin.aclose()
+            try:
+                # Give the process time to exit gracefully after stdin closes
+                with anyio.fail_after(PROCESS_TERMINATION_TIMEOUT):
+                    await process.wait()
+            except TimeoutError:
+                # Process didn't exit from stdin closure, use platform-specific
+                # termination
+                # which handles SIGTERM -> SIGKILL escalation
+                await _terminate_process_tree(process)
+            except ProcessLookupError:
+                # Process already exited, which is fine
+                pass
+            await read_stream.aclose()
+            await write_stream.aclose()
+            await read_stream_writer.aclose()
+            await write_stream_reader.aclose()
     logger.error("Process %s closed", "xx")
 
 
@@ -189,3 +211,23 @@ async def _create_platform_compatible_process(
     logger.info("launched process: %s, pid: %s", command, process.pid)
 
     return process
+
+
+async def _terminate_process_tree(
+    process: Process | FallbackProcess, timeout_seconds: float = 2.0
+) -> None:
+    """Terminate a process and all its children using platform-specific methods.
+
+    Unix: Uses os.killpg() for atomic process group termination
+    Windows: Uses Job Objects via pywin32 for reliable child process cleanup
+
+    Args:
+        process: The process to terminate
+        timeout_seconds: Timeout in seconds before force killing (default: 2.0)
+    """
+    if sys.platform == "win32":
+        await terminate_windows_process_tree(process, timeout_seconds)
+    else:
+        # FallbackProcess should only be used for Windows compatibility
+        assert isinstance(process, Process)
+        await terminate_posix_process_tree(process, timeout_seconds)

dive_mcp_host/host/tools/local_http_server.py

@@ -4,6 +4,7 @@
 import sys
 from collections.abc import AsyncGenerator
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
+from time import time
 from typing import Any
 
 import httpx
@@ -30,7 +31,6 @@ async def local_http_server(  # noqa: C901, PLR0913, PLR0915
     command: str | None = None,
     args: list[str] | None = None,
     env: dict[str, str] | None = None,
-    max_connection_retries: int = 10,
     headers: dict[str, Any] | None = None,
 ) -> AsyncGenerator[tuple[InitializeResult, ListToolsResult, int], None]:
     """Create a local MCP server client.
@@ -65,7 +65,13 @@ def _sse_client(
             value = headers[key]
             if isinstance(value, SecretStr):
                 headers[key] = value.get_secret_value()
-        return sse_client(url=url, headers=headers)
+        logger.debug("Connecting sse client with timeout: %s", config.initial_timeout)
+        return sse_client(
+            url=url,
+            headers=headers,
+            sse_read_timeout=config.initial_timeout,
+            timeout=config.initial_timeout,
+        )
 
     get_client = _sse_client if config.transport == "sse" else websocket_client
     logger.debug("Starting local MCP server %s with command: %s", config.name, command)
@@ -80,7 +86,6 @@ def _sse_client(
     ):
         logger.error("Failed to start subprocess for %s", config.name)
         raise RuntimeError("failed to start subprocess")
-    retried = 0
 
     # it tooks time to start the server, so we need to retry
     async def _read_stdout(
@@ -114,38 +119,40 @@ async def _read_stderr(
         name="read-stdout",
     )
 
+    start_time = time()
     try:
-        while retried < max_connection_retries:
-            await asyncio.sleep(0.3 if retried == 0 else 1)
-            logger.debug(
-                "Attempting to connect to server %s (attempt %d/%d)",
-                config.name,
-                retried + 1,
-                max_connection_retries,
-            )
+        logger.debug(
+            "Server %s initalizing with timeout: %s",
+            config.name,
+            config.initial_timeout,
+        )
+        while (time() - start_time) < config.initial_timeout:
             with suppress(TimeoutError, httpx.HTTPError):
                 async with (
                     get_client(url=config.url) as streams,
                     ClientSession(*streams) as session,
                 ):
-                    async with asyncio.timeout(10):
+                    async with asyncio.timeout(config.initial_timeout):
                         initialize_result = await session.initialize()
                         tools = await session.list_tools()
                         logger.info(
                             "Successfully connected to server %s, got tools: %s",
                             config.name,
                             tools,
                         )
-                        break
-            retried += 1
+                    logger.info("Connected to the server %s", config.name)
+                    yield initialize_result, tools, subprocess.pid
+                    break
+            await asyncio.sleep(0.3)
         else:
-            raise InvalidMcpServerError(config.name)
-        logger.info(
-            "Connected to the server %s after %d attempts", config.name, retried
-        )
-        yield initialize_result, tools, subprocess.pid
+            logger.warning(
+                "Connected to the server %s failed after %s seconds",
+                config.name,
+                config.initial_timeout,
+            )
+            raise InvalidMcpServerError(config.name, reason="failed to initalize")
     finally:
-        with suppress(TimeoutError):
+        with suppress(TimeoutError, ProcessLookupError, asyncio.CancelledError):
             logger.debug("Terminating subprocess for %s", config.name)
             read_stderr_task.cancel()
             read_stdout_task.cancel()
@@ -158,7 +165,7 @@ async def _read_stderr(
             subprocess = None
         if subprocess:
             logger.info("Timeout to terminate mcp-server %s. Kill it.", config.name)
-            with suppress(TimeoutError):
+            with suppress(TimeoutError, ProcessLookupError, asyncio.CancelledError):
                 read_stderr_task.cancel()
                 read_stdout_task.cancel()
                 subprocess.kill()

dive_mcp_host/host/tools/mcp_server.py

@@ -111,7 +111,7 @@ class McpServer(ContextProtocol):
 
     RETRY_LIMIT: int = 3
     KEEP_ALIVE_INTERVAL: float = 60
-    RESTART_INTERVAL: float = 3
+    RESTART_INTERVAL: float = 1
 
     def __init__(
         self,
@@ -214,11 +214,16 @@ async def _message_handler(
 
     async def _init_tool_info(self, session: ClientSession) -> None:
         """Initialize the session."""
-        async with asyncio.timeout(10):
+        logger.debug(
+            "Client %s initalizing with timeout: %s",
+            self.name,
+            self.config.initial_timeout,
+        )
+        async with asyncio.timeout(self.config.initial_timeout):
             # When using stdio, the initialize call may block indefinitely
             self._initialize_result = await session.initialize()
             logger.debug(
-                "Client %s initializing, result: %s",
+                "Client %s initialize result: %s",
                 self.name,
                 self._initialize_result,
             )
@@ -428,15 +433,19 @@ async def _session_ctx_mgr_wrapper(
             async with self._cond:
                 self._cond.notify_all()
 
-    async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915
+    async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915, PLR0912
         """Client watcher task.
 
         Restart the client if need.
         Only this watcher can set the client status to RUNNING / FAILED.
         """
         env = os.environ.copy()
         env.update(self.config.env)
-        while True:
+        start_time = time.time()
+        while (
+            self._retries == 0
+            or (time.time() - start_time) < self.config.initial_timeout
+        ):
             should_break = False
             try:
                 logger.debug("Attempting to initialize client %s", self.name)
@@ -466,6 +475,11 @@ async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915
                             self.name,
                             self._client_status,
                         )
+                        if self._client_status == ClientState.RESTARTING:
+                            self._retries = 0
+                            start_time = time.time()
+                            continue
+                        return
             except* ProcessLookupError as eg:
                 # this raised when a stdio process is exited
                 # and the initialize call is timeout
@@ -522,25 +536,25 @@ async def _stdio_client_watcher(self) -> None:  # noqa: C901, PLR0915
             if self._client_status == ClientState.CLOSED:
                 logger.info("Client %s closed, stopping watcher", self.name)
                 return
-            if self._retries >= self.RETRY_LIMIT or should_break:
-                logger.warning(
-                    "client for [%s] failed after %d retries %s",
-                    self.name,
-                    self._retries,
-                    self._exception,
-                )
-                async with self._cond:
-                    if self._client_status != ClientState.CLOSED:
-                        await self.__change_state(ClientState.FAILED, None, False)
-                return
+            if should_break:
+                break
             logger.debug(
-                "Retrying client initialization for %s (attempt %d/%d)",
+                "Retrying client initialization for %s (attempt %d)",
                 self.name,
                 self._retries,
-                self.RETRY_LIMIT,
             )
             await asyncio.sleep(self.RESTART_INTERVAL)
 
+        logger.warning(
+            "client for [%s] failed after %d retries %s",
+            self.name,
+            self._retries,
+            self._exception,
+        )
+        async with self._cond:
+            if self._client_status != ClientState.CLOSED:
+                await self.__change_state(ClientState.FAILED, None, False)
+
     async def _stdio_setup(self) -> None:
         """Setup the stdio client."""
         self._server_task = asyncio.create_task(
@@ -714,7 +728,10 @@ def _http_get_client(
     async def _http_init_client(self) -> None:
         """Initialize the HTTP client."""
         async with (
-            self._http_get_client(sse_read_timeout=30) as streams,
+            self._http_get_client(
+                sse_read_timeout=self.config.initial_timeout,
+                timeout=self.config.initial_timeout,
+            ) as streams,
             ClientSession(
                 *[streams[0], streams[1]], message_handler=self._message_handler
             ) as session,
@@ -724,7 +741,11 @@ async def _http_init_client(self) -> None:
     async def _http_setup(self) -> None:
         """Setup the http client."""
         self._retries = 0
-        for _ in range(self.RETRY_LIMIT):
+        start_time = time.time()
+        while (
+            self._retries == 0
+            or (time.time() - start_time) < self.config.initial_timeout
+        ):
             should_break = False
             try:
                 await self._http_init_client()

dive_mcp_host/httpd/conf/mcp_servers.py

@@ -6,7 +6,14 @@
 from pathlib import Path
 from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, BeforeValidator, Field, SecretStr, field_serializer
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    SecretStr,
+    field_serializer,
+)
 
 from dive_mcp_host.env import DIVE_CONFIG_DIR
 from dive_mcp_host.host.conf import ProxyUrl
@@ -34,6 +41,13 @@ class MCPServerConfig(BaseModel):
     proxy: ProxyUrl | None = None
     headers: dict[str, SecretStr] | None = Field(default_factory=dict)
     exclude_tools: list[str] = Field(default_factory=list)
+    initial_timeout: float = Field(default=10, ge=10, alias="initialTimeout")
+
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_by_alias=True,
+        serialize_by_alias=True,
+    )
 
     def model_post_init(self, _: Any) -> None:
         """Post-initialization hook."""
@@ -56,6 +70,12 @@ class Config(BaseModel):
         alias="mcpServers", default_factory=dict
     )
 
+    model_config = ConfigDict(
+        validate_by_name=True,
+        validate_by_alias=True,
+        serialize_by_alias=True,
+    )
+
 
 type McpServerConfigCallback = Callable[[Config], Config | Coroutine[Any, Any, Config]]
 UpdateAllConfigsHookName = "httpd.config.mcp_servers.update_all_configs"

dive_mcp_host/httpd/server.py

@@ -236,6 +236,7 @@ async def load_host_config(self) -> HostConfig:
                 headers=server_config.headers or {},
                 proxy=server_config.proxy or None,
                 exclude_tools=server_config.exclude_tools,
+                initial_timeout=server_config.initial_timeout,
             )
 
         logger.debug("got %s mcp servers in config", len(mcp_servers))