refactor: solve runtime code inconsistencies

Author: radugheo

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -295,8 +295,8 @@ async def _run_server(self) -> UiPathRuntimeResult:
                     )
                     self._cancel_event.set()
                 finally:
-                    # Cancel any pending tasks gracefully
-                    for task in [run_task, cancel_task, self._keep_alive_task]:
+                    # Cancel pending tasks
+                    for task in [run_task, cancel_task]:
                         if task and not task.done():
                             task.cancel()
                             try:
@@ -322,7 +322,7 @@ async def _run_server(self) -> UiPathRuntimeResult:
         except Exception as e:
             if isinstance(e, UiPathMcpRuntimeError):
                 raise
-            detail = f"Error: {str(e)}"
+            detail = f"Error: {e}"
             raise UiPathMcpRuntimeError(
                 UiPathErrorCode.EXECUTION_ERROR,
                 "MCP Runtime execution failed",
@@ -370,16 +370,17 @@ async def _cleanup(self) -> None:
                 try:
                     await transport._ws.close()
                 except Exception as e:
-                    logger.error(f"Error closing SignalR WebSocket: {str(e)}")
+                    logger.error(f"Error closing SignalR WebSocket: {e}")
 
         # Add a small delay to allow the server to shut down gracefully
         if sys.platform == "win32":
             await asyncio.sleep(0.5)
 
     async def _handle_signalr_session_closed(self, args: list[str]) -> None:
-        """
-        Handle session closed by server.
-        """
+        """Handle session closed by server."""
+        if self._cleanup_done:
+            return
+
         if len(args) < 1:
             logger.error(f"Received invalid websocket message arguments: {args}")
             return
@@ -389,10 +390,7 @@ async def _handle_signalr_session_closed(self, args: list[str]) -> None:
         await self.remove_session(session_id, reason="server closed")
 
     async def _handle_signalr_message(self, args: list[str]) -> None:
-        """
-        Handle incoming SignalR messages.
-        """
-
+        """Handle incoming SignalR messages."""
         if self._cleanup_done:
             return
 
@@ -421,7 +419,7 @@ async def _handle_signalr_message(self, args: list[str]) -> None:
                     await session_server.start()
                 except Exception as e:
                     logger.error(
-                        f"Error starting session server for session {session_id}: {str(e)}"
+                        f"Error starting session server for session {session_id}: {e}"
                     )
                     await self._on_session_start_error(session_id)
                     raise
@@ -435,7 +433,7 @@ async def _handle_signalr_message(self, args: list[str]) -> None:
 
         except Exception as e:
             logger.error(
-                f"Error handling websocket notification for session {session_id}: {str(e)}"
+                f"Error handling websocket notification for session {session_id}: {e}"
             )
 
     async def _handle_signalr_error(self, error: Any) -> None:
@@ -450,17 +448,21 @@ async def _handle_signalr_close(self) -> None:
         """Handle SignalR connection close event."""
         logger.info("Websocket connection closed.")
 
-    async def _start_http_server_process(self) -> None:
-        """Spawn the streamable-http server process.
-
-        The process is started once and shared across all sessions.
-        """
+    def _get_server_env(self) -> dict[str, str]:
+        """Return server env vars, with os.environ merged in for Coded servers."""
         env_vars = self._server.env.copy()
         if self.server_type is UiPathServerType.Coded:
             for name, value in os.environ.items():
                 if name not in env_vars:
                     env_vars[name] = value
+        return env_vars
 
+    async def _start_http_server_process(self) -> None:
+        """Spawn the streamable-http server process.
+
+        The process is started once and shared across all sessions.
+        """
+        env_vars = self._get_server_env()
         merged_env = {**os.environ, **env_vars} if env_vars else None
         self._http_server_stderr_lines = []
         self._http_server_process = await asyncio.create_subprocess_exec(
@@ -501,7 +503,12 @@ async def _wait_for_http_server_ready(
 
         url = self._server.url
         if not url:
-            raise ValueError("streamable-http transport requires url in config")
+            raise UiPathMcpRuntimeError(
+                McpErrorCode.CONFIGURATION_ERROR,
+                "Missing URL for streamable-http server",
+                "Please specify a 'url' in the server configuration for streamable-http transport.",
+                UiPathErrorCategory.SYSTEM,
+            )
 
         for attempt in range(max_retries):
             # Check if process has crashed
@@ -602,14 +609,6 @@ async def _register(self) -> None:
         initialization_successful = False
         tools_result: ListToolsResult | None = None
         server_stderr_output = ""
-        env_vars = self._server.env
-
-        # if server is Coded, include environment variables
-        if self.server_type is UiPathServerType.Coded:
-            for name, value in os.environ.items():
-                # config env variables should have precedence over system ones
-                if name not in env_vars:
-                    env_vars[name] = value
 
         try:
             if self._server.is_streamable_http:
@@ -649,7 +648,7 @@ async def _register(self) -> None:
                 server_params = StdioServerParameters(
                     command=self._server.command,
                     args=self._server.args,
-                    env=env_vars,
+                    env=self._get_server_env(),
                 )
 
                 with tempfile.TemporaryFile(mode="w+b") as stderr_temp_binary:
@@ -779,49 +778,39 @@ async def _on_session_start_error(self, session_id: str) -> None:
                 f"Error sending session dispose signal to UiPath MCP Server: {e}"
             )
 
+    async def _on_keep_alive_response(self, response: CompletionMessage) -> None:
+        """Handle keep-alive response: log session state, detect orphaned sandboxed runtimes."""
+        if response.error:
+            logger.error(f"Error during keep-alive: {response.error}")
+            return
+        session_ids = response.result
+        logger.info(f"Server active sessions: {session_ids}")
+        runtime_sessions = {}
+        for sid, s in self._session_servers.items():
+            health = s.get_health_info()
+            runtime_sessions[sid] = {
+                "task_done": health.task_done,
+                "active_requests": len(s._active_requests),
+            }
+        logger.info(f"Runtime active sessions: {runtime_sessions}")
+        # If there are no active sessions and this is a sandbox environment
+        # We need to cancel the runtime
+        # eg: when user kills the agent that triggered the runtime, before we subscribe to events
+        if not session_ids and self.sandboxed and not self._cancel_event.is_set():
+            logger.warning("No active sessions, cancelling sandboxed runtime...")
+            self._cancel_event.set()
+
     async def _keep_alive(self) -> None:
-        """
-        Heartbeat to keep the runtime available.
-        """
+        """Heartbeat to keep the runtime available."""
         try:
             while not self._cancel_event.is_set():
                 try:
-
-                    async def on_keep_alive_response(
-                        response: CompletionMessage,
-                    ) -> None:
-                        if response.error:
-                            logger.error(f"Error during keep-alive: {response.error}")
-                            return
-                        session_ids = response.result
-                        logger.info(f"Server active sessions: {session_ids}")
-                        runtime_sessions = {}
-                        for sid, s in self._session_servers.items():
-                            health = s.get_health_info()
-                            runtime_sessions[sid] = {
-                                "task_done": health.task_done,
-                                "active_requests": len(s._active_requests),
-                            }
-                        logger.info(f"Runtime active sessions: {runtime_sessions}")
-                        # If there are no active sessions and this is a sandbox environment
-                        # We need to cancel the runtime
-                        # eg: when user kills the agent that triggered the runtime, before we subscribe to events
-                        if (
-                            not session_ids
-                            and self.sandboxed
-                            and not self._cancel_event.is_set()
-                        ):
-                            logger.error(
-                                "No active sessions, cancelling sandboxed runtime..."
-                            )
-                            self._cancel_event.set()
-
                     if self._signalr_client:
                         logger.info("Sending keep-alive ping...")
                         await self._signalr_client.send(
                             method="OnKeepAlive",
                             arguments=[],
-                            on_invocation=on_keep_alive_response,  # type: ignore
+                            on_invocation=self._on_keep_alive_response,  # type: ignore
                         )
                     else:
                         logger.error("SignalR client not initialized during keep-alive")
@@ -839,9 +828,7 @@ async def on_keep_alive_response(
             raise
 
     async def _on_runtime_abort(self) -> None:
-        """
-        Sends a runtime abort signalr to terminate all connected sessions.
-        """
+        """Send a runtime abort request to terminate all connected sessions."""
         try:
             response = await self._uipath.api_client.request_async(
                 "POST",
@@ -854,7 +841,7 @@ async def _on_runtime_abort(self) -> None:
                 )
             else:
                 logger.error(
-                    f"Error sending runtime abort signalr to UiPath MCP Server: {response.status_code} - {response.text}"
+                    f"Error sending runtime abort to UiPath MCP Server: {response.status_code} - {response.text}"
                 )
         except Exception as e:
             logger.error(