feat: add token refresher

Author: radugheo

.github/workflows/ci.yml

@@ -20,3 +20,6 @@ jobs:
 
   lint:
     uses: ./.github/workflows/lint.yml
+
+  test:
+    uses: ./.github/workflows/test.yml

.github/workflows/test.yml

@@ -0,0 +1,37 @@
+name: Test
+
+on:
+  workflow_call
+
+jobs:
+  test:
+    name: Test
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 10
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+        os: [ubuntu-latest, windows-latest]
+
+    permissions:
+      contents: read
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --all-extras
+
+      - name: Run tests
+        run: uv run pytest
+
+    continue-on-error: true

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-mcp"
-version = "0.1.1"
+version = "0.1.2"
 description = "UiPath MCP SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
@@ -42,6 +42,7 @@ dev = [
     "mypy>=1.14.1",
     "ruff>=0.9.4",
     "pytest>=7.4.0",
+    "pytest-asyncio>=0.23.0",
     "pytest-cov>=4.1.0",
     "pytest-mock>=3.11.1",
     "pre-commit>=4.5.1",
@@ -84,6 +85,7 @@ disallow_untyped_defs = false
 testpaths = ["tests"]
 python_files = "test_*.py"
 addopts = "-ra -q"
+asyncio_mode = "auto"
 
 [[tool.uv.index]]
 name = "testpypi"

src/uipath_mcp/_cli/_runtime/_runtime.py

@@ -38,6 +38,7 @@
 from ._context import UiPathServerType
 from ._exception import McpErrorCode, UiPathMcpRuntimeError
 from ._session import BaseSessionServer, StdioSessionServer, StreamableHttpSessionServer
+from ._token_refresh import TokenRefresher
 
 logger = logging.getLogger(__name__)
 tracer = trace.get_tracer(__name__)
@@ -85,6 +86,7 @@ def __init__(
         self._http_stderr_drain_task: asyncio.Task[None] | None = None
         self._http_server_stderr_lines: list[str] = []
         self._uipath = UiPath()
+        self._token_refresher: TokenRefresher | None = None
         self._cleanup_done = False
 
         # Context fields from UiPathConfig
@@ -206,15 +208,19 @@ async def _run_server(self) -> UiPathRuntimeResult:
                 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")
-                bearer_token = self._uipath._config.secret
+
+                signalr_headers = {
+                    "X-UiPath-Internal-TenantId": str(self._tenant_id),
+                    "X-UiPath-Internal-AccountId": str(self._org_id),
+                    "X-UIPATH-FolderKey": self._folder_key,
+                    "Authorization": f"Bearer {self._uipath._config.secret}",
+                }
+
+                self._token_refresher = TokenRefresher(self._uipath)
+
                 self._signalr_client = SignalRClient(
                     signalr_url,
-                    headers={
-                        "X-UiPath-Internal-TenantId": str(self._tenant_id),
-                        "X-UiPath-Internal-AccountId": str(self._org_id),
-                        "X-UIPATH-FolderKey": self._folder_key,
-                        "Authorization": f"Bearer {bearer_token}",
-                    },
+                    headers=signalr_headers,
                 )
                 self._signalr_client.on("MessageReceived", self._handle_signalr_message)
                 self._signalr_client.on(
@@ -236,6 +242,7 @@ 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._token_refresher.start()
 
                 try:
                     # Wait for either the run to complete or cancellation
@@ -297,6 +304,9 @@ async def _cleanup(self) -> None:
 
         await self._on_runtime_abort()
 
+        if self._token_refresher:
+            await self._token_refresher.stop()
+
         if self._keep_alive_task:
             self._keep_alive_task.cancel()
             try:
@@ -374,11 +384,11 @@ async def _handle_signalr_message(self, args: list[str]) -> None:
                 session_server: BaseSessionServer
                 if self._server.is_streamable_http:
                     session_server = StreamableHttpSessionServer(
-                        self._server, self.slug, session_id
+                        self._server, self.slug, session_id, self._uipath
                     )
                 else:
                     session_server = StdioSessionServer(
-                        self._server, self.slug, session_id
+                        self._server, self.slug, session_id, self._uipath
                     )
                 try:
                     await session_server.start()

src/uipath_mcp/_cli/_runtime/_session.py

@@ -31,7 +31,13 @@
 class BaseSessionServer(ABC):
     """Base class with transport-agnostic message relay logic."""
 
-    def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
+    def __init__(
+        self,
+        server_config: McpServer,
+        server_slug: str,
+        session_id: str,
+        uipath: UiPath,
+    ):
         self._server_config = server_config
         self._server_slug = server_slug
         self._session_id = session_id
@@ -42,7 +48,7 @@ def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
         self._active_requests: dict[str, str] = {}
         self._last_request_id: str | None = None
         self._last_message_id: str | None = None
-        self._uipath = UiPath()
+        self._uipath = uipath
         self._mcp_tracer = McpTracer(tracer, logger)
 
     @property
@@ -284,8 +290,14 @@ def _get_message_id(self, message: JSONRPCMessage) -> str:
 class StdioSessionServer(BaseSessionServer):
     """Manages a stdio server process for a specific session."""
 
-    def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
-        super().__init__(server_config, server_slug, session_id)
+    def __init__(
+        self,
+        server_config: McpServer,
+        server_slug: str,
+        session_id: str,
+        uipath: UiPath,
+    ):
+        super().__init__(server_config, server_slug, session_id, uipath)
         self._server_stderr_output: str | None = None
 
     @property