feat: style and cleanup

Author: edis-uipath

src/uipath_mcp/_cli/_runtime/_exception.py

@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Optional, Union
 
 from uipath.runtime.errors import (
     UiPathBaseRuntimeError,
@@ -20,11 +19,11 @@ class UiPathMcpRuntimeError(UiPathBaseRuntimeError):
 
     def __init__(
         self,
-        code: Union[McpErrorCode, UiPathErrorCode],
+        code: McpErrorCode | UiPathErrorCode,
         title: str,
         detail: str,
         category: UiPathErrorCategory = UiPathErrorCategory.UNKNOWN,
-        status: Optional[int] = None,
+        status: int | None = None,
     ):
         super().__init__(
             code.value, title, detail, category, status, prefix="UiPathMCP"

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -6,7 +6,8 @@
 import sys
 import tempfile
 import uuid
-from typing import Any, AsyncGenerator, Dict, List, Optional
+from collections.abc import AsyncGenerator
+from typing import Any
 
 from httpx import HTTPStatusError
 from mcp import ClientSession, StdioServerParameters, stdio_client
@@ -74,11 +75,11 @@ def __init__(
         self._server_id = server_id
         self._server_slug = server_slug
 
-        self._signalr_client: Optional[SignalRClient] = None
-        self._session_servers: Dict[str, SessionServer] = {}
-        self._session_output: Optional[str] = None
+        self._signalr_client: SignalRClient | None = None
+        self._session_servers: dict[str, SessionServer] = {}
+        self._session_output: str | None = None
         self._cancel_event = asyncio.Event()
-        self._keep_alive_task: Optional[asyncio.Task[None]] = None
+        self._keep_alive_task: asyncio.Task[None] | None = None
         self._uipath = UiPath()
         self._cleanup_done = False
 
@@ -396,7 +397,7 @@ async def _register(self) -> None:
         """Register the MCP server with UiPath."""
 
         initialization_successful = False
-        tools_result: Optional[ListToolsResult] = None
+        tools_result: ListToolsResult | None = None
         server_stderr_output = ""
         env_vars = self._server.env
 
@@ -433,7 +434,6 @@ async def _register(self) -> None:
 
                             # Only proceed if initialization was successful
                             tools_result = await session.list_tools()
-                            # logger.info(tools_result)
                         except Exception as err:
                             logger.error(f"Initialization error: {err}")
                             # Capture stderr output here, after the timeout
@@ -472,7 +472,7 @@ async def _register(self) -> None:
                     UiPathErrorCategory.DEPLOYMENT,
                 )
 
-            tools_list: List[Dict[str, str | None]] = []
+            tools_list: list[dict[str, str | None]] = []
             client_info = {
                 "server": {
                     "Name": self.slug,
@@ -641,12 +641,9 @@ def packaged(self) -> bool:
         Returns:
             bool: True if this is a packaged runtime (has a process), False otherwise.
         """
-        process_key = None
-        process_key = self._process_key
-
         return (
-            process_key is not None
-            and process_key != "00000000-0000-0000-0000-000000000000"
+            self._process_key is not None
+            and self._process_key != "00000000-0000-0000-0000-000000000000"
         )
 
     @property

src/uipath_mcp/_cli/_runtime/_session.py

@@ -2,7 +2,7 @@
 import io
 import logging
 import tempfile
-from typing import Dict, Optional
+from typing import Any
 
 from mcp import StdioServerParameters, stdio_client
 from mcp.shared.message import SessionMessage
@@ -33,20 +33,20 @@ def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
         self._server_config = server_config
         self._server_slug = server_slug
         self._session_id = session_id
-        self._read_stream = None
-        self._write_stream = None
-        self._mcp_session = None
-        self._run_task: Optional[asyncio.Task[None]] = None
+        self._read_stream: Any = None
+        self._write_stream: Any = None
+        self._mcp_session: Any = None
+        self._run_task: asyncio.Task[None] | None = None
         self._message_queue: asyncio.Queue[JSONRPCMessage] = asyncio.Queue()
-        self._active_requests: Dict[str, str] = {}
-        self._last_request_id: Optional[str] = None
-        self._last_message_id: Optional[str] = None
+        self._active_requests: dict[str, str] = {}
+        self._last_request_id: str | None = None
+        self._last_message_id: str | None = None
         self._uipath = UiPath()
         self._mcp_tracer = McpTracer(tracer, logger)
-        self._server_stderr_output: Optional[str] = None
+        self._server_stderr_output: str | None = None
 
     @property
-    def output(self) -> Optional[str]:
+    def output(self) -> str | None:
         """Returns the captured stderr output from the MCP server process."""
         return self._server_stderr_output
 
@@ -226,7 +226,7 @@ async def _consume_messages(self):
                 message = await self._message_queue.get()
                 try:
                     if self._write_stream:
-                        logger.info(
+                        logger.debug(
                             f"Session {self._session_id} - processing queued message: {message}..."
                         )
                         await self._write_stream.send(SessionMessage(message=message))
@@ -279,7 +279,7 @@ async def _send_message_internal(
             json=message.model_dump(),
         )
         if response.status_code == 202:
-            logger.info(f"Outgoing message sent to UiPath MCP Server: {message}")
+            logger.debug(f"Outgoing message sent to UiPath MCP Server: {message}")
         elif 500 <= response.status_code < 600:
             raise Exception(f"{response.status_code} - {response.text}")
 
@@ -292,7 +292,7 @@ async def _get_messages_internal(self, request_id: str) -> None:
             self._last_request_id = request_id
             messages = response.json()
             for message in messages:
-                logger.info(f"Received message: {message}")
+                logger.debug(f"Received message: {message}")
                 json_message = JSONRPCMessage.model_validate(message)
                 if self._is_request(json_message):
                     message_id = self._get_message_id(json_message)

src/uipath_mcp/_cli/_runtime/_tracer.py

@@ -1,26 +1,23 @@
 import json
 import logging
-from typing import Dict, Optional, TypeVar
 
 import mcp.types as types
 from opentelemetry import trace
 from opentelemetry.trace import Span, StatusCode
 
-T = TypeVar("T")
-
 
 class McpTracer:
     """Helper class to create and manage spans for MCP messages."""
 
     def __init__(
         self,
-        tracer: Optional[trace.Tracer] = None,
-        logger: Optional[logging.Logger] = None,
+        tracer: trace.Tracer | None = None,
+        logger: logging.Logger | None = None,
     ):
         self._tracer = tracer or trace.get_tracer(__name__)
         self._logger = logger or logging.getLogger(__name__)
         # Dictionary to store active request spans
-        self._active_request_spans: Dict[str, Span] = {}
+        self._active_request_spans: dict[str, Span] = {}
 
     def create_span_for_message(self, message: types.JSONRPCMessage, **context) -> Span:
         """Create and configure a span for a message.

src/uipath_mcp/_cli/_utils/_config.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
@@ -12,7 +12,7 @@ class McpServer:
     def __init__(
         self,
         name: str,
-        server_config: Dict[str, Any],
+        server_config: dict[str, Any],
     ):
         self.name = name
         self.type = server_config.get("type")
@@ -24,11 +24,11 @@ def __init__(
                 self.env[key] = os.environ[key]
 
     @property
-    def file_path(self) -> Optional[str]:
+    def file_path(self) -> str | None:
         """Get the file path from args if available."""
         return self.args[0] if self.args and len(self.args) > 0 else None
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Convert the server model back to a dictionary."""
         return {"type": self.type, "command": self.command, "args": self.args}
 
@@ -39,8 +39,8 @@ def __repr__(self) -> str:
 class McpConfig:
     def __init__(self, config_path: str = "mcp.json"):
         self.config_path = config_path
-        self._servers: Dict[str, McpServer] = {}
-        self._raw_config: Dict[str, Any] = {}
+        self._servers: dict[str, McpServer] = {}
+        self._raw_config: dict[str, Any] = {}
 
         if self.exists:
             self._load_config()
@@ -68,11 +68,11 @@ def _load_config(self) -> None:
             logger.error(f"Failed to load mcp.json: {str(e)}")
             raise
 
-    def get_servers(self) -> List[McpServer]:
+    def get_servers(self) -> list[McpServer]:
         """Get list of all server models."""
         return list(self._servers.values())
 
-    def get_server(self, name: str) -> Optional[McpServer]:
+    def get_server(self, name: str) -> McpServer | None:
         """
         Get a server model by name.
         If there's only one server available, return that one regardless of name.
@@ -85,11 +85,11 @@ def get_server(self, name: str) -> Optional[McpServer]:
         # Otherwise, fall back to looking up by name
         return self._servers.get(name)
 
-    def get_server_names(self) -> List[str]:
+    def get_server_names(self) -> list[str]:
         """Get list of all server names."""
         return list(self._servers.keys())
 
-    def load_config(self) -> Dict[str, Any]:
+    def load_config(self) -> dict[str, Any]:
         """Load and validate MCP servers configuration."""
         if not self.exists:
             raise FileNotFoundError(f"Config file not found: {self.config_path}")