fix: runtime error handliing

Author: radugheo

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -363,7 +363,6 @@ async def _register(self) -> None:
         """Register the MCP server with UiPath."""
         server = cast(McpServer, self._server)
 
-        initialization_successful = False
         tools_result: Optional[ListToolsResult] = None
         server_stderr_output = ""
         env_vars = server.env
@@ -375,6 +374,7 @@ async def _register(self) -> None:
                 if name not in env_vars:
                     env_vars[name] = value
 
+        # Create a temporary session to get tools
         try:
             # Create a temporary session to get tools
             server_params = StdioServerParameters(
@@ -387,49 +387,34 @@ async def _register(self) -> None:
             # Use a temporary file to capture stderr
             with tempfile.TemporaryFile(mode="w+b") as stderr_temp_binary:
                 stderr_temp = io.TextIOWrapper(stderr_temp_binary, encoding="utf-8")
-                async with stdio_client(server_params, errlog=stderr_temp) as (
-                    read,
-                    write,
-                ):
-                    async with ClientSession(read, write) as session:
-                        logger.info("Initializing client session...")
-                        # Try to initialize with timeout
-                        try:
+                try:
+                    async with stdio_client(server_params, errlog=stderr_temp) as (
+                        read,
+                        write,
+                    ):
+                        async with ClientSession(read, write) as session:
+                            logger.info("Initializing client session...")
                             await asyncio.wait_for(session.initialize(), timeout=30)
-                            initialization_successful = True
                             logger.info("Initialization successful")
-
-                            # Only proceed if initialization was successful
                             tools_result = await session.list_tools()
-                            # logger.info(tools_result)
-                        except asyncio.TimeoutError:
-                            logger.error("Initialization timed out")
-                            # Capture stderr output here, after the timeout
-                            stderr_temp.seek(0)
-                            server_stderr_output = stderr_temp.read()
-
-        except* Exception as eg:
-            for e in eg.exceptions:
-                logger.error(
-                    f"Unexpected error: {e}",
-                    exc_info=True,
-                )
-
-        # Now that we're outside the context managers, check if initialization succeeded
-        if not initialization_successful:
+                except Exception:
+                    stderr_temp.seek(0)
+                    server_stderr_output = stderr_temp.read()
+                    raise
+        except Exception as e:
             await self._on_runtime_abort()
-            error_message = "The server process failed to initialize. Verify environment variables are set correctly."
+            error_detail = "The server process failed to initialize."
             if server_stderr_output:
-                error_message += f"\nServer error output:\n{server_stderr_output}"
+                error_detail += f"\n\nServer error output:\n{server_stderr_output}"
+
             raise UiPathMcpRuntimeError(
                 McpErrorCode.INITIALIZATION_ERROR,
                 "Server initialization failed",
-                error_message,
+                error_detail,
                 UiPathErrorCategory.DEPLOYMENT,
-            )
+            ) from e
 
-        # If we got here, initialization was successful and we have the tools
-        # Now continue with registration
+        # If we got here, initialization was successful - continue with registration
         logger.info("Registering server runtime ...")
         try:
             if not tools_result: