Skip to content

Commit 16bf20d

Browse files
cristipufucristipufu
authored
Merge pull request #39 from UiPath/fix/samples
fix: collect server output
2 parents 32e434f + f2e7777 commit 16bf20d

File tree

3 files changed

+70
-86
lines changed

3 files changed

+70
-86
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-mcp"
3-
version = "0.0.35"
3+
version = "0.0.36"
44
description = "UiPath MCP SDK"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.10"

‎src/uipath_mcp/_cli/_runtime/_runtime.py‎

Lines changed: 21 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,26 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
101101
for task in pending:
102102
task.cancel()
103103

104-
return UiPathRuntimeResult()
104+
session_outputs = {}
105+
for session_id, session_server in self.session_servers.items():
106+
try:
107+
await session_server.cleanup()
108+
stderr_output = session_server.get_server_stderr()
109+
if stderr_output:
110+
session_outputs[session_id] = stderr_output
111+
except Exception as e:
112+
logger.error(f"Error stopping session {session_id}: {str(e)}")
113+
114+
output_result = {}
115+
if len(session_outputs) == 1:
116+
# If there's only one session, use a single "output" key
117+
first_session_id = next(iter(session_outputs))
118+
output_result["output"] = session_outputs[first_session_id]
119+
elif session_outputs:
120+
# If there are multiple sessions, use the sessionId as the key
121+
output_result = session_outputs
122+
123+
return UiPathRuntimeResult(output=output_result)
105124

106125
except Exception as e:
107126
if isinstance(e, UiPathMcpRuntimeError):
@@ -248,49 +267,8 @@ async def _register(self) -> None:
248267
# We don't continue with registration here - we'll do it after the context managers
249268

250269
except BaseException as e:
251-
is_process_lookup = False
252-
is_exception_group_with_lookup = False
270+
logger.error(f"Error during server initialization: {e}")
253271

254-
if isinstance(e, ProcessLookupError):
255-
is_process_lookup = True
256-
else:
257-
try:
258-
from exceptiongroup import ExceptionGroup # Backport for < 3.11
259-
260-
if sys.version_info >= (3, 11) and isinstance(e, ExceptionGroup):
261-
for sub_e in e.exceptions:
262-
if isinstance(sub_e, ProcessLookupError):
263-
is_exception_group_with_lookup = True
264-
break
265-
except ImportError:
266-
# ExceptionGroup is not available (Python < 3.11)
267-
pass
268-
269-
if is_process_lookup or is_exception_group_with_lookup:
270-
logger.info(
271-
"Process already terminated during cleanup - this is expected"
272-
)
273-
else:
274-
logger.error(f"Error during server initialization: {e}")
275-
raise UiPathMcpRuntimeError(
276-
"SERVER_ERROR",
277-
"Server initialization failed",
278-
str(e),
279-
UiPathErrorCategory.DEPLOYMENT,
280-
) from e
281-
282-
if is_process_lookup or is_exception_group_with_lookup:
283-
logger.info(
284-
"Process already terminated during cleanup - this is expected"
285-
)
286-
else:
287-
logger.error(f"Error during server initialization: {e}")
288-
raise UiPathMcpRuntimeError(
289-
"INITIALIZATION_ERROR",
290-
"Server initialization failed",
291-
str(e),
292-
UiPathErrorCategory.DEPLOYMENT,
293-
) from e
294272

295273
# Now that we're outside the context managers, check if initialization succeeded
296274
if not initialization_successful:
@@ -347,13 +325,6 @@ async def cleanup(self) -> None:
347325
"""Clean up all resources."""
348326
logger.info("Cleaning up all resources")
349327

350-
# Clean up all session servers
351-
for session_id, session_server in list(self.session_servers.items()):
352-
try:
353-
await session_server.cleanup()
354-
except Exception as e:
355-
logger.error(f"Error cleaning up session {session_id}: {str(e)}")
356-
357328
self.session_servers.clear()
358329

359330
if self.signalr_client and hasattr(self.signalr_client, "_transport"):

‎src/uipath_mcp/_cli/_runtime/_session.py‎

