chore: fix types and linting errors

Author: radu-mocanu

pyproject.toml

@@ -64,6 +64,9 @@ line-ending = "auto"
 plugins = [
   "pydantic.mypy"
 ]
+exclude = [
+    "samples/.*"
+]
 
 follow_imports = "silent"
 warn_redundant_casts = true

samples/github-slack-agent/main.py

@@ -108,7 +108,9 @@ def system_prompt(state: AgentState) -> AgentState:
 _This review was generated automatically._
 """
 
-                        return [{"role": "system", "content": system_message}] + state["messages"]
+                        return [{"role": "system", "content": system_message}] + state[
+                            "messages"
+                        ]
 
                     agent = create_react_agent(
                         model,

samples/mcp-functions-agent/builder.py

@@ -180,9 +180,7 @@ async def validator_agent(state: GraphState) -> GraphState:
                 **Example use case:** If the function reads files from disk, your setup function should create a temporary folder and write some files into it. Then, test the function against that folder path.
                 """
 
-                seeder = create_react_agent(
-                    model, tools=tools, prompt=test_case_prompt
-                )
+                seeder = create_react_agent(model, tools=tools, prompt=test_case_prompt)
 
                 test_result = await seeder.ainvoke(state)
 

samples/mcp-functions-server/server.py

@@ -5,6 +5,7 @@
 # Initialize the MCP server
 mcp = FastMCP("Code Functions MCP Server")
 
+
 # Functions registry to track dynamically added code functions
 class FunctionRegistry:
     def __init__(self):

src/uipath_mcp/_cli/_runtime/_context.py

@@ -53,10 +53,10 @@ class UiPathServerType(Enum):
         SelfHosted (3): Tunnel to externally hosted server
     """
 
-    UiPath = 0  # type: int # Processes, Agents, Activities
-    Command = 1  # type: int # npx, uvx
-    Coded = 2  # type: int # PackageType.MCPServer
-    SelfHosted = 3  # type: int # tunnel to externally hosted server
+    UiPath = 0  # Processes, Agents, Activities
+    Command = 1  # npx, uvx
+    Coded = 2  # PackageType.MCPServer
+    SelfHosted = 3  # tunnel to externally hosted server
 
     @classmethod
     def from_string(cls, name: str) -> "UiPathServerType":

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -1,24 +1,27 @@
 import asyncio
+import io
 import json
 import logging
 import os
 import sys
 import tempfile
 import uuid
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, cast
 
 from httpx import HTTPStatusError
 from mcp import ClientSession, StdioServerParameters, stdio_client
-from mcp.types import JSONRPCResponse
+from mcp.types import JSONRPCResponse, ListToolsResult
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from pysignalr.client import CompletionMessage, SignalRClient
+from pysignalr.client import SignalRClient
+from pysignalr.messages import CompletionMessage
 from uipath import UiPath
 from uipath._cli._runtime._contracts import (
     UiPathBaseRuntime,
     UiPathErrorCategory,
     UiPathRuntimeResult,
+    UiPathTraceContext,
 )
 from uipath.tracing import LlmOpsHttpExporter
 
@@ -47,8 +50,75 @@ def __init__(self, context: UiPathMcpRuntimeContext):
         self._session_servers: Dict[str, SessionServer] = {}
         self._session_output: Optional[str] = None
         self._cancel_event = asyncio.Event()
-        self._keep_alive_task: Optional[asyncio.Task] = None
+        self._keep_alive_task: Optional[asyncio.Task[None]] = None
         self._uipath = UiPath()
+        self.trace_provider: Optional[TracerProvider] = None
+
+    async def validate(self) -> None:
+        """Validate runtime inputs and load MCP server configuration."""
+        if self.context.config is None:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing configuration",
+                "Configuration is required.",
+                UiPathErrorCategory.SYSTEM,
+            )
+
+        if self.context.entrypoint is None:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing entrypoint",
+                "Entrypoint is required.",
+                UiPathErrorCategory.SYSTEM,
+            )
+
+        self._server = self.context.config.get_server(self.context.entrypoint)
+        if not self._server:
+            raise UiPathMcpRuntimeError(
+                "SERVER_NOT_FOUND",
+                "MCP server not found",
+                f"Server '{self.context.entrypoint}' not found in configuration",
+                UiPathErrorCategory.DEPLOYMENT,
+            )
+
+    def _validate_auth(self) -> None:
+        """Validate authentication-related configuration.
+
+        Raises:
+            UiPathMcpRuntimeError: If any required authentication values are missing.
+        """
+        uipath_url = os.environ.get("UIPATH_URL")
+        if not uipath_url:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing UIPATH_URL environment variable",
+                "Please run 'uipath auth'.",
+                UiPathErrorCategory.USER,
+            )
+
+        if not self.context.trace_context:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing trace context",
+                "Trace context is required for SignalR connection.",
+                UiPathErrorCategory.SYSTEM,
+            )
+
+        if not self.context.trace_context.tenant_id:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing tenant ID",
+                "Please run 'uipath auth'.",
+                UiPathErrorCategory.SYSTEM,
+            )
+
+        if not self.context.trace_context.org_id:
+            raise UiPathMcpRuntimeError(
+                "CONFIGURATION_ERROR",
+                "Missing organization ID",
+                "Please run 'uipath auth'.",
+                UiPathErrorCategory.SYSTEM,
+            )
 
     async def execute(self) -> Optional[UiPathRuntimeResult]:
         """
