feat: add watchdog for runtime session handling

Author: radugheo

.gitignore

@@ -165,7 +165,11 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
+
+.vscode/
+
+.claude/
 
 # Ruff stuff:
 .ruff_cache/

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -37,8 +37,14 @@
 from .._utils._config import McpServer
 from ._context import UiPathServerType
 from ._exception import McpErrorCode, UiPathMcpRuntimeError
-from ._session import BaseSessionServer, StdioSessionServer, StreamableHttpSessionServer
+from ._session import (
+    BaseSessionServer,
+    SessionHealthInfo,
+    StdioSessionServer,
+    StreamableHttpSessionServer,
+)
 from ._token_refresh import TokenRefresher
+from ._watchdog import SessionWatchdog
 
 logger = logging.getLogger(__name__)
 tracer = trace.get_tracer(__name__)
@@ -86,6 +92,7 @@ def __init__(
         self._http_stderr_drain_task: asyncio.Task[None] | None = None
         self._http_server_stderr_lines: list[str] = []
         self._uipath = UiPath()
+        self._watchdog: SessionWatchdog | None = None
         self._token_refresher: TokenRefresher | None = None
         self._cleanup_done = False
 
@@ -118,6 +125,46 @@ def _validate_auth(self) -> None:
                 UiPathErrorCategory.SYSTEM,
             )
 
+    def get_sessions(self) -> dict[str, SessionHealthInfo]:
+        """Return health info for all active sessions (SessionProvider protocol)."""
+        return {
+            sid: session.get_health_info()
+            for sid, session in self._session_servers.items()
+        }
+
+    async def remove_session(self, session_id: str, reason: str) -> None:
+        """Remove and stop a session by ID (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)
+
+    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},
+            )
+            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")
+            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}")
+
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Get schema for this MCP runtime.
 
@@ -240,6 +287,9 @@ async def _run_server(self) -> UiPathRuntimeResult:
                 run_task = asyncio.create_task(self._signalr_client.run())
                 cancel_task = asyncio.create_task(self._cancel_event.wait())
                 self._keep_alive_task = asyncio.create_task(self._keep_alive())
+
+                self._watchdog = SessionWatchdog(self)
+                self._watchdog.start()
                 self._token_refresher.start()
 
                 try:
@@ -312,6 +362,10 @@ async def _cleanup(self) -> None:
             except asyncio.CancelledError:
                 pass
 
+        if self._watchdog:
+            await self._watchdog.stop()
+            self._watchdog = None
+
         for session_id, session_server in list(self._session_servers.items()):
             try:
                 await session_server.stop()
@@ -367,6 +421,10 @@ async def _handle_signalr_message(self, args: list[str]) -> None:
         """
         Handle incoming SignalR messages.
         """
+
+        if self._cleanup_done:
+            return
+
         if len(args) < 2:
             logger.error(f"Received invalid websocket message arguments: {args}")
             return
@@ -769,7 +827,15 @@ async def on_keep_alive_response(
                             logger.error(f"Error during keep-alive: {response.error}")
                             return
                         session_ids = response.result
-                        logger.info(f"Active sessions: {session_ids}")
+                        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": health.active_request_count,
+                            }
+                        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
@@ -783,6 +849,7 @@ 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

@@ -2,9 +2,12 @@
 import io
 import logging
 import tempfile
+import time
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Any
 
+from anyio import EndOfStream
 from mcp import StdioServerParameters, stdio_client
 from mcp.client.streamable_http import streamable_http_client
 from mcp.shared.message import SessionMessage
@@ -28,6 +31,19 @@
 RETRY_DELAY = 1
 
 
+@dataclass
+class SessionHealthInfo:
+    """Health information for a session, used by the watchdog."""
+
+    session_id: str
+    transport_type: str
+    task_done: bool
+    task_exception: BaseException | None
+    last_activity_time: float
+    queue_size: int
+    active_request_count: int
+
+
 class BaseSessionServer(ABC):
     """Base class with transport-agnostic message relay logic."""
 
@@ -48,9 +64,16 @@ def __init__(
         self._active_requests: dict[str, str] = {}
         self._last_request_id: str | None = None
         self._last_message_id: str | None = None
+        self._last_activity_time: float = time.monotonic()
         self._uipath = uipath
         self._mcp_tracer = McpTracer(tracer, logger)
 
+    @property
+    @abstractmethod
+    def transport_type(self) -> str:
+        """Returns the transport type identifier (e.g. 'stdio', 'streamable-http')."""
+        ...
+
     @property
     @abstractmethod
     def output(self) -> str | None:
@@ -79,8 +102,28 @@ async def stop(self) -> None:
         self._read_stream = None
         self._write_stream = None
 
+    def get_health_info(self) -> SessionHealthInfo:
+        """Return health information for this session."""
+        task_done = self._run_task.done() if self._run_task else True
+        task_exception: BaseException | None = None
+        if task_done and self._run_task is not None:
+            try:
+                task_exception = self._run_task.exception()
+            except (asyncio.CancelledError, asyncio.InvalidStateError):
+                pass
+        return SessionHealthInfo(
+            session_id=self._session_id,
+            transport_type=self.transport_type,
+            task_done=task_done,
+            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:
         """Get new incoming messages from UiPath MCP Server."""
+        self._last_activity_time = time.monotonic()
         for attempt in range(MAX_RETRIES + 1):
             try:
                 await self._get_messages_internal(request_id)
@@ -115,6 +158,7 @@ async def _relay_messages(self) -> None:
                         break
 
                     session_message = await self._read_stream.receive()
+                    self._last_activity_time = time.monotonic()
                     if isinstance(session_message, Exception):
                         logger.error(f"Received error: {session_message}")
                         continue
@@ -137,6 +181,11 @@ async def _relay_messages(self) -> None:
                         # For non-responses, use the last known request_id
                         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}"
+                    )
+                    break
                 except Exception as e:
                     if session_message:
                         logger.info(session_message)
@@ -292,6 +341,10 @@ class StdioSessionServer(BaseSessionServer):
 
     _server_stderr_output: str | None = None
 
+    @property
+    def transport_type(self) -> str:
+        return "stdio"
+
     @property
     def output(self) -> str | None:
         """Returns the captured stderr output from the MCP server process."""
@@ -345,6 +398,10 @@ async def _run_server(self, server_params: StdioServerParameters) -> None:
 class StreamableHttpSessionServer(BaseSessionServer):
     """Manages an HTTP connection to a shared streamable-http server for a specific session."""
 
+    @property
+    def transport_type(self) -> str:
+        return "streamable-http"
+
     @property
     def output(self) -> str | None:
         """Returns captured output from the server process, if any."""