Merge pull request #100 from UiPath/fix/tracing

feat(tracing): fix tracing

Author: ionmincu

pyproject.toml

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

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -10,14 +10,16 @@
 from mcp import ClientSession, StdioServerParameters, stdio_client
 from mcp.types import JSONRPCResponse
 from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from pysignalr.client import CompletionMessage, SignalRClient
 from uipath import UiPath
 from uipath._cli._runtime._contracts import (
     UiPathBaseRuntime,
     UiPathErrorCategory,
     UiPathRuntimeResult,
 )
-from uipath.tracing import wait_for_tracers
+from uipath.tracing import LlmOpsHttpExporter
 
 from .._utils._config import McpServer
 from ._context import UiPathMcpRuntimeContext, UiPathServerType
@@ -62,6 +64,12 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         try:
             if self._server is None:
                 return None
+            
+            self.trace_provider = TracerProvider()
+            trace.set_tracer_provider(self.trace_provider)
+            self.trace_provider.add_span_processor(
+                BatchSpanProcessor(LlmOpsHttpExporter())
+            )  # type: ignore
 
             # Set up SignalR client
             signalr_url = f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self._server.name}&runtimeId={self._runtime_id}"
@@ -125,7 +133,7 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 UiPathErrorCategory.USER,
             ) from e
         finally:
-            wait_for_tracers()
+            self.trace_provider.shutdown()
 
     async def validate(self) -> None:
         """Validate runtime inputs and load MCP server configuration."""