fix: refactor, namings, private vars

Author: cristipufu

pyproject.toml

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

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -29,21 +29,22 @@
 
 class UiPathMcpRuntime(UiPathBaseRuntime):
     """
-    A runtime class implementing the async context manager protocol.
-    This allows using the class with 'async with' statements.
+    A runtime class for hosting UiPath MCP servers.
     """
 
     def __init__(self, context: UiPathMcpRuntimeContext):
         super().__init__(context)
         self.context: UiPathMcpRuntimeContext = context
-        self.server: Optional[McpServer] = None
-        self.signalr_client: Optional[SignalRClient] = None
-        self.session_servers: Dict[str, SessionServer] = {}
+        self._server: Optional[McpServer] = None
+        self._signalr_client: Optional[SignalRClient] = None
+        self._session_servers: Dict[str, SessionServer] = {}
+        self._session_outputs: Dict[str, str] = {}
+        self._cancel_event = asyncio.Event()
         self._uipath = UiPath()
 
     async def execute(self) -> Optional[UiPathRuntimeResult]:
         """
-        Start the runtime and connect to SignalR.
+        Start the MCP Server runtime.
 
         Returns:
             Dictionary with execution results
@@ -54,46 +55,41 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         await self.validate()
 
         try:
-            if self.server is None:
+            if self._server is None:
                 return None
 
             # Set up SignalR client
-            signalr_url = f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self.server.name}&sessionId={self.server.session_id}"
+            signalr_url = f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self._server.name}&sessionId={self._server.session_id}"
 
-            self.cancel_event = asyncio.Event()
-
-            with tracer.start_as_current_span(self.server.name) 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)
+            with tracer.start_as_current_span(self._server.name) 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("span_type", "MCP Server")
-                self.signalr_client = SignalRClient(
+                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("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)
+                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
+                # Register the local 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())
+                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())
+                cancel_task = asyncio.create_task(self._cancel_event.wait())
 
+                # Keep the runtime alive
                 # Wait for either the run to complete or cancellation
                 done, pending = await asyncio.wait(
                     [run_task, cancel_task], return_when=asyncio.FIRST_COMPLETED
@@ -103,24 +99,14 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 for task in pending:
                     task.cancel()
 
-                session_outputs = {}
-                for session_id, session_server in self.session_servers.items():
-                    try:
-                        await session_server.cleanup()
-                        stderr_output = session_server.get_server_stderr()
-                        if stderr_output:
-                            session_outputs[session_id] = stderr_output
-                    except Exception as e:
-                        logger.error(f"Error stopping session {session_id}: {str(e)}")
-
                 output_result = {}
-                if len(session_outputs) == 1:
-                    # If there's only one session, use a single "output" key
-                    first_session_id = next(iter(session_outputs))
-                    output_result["content"] = session_outputs[first_session_id]
-                elif session_outputs:
-                    # If there are multiple sessions, use the sessionId as the key
-                    output_result = session_outputs
+                if len(self._session_outputs) == 1:
+                    # If there's only one session, use a single "content" key
+                    single_session_id = next(iter(self._session_outputs))
+                    output_result["content"] = self._session_outputs[single_session_id]
+                elif self._session_outputs:
+                    # If there are multiple sessions, use the session_id as the key
+                    output_result = self._session_outputs
 
                 self.context.result = UiPathRuntimeResult(output=output_result)
 
@@ -129,55 +115,72 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         except Exception as e:
             if isinstance(e, UiPathMcpRuntimeError):
                 raise
-
             detail = f"Error: {str(e)}"
-
             raise UiPathMcpRuntimeError(
                 "EXECUTION_ERROR",
                 "MCP Runtime execution failed",
                 detail,
                 UiPathErrorCategory.USER,
             ) from e
-
         finally:
             wait_for_tracers()
 
     async def validate(self) -> None:
         """Validate runtime inputs and load MCP server configuration."""
-        self.server = self.context.config.get_server(self.context.entrypoint)
-        if not self.server:
+        self._server = self.context.config.get_server(self.context.entrypoint)
+        if not self._server:
             raise UiPathMcpRuntimeError(
                 "SERVER_NOT_FOUND",
                 "MCP server not found",
                 f"Server '{self.context.entrypoint}' not found in configuration",
                 UiPathErrorCategory.DEPLOYMENT,
             )
 
-    async def handle_signalr_session_closed(self, args: list) -> None:
+    async def cleanup(self) -> None:
+        """Clean up all resources."""
+        if self._signalr_client and hasattr(self._signalr_client, "_transport"):
+            transport = self._signalr_client._transport
+            if transport and hasattr(transport, "_ws") and transport._ws:
+                try:
+                    await transport._ws.close()
+                except Exception as e:
+                    logger.error(f"Error closing SignalR WebSocket: {str(e)}")
+
+        # Add a small delay to allow the server to shut down gracefully
+        if sys.platform == "win32":
+            await asyncio.sleep(0.1)
+
+    async def _handle_signalr_session_closed(self, args: list) -> None:
         """
         Handle session closed by server.
         """
         if len(args) < 1:
-            logger.error(f"Received invalid SignalR message arguments: {args}")
+            logger.error(f"Received invalid websocket message arguments: {args}")
             return
 
         session_id = args[0]
 
         logger.info(f"Received closed signal for session {session_id}")
 
         try:
-            self.cancel_event.set()
+            session_server = self._session_servers.pop(session_id, None)
+            if session_server:
+                await session_server.cleanup()
+                if session_server.output:
+                    self._session_outputs[session_id] = session_server.output
+
+            if len(self._session_servers) == 0:
+                self._cancel_event.set()
 
         except Exception as e:
             logger.error(f"Error terminating session {session_id}: {str(e)}")
 
-    async def handle_signalr_message(self, args: list) -> None:
+    async def _handle_signalr_message(self, args: list) -> None:
         """
         Handle incoming SignalR messages.
