Merge pull request #116 from UiPath/fix/error_handling_invalid_messages

fix: better error handling for invalid messages

Author: cristipufu

pyproject.toml

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

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -66,11 +66,12 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
             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
+            if self.context.job_id:
+                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')}/agenthub_/wsstunnel?slug={self._server.name}&runtimeId={self._runtime_id}"

src/uipath_mcp/_cli/_runtime/_session.py

@@ -119,26 +119,42 @@ async def _run_server(self, server_params: StdioServerParameters) -> None:
                     try:
                         while True:
                             # Get message from local server
-                            session_message = await self._read_stream.receive()
-                            message = session_message.message
-                            # For responses, determine which request_id to use
-                            if self._is_response(message):
-                                message_id = self._get_message_id(message)
-                                if message_id and message_id in self._active_requests:
-                                    # Use the stored request_id for this response
-                                    request_id = self._active_requests[message_id]
-                                    # Send with the matched request_id
-                                    await self._send_message(message, request_id)
-                                    # Clean up the mapping after use
-                                    del self._active_requests[message_id]
+                            session_message = None
+                            try:
+                                session_message = await self._read_stream.receive()
+                                message = session_message.message
+                                # For responses, determine which request_id to use
+                                if self._is_response(message):
+                                    message_id = self._get_message_id(message)
+                                    if message_id and message_id in self._active_requests:
+                                        # Use the stored request_id for this response
+                                        request_id = self._active_requests[message_id]
+                                        # Send with the matched request_id
+                                        await self._send_message(message, request_id)
+                                        # Clean up the mapping after use
+                                        del self._active_requests[message_id]
+                                    else:
+                                        # If no mapping found, use the last known request_id
+                                        await self._send_message(
+                                            message, self._last_request_id
+                                        )
                                 else:
-                                    # If no mapping found, use the last known request_id
-                                    await self._send_message(
-                                        message, self._last_request_id
-                                    )
-                            else:
-                                # For non-responses, use the last known request_id
-                                await self._send_message(message, self._last_request_id)
+                                    # For non-responses, use the last known request_id
+                                    await self._send_message(message, self._last_request_id)
+                            except Exception as e:
+                                if session_message:
+                                    logger.info(session_message)
+                                logger.error(
+                                    f"Error processing message for session {self._session_id}: {e}",
+                                    exc_info=True,
+                                )
+                                await self._send_message(
+                                    JSONRPCError(
+                                        code=-32000,
+                                        message=f"Error processing message: {session_message} {e} ",
+                                    ),
+                                    self._last_request_id,
+                                )
                     finally:
                         # Cancel the consumer when we exit the loop
                         consumer_task.cancel()
@@ -158,6 +174,7 @@ async def _run_server(self, server_params: StdioServerParameters) -> None:
                 self._server_stderr_output = stderr_temp.read().decode(
                     "utf-8", errors="replace"
                 )
+                logger.error(self._server_stderr_output)
                 # The context managers will handle cleanup of resources
 
     def _run_server_callback(self, task):