Merge pull request #23 from UiPath/fix/samples

feat: add tracing

Author: cristipufu

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "uipath-mcp"
-version = "0.0.18"
+version = "0.0.19"
 description = "UiPath MCP SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
     "mcp==1.6.0",
     "pysignalr==1.2.0",
-    "uipath==2.0.1",
+    "uipath==2.0.11",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",

src/uipath_mcp/_cli/_runtime/_session.py

@@ -4,12 +4,15 @@
 import mcp.types as types
 from mcp import StdioServerParameters
 from mcp.client.stdio import stdio_client
+from opentelemetry import trace
 from uipath import UiPath
+from uipath.tracing import traced, wait_for_tracers
 
 from .._utils._config import McpServer
 from ._logger import LoggerAdapter
 
 logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
 
 
 class SessionServer:
@@ -57,33 +60,38 @@ async def _run_server(self, server_params: StdioServerParameters) -> None:
         logger.info(f"Starting server process for session {self.session_id}")
         try:
             stderr_adapter = LoggerAdapter(logger)
-            async with stdio_client(server_params, errlog=stderr_adapter) as (
-                read,
-                write,
-            ):
-                self.read_stream, self.write_stream = read, write
-                logger.info(f"Session {self.session_id} - stdio client started")
+            with tracer.start_as_current_span(self.server_config.name) as root_span:
+                root_span.set_attribute("session_id", self.session_id)
+                root_span.set_attribute(
+                    "server_params", server_params.model_dump_json()
+                )
+                async with stdio_client(server_params, errlog=stderr_adapter) as (
+                    read,
+                    write,
+                ):
+                    self.read_stream, self.write_stream = read, write
+                    logger.info(f"Session {self.session_id} - stdio client started")
 
-                logger.info(f"Session {self.session_id} - MCP session initialized")
+                    logger.info(f"Session {self.session_id} - MCP session initialized")
 
-                # Start the message consumer task
-                consumer_task = asyncio.create_task(self._consume_messages())
+                    # Start the message consumer task
+                    consumer_task = asyncio.create_task(self._consume_messages())
 
-                # Process incoming messages from the server
-                try:
-                    while True:
-                        print("Waiting for messages...")
-                        message = await self.read_stream.receive()
-                        json_str = message.model_dump_json()
-                        print(f"Received message from local server: {json_str}")
-                        await self.send_outgoing_message(message)
-                finally:
-                    # Cancel the consumer when we exit the loop
-                    consumer_task.cancel()
+                    # Process incoming messages from the server
                     try:
-                        await consumer_task
-                    except asyncio.CancelledError:
-                        pass
+                        while True:
+                            print("Waiting for messages...")
+                            message = await self.read_stream.receive()
+                            json_str = message.model_dump_json()
+                            print(f"Received message from local server: {json_str}")
+                            await self.send_outgoing_message(message)
+                    finally:
+                        # Cancel the consumer when we exit the loop
+                        consumer_task.cancel()
+                        try:
+                            await consumer_task
+                        except asyncio.CancelledError:
+                            pass
 
         except Exception as e:
             logger.error(
@@ -155,6 +163,7 @@ async def send_message(self, message: types.JSONRPCMessage) -> None:
         await self._message_queue.put(message)
         logger.debug(f"Session {self.session_id} - message queued for processing")
 
+    @traced
     async def get_incoming_messages(self) -> None:
         """Get new messages from UiPath MCP Server."""
         response = self._uipath.api_client.request(
@@ -169,6 +178,7 @@ async def get_incoming_messages(self) -> None:
                 logger.info(f"Forwarding message to local MCP Server: {message}")
                 await self.send_message(json_message)
 
+    @traced
     async def send_outgoing_message(self, message: types.JSONRPCMessage) -> None:
         """Send new message to UiPath MCP Server."""
         response = self._uipath.api_client.request(
@@ -210,4 +220,5 @@ async def cleanup(self) -> None:
         self.read_stream = None
         self.write_stream = None
         self.mcp_session = None
+        wait_for_tracers()
         logger.info(f"Cleanup completed for session {self.session_id}")