diff --git a/pyproject.toml b/pyproject.toml index 04179a0..1031cd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-mcp" -version = "0.0.96" +version = "0.0.97" description = "UiPath MCP SDK" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/src/uipath_mcp/_cli/_runtime/_context.py b/src/uipath_mcp/_cli/_runtime/_context.py index 686b036..7864334 100644 --- a/src/uipath_mcp/_cli/_runtime/_context.py +++ b/src/uipath_mcp/_cli/_runtime/_context.py @@ -10,6 +10,7 @@ class UiPathMcpRuntimeContext(UiPathRuntimeContext): """Context information passed throughout the runtime execution.""" config: Optional[McpConfig] = None + folder_key: Optional[str] = None class UiPathServerType(Enum): diff --git a/src/uipath_mcp/_cli/_runtime/_runtime.py b/src/uipath_mcp/_cli/_runtime/_runtime.py index 9943b58..9f1c0ad 100644 --- a/src/uipath_mcp/_cli/_runtime/_runtime.py +++ b/src/uipath_mcp/_cli/_runtime/_runtime.py @@ -14,6 +14,7 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor from pysignalr.client import CompletionMessage, SignalRClient +from uipath import UiPath from uipath._cli._runtime._contracts import ( UiPathBaseRuntime, UiPathErrorCategory, @@ -25,7 +26,6 @@ from ._context import UiPathMcpRuntimeContext, UiPathServerType from ._exception import UiPathMcpRuntimeError from ._session import SessionServer -from uipath import UiPath logger = logging.getLogger(__name__) tracer = trace.get_tracer(__name__) @@ -76,6 +76,26 @@ async def execute(self) -> Optional[UiPathRuntimeResult]: # Set up SignalR client signalr_url = f"{os.environ.get('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") + if not folder_path: + raise UiPathMcpRuntimeError( + "REGISTRATION_ERROR", + "No UIPATH_FOLDER_PATH or UIPATH_FOLDER_KEY environment variable set.", + "Please set the UIPATH_FOLDER_PATH or UIPATH_FOLDER_KEY environment variable.", + UiPathErrorCategory.USER, + ) + self.context.folder_key = self._uipath.folders.retrieve_key(folder_path=folder_path) + if not self.context.folder_key: + raise UiPathMcpRuntimeError( + "REGISTRATION_ERROR", + "Folder NOT FOUND. Invalid UIPATH_FOLDER_PATH environment variable.", + "Please set the UIPATH_FOLDER_PATH or UIPATH_FOLDER_KEY environment variable.", + UiPathErrorCategory.USER, + ) + + logger.info(f"Folder key: {self.context.folder_key}") + with tracer.start_as_current_span(self._server.name) as root_span: root_span.set_attribute("runtime_id", self._runtime_id) root_span.set_attribute("command", self._server.command) @@ -86,6 +106,7 @@ async def execute(self) -> Optional[UiPathRuntimeResult]: headers={ "X-UiPath-Internal-TenantId": self.context.trace_context.tenant_id, "X-UiPath-Internal-AccountId": self.context.trace_context.org_id, + "X-UIPATH-FolderKey": self.context.folder_key, }, ) self._signalr_client.on("MessageReceived", self._handle_signalr_message) @@ -273,21 +294,6 @@ async def _handle_signalr_close(self) -> None: async def _register(self) -> None: """Register the MCP server with UiPath.""" - folder_key = os.environ.get("UIPATH_FOLDER_KEY") - folder_path = os.environ.get("UIPATH_FOLDER_PATH") - if not folder_key and not folder_path: - raise UiPathMcpRuntimeError( - "REGISTRATION_ERROR", - "No UIPATH_FOLDER_PATH or UIPATH_FOLDER_KEY environment variable set.", - "Please set the UIPATH_FOLDER_PATH or UIPATH_FOLDER_KEY environment variable.", - UiPathErrorCategory.USER, - ) - if not folder_key: - uipath = UiPath() - folder_key = uipath.folders.retrieve_key(folder_path=folder_path) - - logger.info(f"Folder key: {folder_key}") - initialization_successful = False tools_result = None server_stderr_output = "" @@ -386,7 +392,7 @@ async def _register(self) -> None: "POST", f"agenthub_/mcp/{self._server.name}/runtime/start?runtimeId={self._runtime_id}", json=client_info, - headers={"X-UIPATH-FolderKey": folder_key}, + headers={"X-UIPATH-FolderKey": self.context.folder_key}, ) logger.info("Registered MCP Server type successfully") except Exception as e: diff --git a/src/uipath_mcp/_cli/cli_run.py b/src/uipath_mcp/_cli/cli_run.py index c2d4a4c..a711aaf 100644 --- a/src/uipath_mcp/_cli/cli_run.py +++ b/src/uipath_mcp/_cli/cli_run.py @@ -39,6 +39,7 @@ async def execute(): context.resume = resume context.logs_min_level = env.get("LOG_LEVEL", "INFO") context.job_id = env.get("UIPATH_JOB_KEY") + context.folder_key = env.get("UIPATH_FOLDER_KEY") context.trace_id = env.get("UIPATH_TRACE_ID") context.tracing_enabled = env.get("UIPATH_TRACING_ENABLED", True) context.trace_context = UiPathTraceContext(