Lines changed: 48 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import asyncio
22
import logging
3+
import tempfile
4+
from typing import Optional
35

46
import mcp.types as types
57
from mcp import StdioServerParameters
@@ -8,7 +10,6 @@
810
from uipath import UiPath
911

1012
from .._utils._config import McpServer
11-
from ._logger import NullLogger
1213
from ._tracer import McpTracer
1314

1415
logger = logging.getLogger(__name__)
@@ -29,6 +30,11 @@ def __init__(self, server_config: McpServer, session_id: str):
2930
self._message_queue = asyncio.Queue()
3031
self._uipath = UiPath()
3132
self._mcp_tracer = McpTracer(tracer, logger)
33+
self._server_stderr_output: Optional[str] = None
34+
35+
def get_server_stderr(self) -> Optional[str]:
36+
"""Returns the captured stderr output from the MCP server process."""
37+
return self._server_stderr_output
3238

3339
async def start(self) -> None:
3440
"""Start the server process in a separate task."""
@@ -59,44 +65,51 @@ async def start(self) -> None:
5965
async def _run_server(self, server_params: StdioServerParameters) -> None:
6066
"""Run the server in proper context managers."""
6167
logger.info(f"Starting server process for session {self.session_id}")
62-
try:
63-
stderr_null = NullLogger()
64-
65-
async with stdio_client(server_params, errlog=stderr_null) as (
66-
read,
67-
write,
68-
):
69-
self.read_stream, self.write_stream = read, write
70-
logger.info(f"Session {self.session_id} - stdio client started")
68+
self._server_stderr_output = None
69+
with tempfile.TemporaryFile(mode="w+") as stderr_temp:
70+
try:
71+
async with stdio_client(server_params, errlog=stderr_temp) as (
72+
read,
73+
write,
74+
):
75+
self.read_stream, self.write_stream = read, write
76+
logger.info(f"Session {self.session_id} - stdio client started")
7177

72-
# Start the message consumer task
73-
consumer_task = asyncio.create_task(self._consume_messages())
78+
# Start the message consumer task
79+
consumer_task = asyncio.create_task(self._consume_messages())
7480

75-
# Process incoming messages from the server
76-
try:
77-
while True:
78-
logger.info("Waiting for messages...")
79-
message = await self.read_stream.receive()
80-
await self.send_outgoing_message(message)
81-
finally:
82-
# Cancel the consumer when we exit the loop
83-
consumer_task.cancel()
81+
# Process incoming messages from the server
8482
try:
85-
await consumer_task
86-
except asyncio.CancelledError:
87-
pass
83+
while True:
84+
logger.info(
85+
f"Session {self.session_id}: Waiting for messages..."
86+
)
87+
message = await self.read_stream.receive()
88+
await self.send_outgoing_message(message)
89+
finally:
90+
stderr_temp.seek(0)
91+
self._server_stderr_output = stderr_temp.read()
92+
logger.debug(
93+
f"Session {self.session_id} - Server stderr output:\n{self._server_stderr_output}"
94+
)
95+
# Cancel the consumer when we exit the loop
96+
consumer_task.cancel()
97+
try:
98+
await consumer_task
99+
except asyncio.CancelledError:
100+
pass
88101

89-
except Exception as e:
90-
logger.error(
91-
f"Error in server process for session {self.session_id}: {e}",
92-
exc_info=True,
93-
)
94-
finally:
95-
# The context managers will handle cleanup of resources
96-
logger.info(f"Server process for session {self.session_id} has ended")
97-
self.read_stream = None
98-
self.write_stream = None
99-
self.mcp_session = None
102+
except Exception as e:
103+
logger.error(
104+
f"Error in server process for session {self.session_id}: {e}",
105+
exc_info=True,
106+
)
107+
finally:
108+
# The context managers will handle cleanup of resources
109+
logger.info(f"Server process for session {self.session_id} has ended")
110+
self.read_stream = None
111+
self.write_stream = None
112+
self.mcp_session = None
100113

101114
def _on_task_done(self, task):
102115
"""Handle task completion."""

0 commit comments

Comments
 (0)