@@ -71,10 +141,14 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 trace.set_tracer_provider(self.trace_provider)
                 self.trace_provider.add_span_processor(
                     BatchSpanProcessor(LlmOpsHttpExporter())
-                )  # type: ignore
+                )
+
+            # Validate authentication configuration
+            self._validate_auth()
 
             # Set up SignalR client
-            signalr_url = f"{os.environ.get('UIPATH_URL')}/agenthub_/wsstunnel?slug={self._server.name}&runtimeId={self._runtime_id}"
+            uipath_url = os.environ.get("UIPATH_URL")
+            signalr_url = f"{uipath_url}/agenthub_/wsstunnel?slug={self._server.name}&runtimeId={self._runtime_id}"
 
             if not self.context.folder_key:
                 folder_path = os.environ.get("UIPATH_FOLDER_PATH")
@@ -98,17 +172,22 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
 
             with tracer.start_as_current_span(self.slug) as root_span:
                 root_span.set_attribute("runtime_id", self._runtime_id)
-                root_span.set_attribute("command", self._server.command)
-                root_span.set_attribute("args", self._server.args)
+                root_span.set_attribute("command", str(self._server.command))
+                root_span.set_attribute("args", json.dumps(self._server.args))
                 root_span.set_attribute("span_type", "MCP Server")
+                trace_context = cast(UiPathTraceContext, self.context.trace_context)
+                tenant_id = str(trace_context.tenant_id)
+                org_id = str(trace_context.org_id)
+
                 self._signalr_client = SignalRClient(
                     signalr_url,
                     headers={
-                        "X-UiPath-Internal-TenantId": self.context.trace_context.tenant_id,
-                        "X-UiPath-Internal-AccountId": self.context.trace_context.org_id,
+                        "X-UiPath-Internal-TenantId": tenant_id,
+                        "X-UiPath-Internal-AccountId": org_id,
                         "X-UIPATH-FolderKey": self.context.folder_key,
                     },
                 )
+
                 self._signalr_client.on("MessageReceived", self._handle_signalr_message)
                 self._signalr_client.on(
                     "SessionClosed", self._handle_signalr_session_closed
@@ -166,20 +245,9 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
             ) from e
         finally:
             await self.cleanup()
-            if hasattr(self, "trace_provider") and self.trace_provider:
+            if self.trace_provider:
                 self.trace_provider.shutdown()
 
-    async def validate(self) -> None:
-        """Validate runtime inputs and load MCP server configuration."""
-        self._server = self.context.config.get_server(self.context.entrypoint)
-        if not self._server:
-            raise UiPathMcpRuntimeError(
-                "SERVER_NOT_FOUND",
-                "MCP server not found",
-                f"Server '{self.context.entrypoint}' not found in configuration",
-                UiPathErrorCategory.DEPLOYMENT,
-            )
-
     async def cleanup(self) -> None:
         """Clean up all resources."""
 
@@ -210,7 +278,7 @@ async def cleanup(self) -> None:
         if sys.platform == "win32":
             await asyncio.sleep(0.5)
 
-    async def _handle_signalr_session_closed(self, args: list) -> None:
+    async def _handle_signalr_session_closed(self, args: list[Any]) -> None:
         """
         Handle session closed by server.
         """
@@ -240,7 +308,7 @@ async def _handle_signalr_session_closed(self, args: list) -> None:
         except Exception as e:
             logger.error(f"Error terminating session {session_id}: {str(e)}")
 
-    async def _handle_signalr_message(self, args: list) -> None:
+    async def _handle_signalr_message(self, args: list[Any]) -> None:
         """
         Handle incoming SignalR messages.
         """
@@ -254,10 +322,12 @@ async def _handle_signalr_message(self, args: list) -> None:
         logger.info(f"Received websocket notification... {session_id}")
 
         try:
+            server = cast(McpServer, self._server)
+
             # Check if we have a session server for this session_id
             if session_id not in self._session_servers:
                 # Create and start a new session server
