refact: remove verbose session deletion code

Author: radugheo

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -133,37 +133,29 @@ def get_sessions(self) -> dict[str, SessionHealthInfo]:
         }
 
     async def remove_session(self, session_id: str, reason: str) -> None:
-        """Remove and stop a session by ID (SessionProvider protocol)."""
+        """Pop, stop, and clean up a single session (SessionProvider protocol)."""
         session_server = self._session_servers.pop(session_id, None)
-        if session_server is not None:
-            logger.warning(
-                f"Removing session {session_id}: {reason}"
-            )
-            try:
-                await session_server.stop()
-            except Exception:
-                logger.error(
-                    f"Error stopping session {session_id} during watchdog removal",
-                    exc_info=True,
-                )
-            await self._close_session_on_server(session_id)
+        if session_server is None:
+            return
+
+        logger.warning(f"Removing session {session_id}: {reason}")
 
-    async def _close_session_on_server(self, session_id: str) -> None:
-        """Notify the UiPath server to remove a session so it stops sending messages."""
         try:
-            await self._uipath.api_client.request_async(
-                "DELETE",
-                f"agenthub_/mcp/{self._folder_key}/{self.slug}",
-                headers={"mcp-session-id": session_id},
+            await session_server.stop()
+        except Exception:
+            logger.error(
+                f"Error stopping session {session_id}",
+                exc_info=True,
             )
-            logger.info(f"Notified server of session closure: {session_id}")
-        except HTTPStatusError as e:
-            if e.response.status_code == 404:
-                logger.info(f"Session {session_id} already removed server-side")
+
+        if session_server.output:
+            if self.sandboxed:
+                self._session_output = session_server.output
             else:
-                logger.error(f"Error closing session {session_id} on server: {e}")
-        except Exception as e:
-            logger.error(f"Error closing session {session_id} on server: {e}")
+                logger.info(f"Session {session_id} output: {session_server.output}")
+
+        if self.sandboxed:
+            self._cancel_event.set()
 
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Get schema for this MCP runtime.
@@ -366,11 +358,8 @@ async def _cleanup(self) -> None:
             await self._watchdog.stop()
             self._watchdog = None
 
-        for session_id, session_server in list(self._session_servers.items()):
-            try:
-                await session_server.stop()
-            except Exception as e:
-                logger.error(f"Error cleaning up session server {session_id}: {str(e)}")
+        for session_id in list(self._session_servers.keys()):
+            await self.remove_session(session_id, reason="runtime shutdown")
 
         # Stop the shared HTTP server process (streamable-http only)
         await self._stop_http_server_process()
@@ -396,26 +385,8 @@ async def _handle_signalr_session_closed(self, args: list[str]) -> None:
             return
 
         session_id = args[0]
-
         logger.info(f"Received closed signal for session {session_id}")
-
-        try:
-            session_server = self._session_servers.pop(session_id, None)
-            if session_server:
-                await session_server.stop()
-                if session_server.output:
-                    if self.sandboxed:
-                        self._session_output = session_server.output
-                    else:
-                        logger.info(
-                            f"Session {session_id} output: {session_server.output}"
-                        )
-            # If this is a sandboxed runtime for a specific session, cancel the execution
-            if self.sandboxed:
-                self._cancel_event.set()
-
-        except Exception as e:
-            logger.error(f"Error terminating session {session_id}: {str(e)}")
+        await self.remove_session(session_id, reason="server closed")
 
     async def _handle_signalr_message(self, args: list[str]) -> None:
         """
@@ -619,13 +590,9 @@ async def _monitor_http_server_process(self) -> None:
                 # Stop all HTTP sessions, they will fail on next request anyway
                 for session_id, session_server in list(self._session_servers.items()):
                     if isinstance(session_server, StreamableHttpSessionServer):
-                        try:
-                            await session_server.stop()
-                        except Exception as e:
-                            logger.error(
-                                f"Error stopping session {session_id} after process crash: {e}"
-                            )
-                        self._session_servers.pop(session_id, None)
+                        await self.remove_session(
+                            session_id, reason="http process crash"
+                        )
         except asyncio.CancelledError:
             pass
 
@@ -833,7 +800,7 @@ async def on_keep_alive_response(
                             health = s.get_health_info()
                             runtime_sessions[sid] = {
                                 "task_done": health.task_done,
-                                "active_requests": health.active_request_count,
+                                "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
@@ -849,7 +816,6 @@ async def on_keep_alive_response(
                             )
                             self._cancel_event.set()
 
-
                     if self._signalr_client:
                         logger.info("Sending keep-alive ping...")
                         await self._signalr_client.send(

src/uipath_mcp/_cli/_runtime/_session.py

@@ -41,7 +41,6 @@ class SessionHealthInfo:
     task_exception: BaseException | None
     last_activity_time: float
     queue_size: int
-    active_request_count: int
 
 
 class BaseSessionServer(ABC):
@@ -118,7 +117,6 @@ def get_health_info(self) -> SessionHealthInfo:
             task_exception=task_exception,
             last_activity_time=self._last_activity_time,
             queue_size=self._message_queue.qsize(),
-            active_request_count=len(self._active_requests),
         )
 
     async def on_message_received(self, request_id: str) -> None:
@@ -182,9 +180,7 @@ async def _relay_messages(self) -> None:
                         if self._last_request_id is not None:
                             await self._send_message(message, self._last_request_id)
                 except EndOfStream:
-                    logger.warning(
-                        f"Read stream closed for session {self._session_id}"
-                    )
+                    logger.warning(f"Read stream closed for session {self._session_id}")
                     break
                 except Exception as e:
                     if session_message:

src/uipath_mcp/_cli/_runtime/_watchdog.py

@@ -104,9 +104,9 @@ async def _check_sessions(self) -> None:
                                 f"with exception: {health.task_exception}"
                             )
                         else:
-                            logger.warning(
+                            logger.info(
                                 f"Watchdog: {transport} session {session_id} task "
-                                f"completed unexpectedly"
+                                f"completed, cleaning up"
                             )
                         await self._provider.remove_session(
                             session_id, reason="dead task"
@@ -135,4 +135,3 @@ async def _check_sessions(self) -> None:
             logger.info(
                 f"Watchdog check: {len(sessions)} session(s), {removed_count} removed"
             )
-

tests/test_watchdog.py

@@ -29,16 +29,16 @@ def _make_health(
     task_exception: BaseException | None = None,
     last_activity_time: float | None = None,
     queue_size: int = 0,
-    active_request_count: int = 0,
 ) -> SessionHealthInfo:
     return SessionHealthInfo(
         session_id=session_id,
         transport_type=transport_type,
         task_done=task_done,
         task_exception=task_exception,
-        last_activity_time=last_activity_time if last_activity_time is not None else time.monotonic(),
+        last_activity_time=last_activity_time
+        if last_activity_time is not None
+        else time.monotonic(),
         queue_size=queue_size,
-        active_request_count=active_request_count,
     )