fix: create tracing spans

Author: cristipufu

pyproject.toml

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

src/uipath_mcp/_cli/_runtime/_logger.py

@@ -1,19 +1,18 @@
-import logging
-import sys
+import os
 
 
 class LoggerAdapter:
-    def __init__(self, logger, level=logging.INFO):
-        self.logger = logger
-        self.level = level
+    def __init__(self):
+        self.devnull = open(os.devnull, "w")
 
     def write(self, message):
-        if message and not message.isspace():
-            self.logger.log(self.level, message.rstrip())
         return len(message)
 
     def flush(self):
-        pass
+        self.devnull.flush()
 
     def fileno(self):
-        return sys.stderr.fileno()
+        return self.devnull.fileno()
+
+    def close(self):
+        self.devnull.close()

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -6,6 +6,7 @@
 
 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
 from uipath._cli._runtime._contracts import (
@@ -20,6 +21,7 @@
 from ._session import SessionServer
 
 logger = logging.getLogger(__name__)
+tracer = trace.get_tracer(__name__)
 
 
 class UiPathMcpRuntime(UiPathBaseRuntime):
@@ -56,41 +58,48 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
 
             self.cancel_event = asyncio.Event()
 
-            self.signalr_client = SignalRClient(
-                signalr_url,
-                headers={
-                    "X-UiPath-Internal-TenantId": self.context.trace_context.tenant_id,
-                    "X-UiPath-Internal-AccountId": self.context.trace_context.org_id,
-                },
-            )
-            self.signalr_client.on("MessageReceived", self.handle_signalr_message)
-            self.signalr_client.on("SessionClosed", self.handle_signalr_session_closed)
-            self.signalr_client.on_error(self.handle_signalr_error)
-            self.signalr_client.on_open(self.handle_signalr_open)
-            self.signalr_client.on_close(self.handle_signalr_close)
-
-            # Register the server with UiPath MCP Server
-            await self._register()
-
-            # Keep the runtime alive
-            # Start SignalR client and keep it running (this is a blocking call)
-            logger.info("Starting websocket client...")
-
-            run_task = asyncio.create_task(self.signalr_client.run())
-
-            # Set up a task to wait for cancellation
-            cancel_task = asyncio.create_task(self.cancel_event.wait())
-
-            # Wait for either the run to complete or cancellation
-            done, pending = await asyncio.wait(
-                [run_task, cancel_task], return_when=asyncio.FIRST_COMPLETED
-            )
-
-            # Cancel any pending tasks
-            for task in pending:
-                task.cancel()
-
-            return UiPathRuntimeResult()
+            with tracer.start_as_current_span("MCP Server") as root_span:
+                root_span.set_attribute("session_id", self.server.session_id)
+                root_span.set_attribute("command", self.server.command)
+                root_span.set_attribute("args", self.server.args)
+                root_span.set_attribute("type", self.server.type)
+                self.signalr_client = SignalRClient(
+                    signalr_url,
+                    headers={
+                        "X-UiPath-Internal-TenantId": self.context.trace_context.tenant_id,
+                        "X-UiPath-Internal-AccountId": self.context.trace_context.org_id,
+                    },
+                )
+                self.signalr_client.on("MessageReceived", self.handle_signalr_message)
+                self.signalr_client.on(
+                    "SessionClosed", self.handle_signalr_session_closed
+                )
+                self.signalr_client.on_error(self.handle_signalr_error)
+                self.signalr_client.on_open(self.handle_signalr_open)
+                self.signalr_client.on_close(self.handle_signalr_close)
+
+                # Register the server with UiPath MCP Server
+                await self._register()
+
+                # Keep the runtime alive
+                # Start SignalR client and keep it running (this is a blocking call)
+                logger.info("Starting websocket client...")
+
+                run_task = asyncio.create_task(self.signalr_client.run())
+
+                # Set up a task to wait for cancellation
+                cancel_task = asyncio.create_task(self.cancel_event.wait())
+
+                # Wait for either the run to complete or cancellation
+                done, pending = await asyncio.wait(
+                    [run_task, cancel_task], return_when=asyncio.FIRST_COMPLETED
+                )
+
+                # Cancel any pending tasks
+                for task in pending:
+                    task.cancel()
+
+                return UiPathRuntimeResult()
 
         except Exception as e:
             if isinstance(e, UiPathMcpRuntimeError):

src/uipath_mcp/_cli/_runtime/_session.py

@@ -6,10 +6,11 @@
 from mcp.client.stdio import stdio_client
 from opentelemetry import trace
 from uipath import UiPath
-from uipath.tracing import traced, wait_for_tracers
+from uipath.tracing import wait_for_tracers
 
 from .._utils._config import McpServer
 from ._logger import LoggerAdapter
+from ._tracer import McpTracer
 
 logger = logging.getLogger(__name__)
 tracer = trace.get_tracer(__name__)
@@ -28,6 +29,7 @@ def __init__(self, server_config: McpServer, session_id: str):
         self.context_task = None
         self._message_queue = asyncio.Queue()
         self._uipath = UiPath()
+        self._mcp_tracer = McpTracer(tracer, logger)
 
     async def start(self) -> None:
         """Start the server process in a separate task."""
@@ -59,39 +61,33 @@ async def _run_server(self, server_params: StdioServerParameters) -> None:
         """Run the server in proper context managers."""
         logger.info(f"Starting server process for session {self.session_id}")
         try:
-            stderr_adapter = LoggerAdapter(logger)
-            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")
+            stderr_null = LoggerAdapter(logger)
 
-                    logger.info(f"Session {self.session_id} - MCP session initialized")
+            async with stdio_client(server_params, errlog=stderr_null) as (
+                read,
+                write,
+            ):
+                self.read_stream, self.write_stream = read, write
+                logger.info(f"Session {self.session_id} - stdio client started")
 
-                    # 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
+                # 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()
                     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()
-                        try:
-                            await consumer_task
-                        except asyncio.CancelledError:
-                            pass
+                        await consumer_task
+                    except asyncio.CancelledError:
+                        pass
 
         except Exception as e:
             logger.error(
@@ -163,7 +159,6 @@ 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(
@@ -175,23 +170,37 @@ async def get_incoming_messages(self) -> None:
             for message in messages:
                 logger.info(f"Incoming message from UiPath MCP Server: {message}")
                 json_message = types.JSONRPCMessage.model_validate(message)
-                logger.info(f"Forwarding message to local MCP Server: {message}")
-                await self.send_message(json_message)
+                with self._mcp_tracer.create_span_for_message(
+                        json_message,
+                        session_id=self.session_id,
+                        server_name=self.server_config.name
+                    ) as _:
+                    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(
-            "POST",
-            f"mcp_/mcp/{self.server_config.name}/out/message?sessionId={self.session_id}",
-            json=message.model_dump(),
-        )
-        if response.status_code == 202:
-            logger.info(f"Outgoing message sent to UiPath MCP Server: {message}")
-        else:
-            logger.error(
-                f"Failed to send outgoing message to UiPath MCP Server: {response.status_code} {response.text}"
-            )
+        with self._mcp_tracer.create_span_for_message(
+            message,
+            session_id=self.session_id,
+            server_name=self.server_config.name
+        ) as span:
+            try:
+                response = self._uipath.api_client.request(
+                    "POST",
+                    f"mcp_/mcp/{self.server_config.name}/out/message?sessionId={self.session_id}",
+                    json=message.model_dump(),
+                )
+
+                span.set_attribute("http.status_code", response.status_code)
+
+                if response.status_code == 202:
+                    logger.info(f"Outgoing message sent to UiPath MCP Server: {message}")
+                else:
+                    self._mcp_tracer.record_http_error(span, response.status_code, response.text)
+            except Exception as e:
+                self._mcp_tracer.record_exception(span, e)
+                raise
 
     async def cleanup(self) -> None:
         """Clean up resources and stop the server."""