Skip to content

Commit 136a9c7

Browse files
cristipufucristipufu
authored
Merge pull request #7 from UiPath/fix/samples
feat: subscribe to signalr events
2 parents 94833af + 003fa7a commit 136a9c7

File tree

3 files changed

+339
-81
lines changed

3 files changed

+339
-81
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-mcp"
3-
version = "0.0.2"
3+
version = "0.0.3"
44
description = "UiPath MCP SDK"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.10"

‎src/uipath_mcp/_cli/_runtime/_runtime.py‎

Lines changed: 144 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import asyncio
22
import logging
3+
import os
34
import sys
4-
from typing import Optional
5+
from typing import Any, Dict, Optional
56

67
from mcp import ClientSession, StdioServerParameters
78
from mcp.client.stdio import stdio_client
9+
from pysignalr.client import SignalRClient
810
from uipath import UiPath
911
from uipath._cli._runtime._contracts import (
1012
UiPathBaseRuntime,
@@ -15,6 +17,7 @@
1517
from .._utils._config import McpServer
1618
from ._context import UiPathMcpRuntimeContext
1719
from ._exception import UiPathMcpRuntimeError
20+
from ._session import SessionServer
1821

1922
logger = logging.getLogger(__name__)
2023

@@ -29,69 +32,150 @@ def __init__(self, context: UiPathMcpRuntimeContext):
2932
super().__init__(context)
3033
self.context: UiPathMcpRuntimeContext = context
3134
self.server: Optional[McpServer] = None
35+
self.signalr_client: Optional[SignalRClient] = None
36+
self.session_servers: Dict[str, SessionServer] = {}
3237

3338
async def execute(self) -> Optional[UiPathRuntimeResult]:
3439
"""
35-
Start the MCP server.
40+
Start the runtime and connect to SignalR.
3641
3742
Returns:
3843
Dictionary with execution results
3944
4045
Raises:
4146
UiPathMcpRuntimeError: If execution fails
4247
"""
43-
4448
await self.validate()
4549

4650
try:
47-
4851
if self.server is None:
4952
return None
5053

51-
server_params = StdioServerParameters(
52-
command=self.server.command,
53-
args=self.server.args,
54-
env=None,
54+
# Set up SignalR client
55+
signalr_url = (
56+
f"{os.environ.get('UIPATH_URL')}/mcp_/wsstunnel?slug={self.server.name}"
5557
)
5658

57-
print(f"Starting MCP server.. {self.server.command} {self.server.args}")
58-
async with stdio_client(server_params) as (read, write):
59-
async with ClientSession(
60-
read, write
61-
) as session:
59+
self.signalr_client = SignalRClient(signalr_url)
60+
self.signalr_client.on("MessageReceived", self.handle_signalr_message)
61+
self.signalr_client.on_error(self.handle_signalr_error)
62+
self.signalr_client.on_open(self.handle_signalr_open)
63+
self.signalr_client.on_close(self.handle_signalr_close)
6264

63-
print("Connected to MCP server")
64-
# Initialize the connection
65-
await session.initialize()
66-
print("MCP server initialized")
67-
# List available prompts
68-
#prompts = await session.list_prompts()
65+
# Register the server with UiPath MCP Server
66+
await self._register()
6967

70-
# Get a prompt
71-
#prompt = await session.get_prompt(
72-
# "example-prompt", arguments={"arg1": "value"}
73-
#)
68+
# Keep the runtime alive
69+
# Start SignalR client and keep it running (this is a blocking call)
70+
logger.info("Starting SignalR client...")
71+
await self.signalr_client.run()
7472

75-
# List available resources
76-
#resources = await session.list_resources()
73+
return UiPathRuntimeResult()
7774

78-
# List available tools
79-
toolsResult = await session.list_tools()
75+
except Exception as e:
76+
if isinstance(e, UiPathMcpRuntimeError):
77+
raise
8078

81-
print(toolsResult)
79+
detail = f"Error: {str(e)}"
8280

83-
# Register with UiPath MCP Server
81+
raise UiPathMcpRuntimeError(
82+
"EXECUTION_ERROR",
83+
"MCP Runtime execution failed",
84+
detail,
85+
UiPathErrorCategory.USER,
86+
) from e
87+
88+
finally:
89+
await self.cleanup()
90+
91+
async def validate(self) -> None:
92+
"""Validate runtime inputs and load MCP server configuration."""
93+
self.server = self.context.config.get_server(self.context.entrypoint)
94+
if not self.server:
95+
raise UiPathMcpRuntimeError(
96+
"SERVER_NOT_FOUND",
97+
"MCP server not found",
98+
f"Server '{self.context.entrypoint}' not found in configuration",
99+
UiPathErrorCategory.DEPLOYMENT,
100+
)
101+
102+
async def handle_signalr_message(self, args: list) -> None:
103+
"""
104+
Handle incoming SignalR messages.
105+
The SignalR client will call this with the arguments from the server.
106+
"""
107+
if len(args) < 2:
108+
logger.error(f"Received invalid SignalR message arguments: {args}")
109+
return
110+
111+
session_id = args[0]
112+
message = args[1]
113+
114+
logger.info(f"Received message for session {session_id}: {message}")
115+
116+
try:
117+
# Check if we have a session server for this session_id
118+
if session_id not in self.session_servers:
119+
# Create and start a new session server
120+
session_server = SessionServer(self.server, session_id)
121+
self.session_servers[session_id] = session_server
122+
await session_server.start(self.signalr_client)
123+
124+
# Get the session server for this session
125+
session_server = self.session_servers[session_id]
126+
127+
# Forward the message to the session's MCP server
128+
await session_server.send_message(message)
129+
130+
except Exception as e:
131+
logger.error(
132+
f"Error handling SignalR message for session {session_id}: {str(e)}"
133+
)
134+
135+
async def handle_signalr_error(self, error: Any) -> None:
136+
"""Handle SignalR errors."""
137+
logger.error(f"SignalR error: {error}")
138+
139+
async def handle_signalr_open(self) -> None:
140+
"""Handle SignalR connection open event."""
141+
logger.info("SignalR connection established")
142+
143+
async def handle_signalr_close(self) -> None:
144+
"""Handle SignalR connection close event."""
145+
logger.info("SignalR connection closed")
146+
147+
# Clean up all session servers when the connection closes
148+
await self.cleanup()
149+
150+
async def _register(self) -> None:
151+
"""Register the MCP server type with UiPath."""
152+
logger.info(f"Registering MCP server type: {self.server.name}")
153+
154+
try:
155+
# Create a temporary session to get tools
156+
server_params = StdioServerParameters(
157+
command=self.server.command,
158+
args=self.server.args,
159+
env=None,
160+
)
161+
162+
# Start a temporary stdio client to get tools
163+
async with stdio_client(server_params) as (read, write):
164+
async with ClientSession(read, write) as session:
165+
await session.initialize()
166+
tools_result = await session.list_tools()
167+
print(tools_result)
84168
client_info = {
85169
"server": {
86170
"Name": self.server.name,
87171
"Slug": self.server.name,
88172
"Version": "1.0.0",
89-
"Type": 1
173+
"Type": 1,
90174
},
91-
"tools": []
175+
"tools": [],
92176
}
93177

94-
for tool in toolsResult.tools:
178+
for tool in tools_result.tools:
95179
tool_info = {
96180
"Type": 1,
97181
"Name": tool.name,
@@ -100,59 +184,39 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
100184
}
101185
client_info["tools"].append(tool_info)
102186

103-
print(client_info)
104-
print("Registering client...")
105-
187+
# Register with UiPath MCP Server
106188
uipath = UiPath()
107-
sseUrl: str
108-
try:
109-
response = uipath.api_client.request(
110-
"POST",
111-
f"/mcp_/api/servers-with-tools/{self.server.name}",
112-
json=client_info,
113-
)
114-
#data = response.json()
115-
#sseUrl = data.get("url")
116-
print("Registered client successfully")
117-
except Exception as e:
118-
raise UiPathMcpRuntimeError(
119-
"NETWORK_ERROR",
120-
"Failed to register with UiPath MCP Server",
121-
str(e),
122-
UiPathErrorCategory.SYSTEM,
123-
) from e
124-
125-
return UiPathRuntimeResult()
189+
uipath.api_client.request(
190+
"POST",
191+
f"mcp_/api/servers-with-tools/{self.server.name}",
192+
json=client_info,
193+
)
194+
logger.info("Registered MCP Server type successfully")
126195

127196
except Exception as e:
128-
if isinstance(e, UiPathMcpRuntimeError):
129-
raise
130-
131-
detail = f"Error: {str(e)}"
132-
133197
raise UiPathMcpRuntimeError(
134-
"EXECUTION_ERROR",
135-
"MCP Server execution failed",
136-
detail,
137-
UiPathErrorCategory.USER,
198+
"NETWORK_ERROR",
199+
"Failed to register with UiPath MCP Server",
200+
str(e),
201+
UiPathErrorCategory.SYSTEM,
138202
) from e
139203

140-
finally:
141-
# Add a small delay to allow the server to shut down gracefully
142-
if sys.platform == 'win32':
143-
await asyncio.sleep(0.1)
204+
async def cleanup(self) -> None:
205+
"""Clean up all resources."""
206+
logger.info("Cleaning up all resources")
144207

145-
async def validate(self) -> None:
146-
"""Validate runtime inputs."""
147-
"""Load and validate the MCP server configuration ."""
148-
self.server = self.context.config.get_server(self.context.entrypoint)
149-
if not self.server:
150-
raise UiPathMcpRuntimeError(
151-
"SERVER_NOT_FOUND",
152-
"MCP server not found",
153-
f"Server '{self.context.entrypoint}' not found in configuration",
154-
UiPathErrorCategory.DEPLOYMENT,
155-
)
208+
# Clean up all session servers
209+
for session_id, session_server in list(self.session_servers.items()):
210+
try:
211+
await session_server.cleanup()
212+
except Exception as e:
213+
logger.error(f"Error cleaning up session {session_id}: {str(e)}")
214+
215+
self.session_servers.clear()
216+
217+
# Close SignalR connection
218+
# self.signalr_client
156219

157-
async def cleanup(self):
158-
pass
220+
# Add a small delay to allow the server to shut down gracefully
221+
if sys.platform == "win32":
222+
await asyncio.sleep(0.1)

0 commit comments

Comments
 (0)