fix: properly terminate subprocess

Author: cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-mcp"
-version = "0.0.57"
+version = "0.0.58"
 description = "UiPath MCP SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -8,7 +8,6 @@
 
 import mcp.types as types
 from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
 from opentelemetry import trace
 from pysignalr.client import SignalRClient
 from uipath import UiPath
@@ -23,6 +22,7 @@
 from ._context import UiPathMcpRuntimeContext
 from ._exception import UiPathMcpRuntimeError
 from ._session import SessionServer
+from ._stdio_client import stdio_client
 
 logger = logging.getLogger(__name__)
 tracer = trace.get_tracer(__name__)

src/uipath_mcp/_cli/_runtime/_stdio_client.py

@@ -0,0 +1,105 @@
+import sys
+from contextlib import asynccontextmanager
+from typing import TextIO
+
+import anyio
+import anyio.lowlevel
+import mcp.types as types
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from anyio.streams.text import TextReceiveStream
+from mcp.client.stdio import (
+    StdioServerParameters,
+    _create_platform_compatible_process,
+    _get_executable_command,
+    get_default_environment,
+)
+
+
+@asynccontextmanager
+async def stdio_client(server: StdioServerParameters, errlog: TextIO = sys.stderr):
+    """
+    Client transport for stdio: this will connect to a server by spawning a
+    process and communicating with it over stdin/stdout.
+    """
+    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+
+    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
+    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+
+    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
+
+    command = _get_executable_command(server.command)
+
+    # Open process with stderr piped for capture
+    process = await _create_platform_compatible_process(
+        command=command,
+        args=server.args,
+        env=(
+            {**get_default_environment(), **server.env}
+            if server.env is not None
+            else get_default_environment()
+        ),
+        errlog=errlog,
+        cwd=server.cwd,
+    )
+
+    async def stdout_reader():
+        assert process.stdout, "Opened process is missing stdout"
+
+        try:
+            async with read_stream_writer:
+                buffer = ""
+                async for chunk in TextReceiveStream(
+                    process.stdout,
+                    encoding=server.encoding,
+                    errors=server.encoding_error_handler,
+                ):
+                    lines = (buffer + chunk).split("\n")
+                    buffer = lines.pop()
+
+                    for line in lines:
+                        try:
+                            message = types.JSONRPCMessage.model_validate_json(line)
+                        except Exception as exc:
+                            await read_stream_writer.send(exc)
+                            continue
+
+                        await read_stream_writer.send(message)
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async def stdin_writer():
+        assert process.stdin, "Opened process is missing stdin"
+
+        try:
+            async with write_stream_reader:
+                async for message in write_stream_reader:
+                    json = message.model_dump_json(by_alias=True, exclude_none=True)
+                    await process.stdin.send(
+                        (json + "\n").encode(
+                            encoding=server.encoding,
+                            errors=server.encoding_error_handler,
+                        )
+                    )
+        except anyio.ClosedResourceError:
+            await anyio.lowlevel.checkpoint()
+
+    async with (
+        anyio.create_task_group() as tg,
+        process,
+    ):
+        tg.start_soon(stdout_reader)
+        tg.start_soon(stdin_writer)
+        try:
+            yield read_stream, write_stream
+        finally:
+            # Clean up process to prevent any dangling orphaned processes
+            try:
+                process.terminate()
+                with anyio.fail_after(2.0):
+                    await process.wait()
+            except TimeoutError:
+                # Force kill if it doesn't terminate
+                process.kill()