-                session_server = SessionServer(self._server, self.slug, session_id)
+                session_server = SessionServer(server, self.slug, session_id)
                 try:
                     await session_server.start()
                 except Exception as e:
@@ -293,11 +363,12 @@ async def _handle_signalr_close(self) -> None:
 
     async def _register(self) -> None:
         """Register the MCP server with UiPath."""
+        server = cast(McpServer, self._server)
 
         initialization_successful = False
-        tools_result = None
+        tools_result: Optional[ListToolsResult] = None
         server_stderr_output = ""
-        env_vars = self._server.env
+        env_vars = dict(server.env)
 
         # if server is Coded, include environment variables
         if self.server_type is UiPathServerType.Coded:
@@ -309,14 +380,15 @@ async def _register(self) -> None:
         try:
             # Create a temporary session to get tools
             server_params = StdioServerParameters(
-                command=self._server.command,
-                args=self._server.args,
+                command=server.command,
+                args=server.args,
                 env=env_vars,
             )
 
             # Start a temporary stdio client to get tools
             # Use a temporary file to capture stderr
-            with tempfile.TemporaryFile(mode="w+b") as stderr_temp:
+            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,
@@ -336,11 +408,7 @@ async def _register(self) -> None:
                             logger.error("Initialization timed out")
                             # Capture stderr output here, after the timeout
                             stderr_temp.seek(0)
-                            server_stderr_output = stderr_temp.read().decode(
-                                "utf-8", errors="replace"
-                            )
-                            # We'll handle this after exiting the context managers
-                        # We don't continue with registration here - we'll do it after the context managers
+                            server_stderr_output = stderr_temp.read()
 
         except* Exception as eg:
             for e in eg.exceptions:
@@ -366,15 +434,24 @@ async def _register(self) -> None:
         # Now continue with registration
         logger.info("Registering server runtime ...")
         try:
+            if not tools_result:
+                raise UiPathMcpRuntimeError(
+                    "INITIALIZATION_ERROR",
+                    "Server initialization failed",
+                    "Failed to get tools list from server",
+                    UiPathErrorCategory.DEPLOYMENT,
+                )
+
+            tools_list: List[Dict[str, str | None]] = []
             client_info = {
                 "server": {
-                    "Id": self.context.server_id,
                     "Name": self.slug,
                     "Slug": self.slug,
+                    "Id": self.context.server_id,
                     "Version": "1.0.0",
                     "Type": self.server_type.value,
                 },
-                "tools": [],
+                "tools": tools_list,
             }
 
             for tool in tools_result.tools:
@@ -386,7 +463,7 @@ async def _register(self) -> None:
                     if tool.inputSchema
                     else "{}",
                 }
-                client_info["tools"].append(tool_info)
+                tools_list.append(tool_info)
 
             # Register with UiPath MCP Server
             await self._uipath.api_client.request_async(
@@ -418,14 +495,14 @@ async def _on_session_start_error(self, session_id: str) -> None:
         try:
             response = await self._uipath.api_client.request_async(
                 "POST",
-                f"agenthub_/mcp/{self.slug}/out/message?sessionId={session_id}",
+                f"agenthub_/mcp/{self.slug}/out/message?sessionId={session_id}",  # type: ignore
                 json=JSONRPCResponse(
                     jsonrpc="2.0",
                     id=0,
                     result={
                         "protocolVersion": "initialize-failure",
                         "capabilities": {},
-                        "serverInfo": {"name": self.slug, "version": "1.0"},
+                        "serverInfo": {"name": self.slug, "version": "1.0"},  # type: ignore
                     },
                 ).model_dump(),
             )
@@ -475,7 +552,7 @@ async def on_keep_alive_response(
                         await self._signalr_client.send(
                             method="OnKeepAlive",
                             arguments=[],
-                            on_invocation=on_keep_alive_response,
+                            on_invocation=on_keep_alive_response,  # type: ignore
                         )
                 except Exception as e:
                     if not self._cancel_event.is_set():
@@ -497,7 +574,7 @@ async def _on_runtime_abort(self) -> None:
         try:
             response = await self._uipath.api_client.request_async(
                 "POST",
-                f"agenthub_/mcp/{self.slug}/runtime/abort?runtimeId={self._runtime_id}",
+                f"agenthub_/mcp/{self.slug}/runtime/abort?runtimeId={self._runtime_id}",  # type: ignore
             )
             if response.status_code == 202:
                 logger.info(
@@ -530,7 +607,9 @@ def packaged(self) -> bool:
         Returns:
             bool: True if this is a packaged runtime (has a process), False otherwise.
         """
-        process_key = self.context.trace_context.process_key
+        process_key = None
+        if self.context.trace_context is not None:
+            process_key = self.context.trace_context.process_key
 
         return (
             process_key is not None