-        The SignalR client will call this with the arguments from the server.
         """
         if len(args) < 1:
-            logger.error(f"Received invalid SignalR message arguments: {args}")
+            logger.error(f"Received invalid websocket message arguments: {args}")
             return
 
         session_id = args[0]
@@ -186,50 +189,50 @@ async def handle_signalr_message(self, args: list) -> None:
 
         try:
             # Check if we have a session server for this session_id
-            if session_id not in self.session_servers:
+            if session_id not in self._session_servers:
                 # Create and start a new session server
-                session_server = SessionServer(self.server, session_id)
-                self.session_servers[session_id] = session_server
+                session_server = SessionServer(self._server, session_id)
+                self._session_servers[session_id] = session_server
                 await session_server.start()
 
             # Get the session server for this session
-            session_server = self.session_servers[session_id]
+            session_server = self._session_servers[session_id]
 
             # Forward the message to the session's MCP server
-            await session_server.get_incoming_messages()
+            await session_server.on_message_received()
 
         except Exception as e:
             logger.error(
                 f"Error handling websocket notification for session {session_id}: {str(e)}"
             )
 
-    async def handle_signalr_error(self, error: Any) -> None:
+    async def _handle_signalr_error(self, error: Any) -> None:
         """Handle SignalR errors."""
-        logger.error(f"SignalR error: {error}")
+        logger.error(f"Websocket error: {error}")
 
-    async def handle_signalr_open(self) -> None:
+    async def _handle_signalr_open(self) -> None:
         """Handle SignalR connection open event."""
 
         logger.info("Websocket connection established.")
-        if self.server.session_id:
+        if self._server.session_id:
             try:
-                session_server = SessionServer(self.server, self.server.session_id)
+                session_server = SessionServer(self._server, self._server.session_id)
                 await session_server.start()
-                self.session_servers[self.server.session_id] = session_server
-                await session_server.get_incoming_messages()
+                self._session_servers[self._server.session_id] = session_server
+                await session_server.on_message_received()
             except Exception as e:
-                await self._dispose_session()
+                await self._on_initialization_failure()
                 logger.error(f"Error starting session server: {str(e)}")
 
-    async def handle_signalr_close(self) -> None:
+    async def _handle_signalr_close(self) -> None:
         """Handle SignalR connection close event."""
-        logger.info("SignalR connection closed.")
+        logger.info("Websocket connection closed.")
         # Clean up all session servers when the connection closes
         await self.cleanup()
 
     async def _register(self) -> None:
-        """Register the MCP server type with UiPath."""
-        logger.info(f"Registering MCP server type: {self.server.name}")
+        """Register the MCP server with UiPath."""
+        logger.info(f"Registering MCP server: {self._server.name}")
 
         initialization_successful = False
         tools_result = None
@@ -238,9 +241,9 @@ async def _register(self) -> None:
         try:
             # Create a temporary session to get tools
             server_params = StdioServerParameters(
-                command=self.server.command,
-                args=self.server.args,
-                env=self.server.env,
+                command=self._server.command,
+                args=self._server.args,
+                env=self._server.env,
             )
 
             # Start a temporary stdio client to get tools
@@ -275,7 +278,7 @@ async def _register(self) -> None:
 
         # Now that we're outside the context managers, check if initialization succeeded
         if not initialization_successful:
-            await self._dispose_session()
+            await self._on_initialization_failure()
             error_message = "The server process failed to initialize. Verify environment variables are set correctly."
             if server_stderr_output:
                 error_message += f"\nServer error output:\n{server_stderr_output}"
@@ -291,8 +294,8 @@ async def _register(self) -> None:
         try:
             client_info = {
                 "server": {
-                    "Name": self.server.name,
-                    "Slug": self.server.name,
+                    "Name": self._server.name,
+                    "Slug": self._server.name,
                     "Version": "1.0.0",
                     "Type": 1,
                 },
@@ -311,7 +314,7 @@ async def _register(self) -> None:
             # Register with UiPath MCP Server
             self._uipath.api_client.request(
                 "POST",
-                f"mcp_/api/servers-with-tools/{self.server.name}",
+                f"mcp_/api/servers-with-tools/{self._server.name}",
                 json=client_info,
             )
             logger.info("Registered MCP Server type successfully")
@@ -324,25 +327,28 @@ async def _register(self) -> None:
                 UiPathErrorCategory.SYSTEM,
             ) from e
 
-    async def _dispose_session(self) -> None:
+    async def _on_initialization_failure(self) -> None:
         """Dispose of the session on the server."""
+        if self._server.session_id is None:
+            return
+
         try:
             response = self._uipath.api_client.request(
                 "POST",
-                f"mcp_/mcp/{self.server.name}/out/message?sessionId={self.server.session_id}",
+                f"mcp_/mcp/{self._server.name}/out/message?sessionId={self._server.session_id}",
                 json=types.JSONRPCResponse(
                     jsonrpc="2.0",
                     id=0,
                     result={
                         "protocolVersion": "initiliaze-failure",
                         "capabilities": {},
-                        "serverInfo": {"name": self.server.name, "version": "1.0"},
+                        "serverInfo": {"name": self._server.name, "version": "1.0"},
                     },
                 ).model_dump(),
             )
             if response.status_code == 202:
                 logger.info(
-                    f"Sent outgoing session dispose message to UiPath MCP Server: {self.server.session_id}"
+                    f"Sent outgoing session dispose message to UiPath MCP Server: {self._server.session_id}"
                 )
             else:
                 logger.error(
@@ -352,21 +358,3 @@ async def _dispose_session(self) -> None:
             logger.error(
                 f"Error sending session dispose signal to UiPath MCP Server: {e}"
             )
-
-    async def cleanup(self) -> None:
-        """Clean up all resources."""
-        logger.info("Cleaning up all resources")
-
-        self.session_servers.clear()
-
-        if self.signalr_client and hasattr(self.signalr_client, "_transport"):
-            transport = self.signalr_client._transport
-            if transport and hasattr(transport, "_ws") and transport._ws:
-                try:
-                    await transport._ws.close()
-                except Exception as e:
-                    logger.error(f"Error closing SignalR WebSocket: {str(e)}")
-
-        # Add a small delay to allow the server to shut down gracefully
-        if sys.platform == "win32":
-            await asyncio.sleep(0.1)