Merge pull request #8 from UiPath/fix/samples

feat: handle close session

Author: cristipufu

pyproject.toml

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

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -53,11 +53,14 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
 
             # Set up SignalR client
             signalr_url = (
-                f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self.server.name}"
+                f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self.server.name}&jobKey={self.context.job_id}"
             )
 
+            self.cancel_event = asyncio.Event()
+
             self.signalr_client = SignalRClient(signalr_url)
             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)
@@ -68,7 +71,21 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
             # Keep the runtime alive
             # Start SignalR client and keep it running (this is a blocking call)
             logger.info("Starting SignalR client...")
-            await 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())
+
+            # Wait for either the run to complete or cancellation
+            done, pending = await asyncio.wait(
+                [run_task, cancel_task],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Cancel any pending tasks
+            for task in pending:
+                task.cancel()
 
             return UiPathRuntimeResult()
 
@@ -99,6 +116,26 @@ async def validate(self) -> None:
                 UiPathErrorCategory.DEPLOYMENT,
             )
 
+    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}")
+            return
+
+        session_id = args[0]
+
+        logger.info(f"Received closed signal for session {session_id}")
+
+        try:
+            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:
         """
         Handle incoming SignalR messages.
@@ -214,8 +251,9 @@ async def cleanup(self) -> None:
 
         self.session_servers.clear()
 
-        # Close SignalR connection
-        # self.signalr_client
+        if self.signalr_client:
+            # Close the SignalR connection
+            await self.signalr_client._transport._ws.close()
 
         # Add a small delay to allow the server to shut down gracefully
         if sys.platform == "win32":

src/uipath_mcp/_cli/_utils/_config.py

@@ -69,7 +69,16 @@ def get_servers(self) -> List[McpServer]:
         return list(self._servers.values())
 
     def get_server(self, name: str) -> Optional[McpServer]:
-        """Get a server model by name."""
+        """
+        Get a server model by name.
+        If there's only one server available, return that one regardless of name.
+        Otherwise, look up the server by the provided name.
+        """
+        # If there's only one server, return it
+        if len(self._servers) == 1:
+            return next(iter(self._servers.values()))
+
+        # Otherwise, fall back to looking up by name
         return self._servers.get(name)
 
     def get_server_names(self) -> List[str]: