Skip to content

Commit affe586

Browse files
committed
feat: add token refresher
1 parent 5717b13 commit affe586

File tree

9 files changed

+590
-8
lines changed

9 files changed

+590
-8
lines changed

.github/workflows/ci.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,3 +20,6 @@ jobs:
2020

2121
lint:
2222
uses: ./.github/workflows/lint.yml
23+
24+
test:
25+
uses: ./.github/workflows/test.yml

.github/workflows/test.yml

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
name: Test
2+
3+
on:
4+
workflow_call
5+
6+
jobs:
7+
test:
8+
name: Test
9+
runs-on: ${{ matrix.os }}
10+
timeout-minutes: 10
11+
strategy:
12+
matrix:
13+
python-version: ["3.11", "3.12", "3.13"]
14+
os: [ubuntu-latest, windows-latest]
15+
16+
permissions:
17+
contents: read
18+
19+
steps:
20+
- name: Checkout
21+
uses: actions/checkout@v4
22+
23+
- name: Setup uv
24+
uses: astral-sh/setup-uv@v5
25+
26+
- name: Setup Python
27+
uses: actions/setup-python@v5
28+
with:
29+
python-version: ${{ matrix.python-version }}
30+
31+
- name: Install dependencies
32+
run: uv sync --all-extras
33+
34+
- name: Run tests
35+
run: uv run pytest
36+
37+
continue-on-error: true

pyproject.toml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-mcp"
3-
version = "0.1.1"
3+
version = "0.1.2"
44
description = "UiPath MCP SDK"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"
@@ -42,6 +42,7 @@ dev = [
4242
"mypy>=1.14.1",
4343
"ruff>=0.9.4",
4444
"pytest>=7.4.0",
45+
"pytest-asyncio>=0.23.0",
4546
"pytest-cov>=4.1.0",
4647
"pytest-mock>=3.11.1",
4748
"pre-commit>=4.5.1",
@@ -84,6 +85,7 @@ disallow_untyped_defs = false
8485
testpaths = ["tests"]
8586
python_files = "test_*.py"
8687
addopts = "-ra -q"
88+
asyncio_mode = "auto"
8789

8890
[[tool.uv.index]]
8991
name = "testpypi"

src/uipath_mcp/_cli/_runtime/_runtime.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
from ._context import UiPathServerType
3939
from ._exception import McpErrorCode, UiPathMcpRuntimeError
4040
from ._session import BaseSessionServer, StdioSessionServer, StreamableHttpSessionServer
41+
from ._token_refresh import TokenRefresher
4142

4243
logger = logging.getLogger(__name__)
4344
tracer = trace.get_tracer(__name__)
@@ -85,6 +86,7 @@ def __init__(
8586
self._http_stderr_drain_task: asyncio.Task[None] | None = None
8687
self._http_server_stderr_lines: list[str] = []
8788
self._uipath = UiPath()
89+
self._token_refresher: TokenRefresher | None = None
8890
self._cleanup_done = False
8991

9092
# Context fields from UiPathConfig
@@ -207,6 +209,8 @@ async def _run_server(self) -> UiPathRuntimeResult:
207209
root_span.set_attribute("args", json.dumps(self._server.args))
208210
root_span.set_attribute("span_type", "MCP Server")
209211
bearer_token = self._uipath._config.secret
212+
self._token_refresher = TokenRefresher(self._uipath)
213+
210214
self._signalr_client = SignalRClient(
211215
signalr_url,
212216
headers={
@@ -236,6 +240,7 @@ async def _run_server(self) -> UiPathRuntimeResult:
236240
run_task = asyncio.create_task(self._signalr_client.run())
237241
cancel_task = asyncio.create_task(self._cancel_event.wait())
238242
self._keep_alive_task = asyncio.create_task(self._keep_alive())
243+
self._token_refresher.start()
239244

240245
try:
241246
# Wait for either the run to complete or cancellation
@@ -297,6 +302,9 @@ async def _cleanup(self) -> None:
297302

298303
await self._on_runtime_abort()
299304

305+
if self._token_refresher:
306+
await self._token_refresher.stop()
307+
300308
if self._keep_alive_task:
301309
self._keep_alive_task.cancel()
302310
try:
@@ -374,11 +382,11 @@ async def _handle_signalr_message(self, args: list[str]) -> None:
374382
session_server: BaseSessionServer
375383
if self._server.is_streamable_http:
376384
session_server = StreamableHttpSessionServer(
377-
self._server, self.slug, session_id
385+
self._server, self.slug, session_id, self._uipath
378386
)
379387
else:
380388
session_server = StdioSessionServer(
381-
self._server, self.slug, session_id
389+
self._server, self.slug, session_id, self._uipath
382390
)
383391
try:
384392
await session_server.start()

src/uipath_mcp/_cli/_runtime/_session.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,13 @@
3131
class BaseSessionServer(ABC):
3232
"""Base class with transport-agnostic message relay logic."""
3333

34-
def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
34+
def __init__(
35+
self,
36+
server_config: McpServer,
37+
server_slug: str,
38+
session_id: str,
39+
uipath: UiPath,
40+
):
3541
self._server_config = server_config
3642
self._server_slug = server_slug
3743
self._session_id = session_id
@@ -42,7 +48,7 @@ def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
4248
self._active_requests: dict[str, str] = {}
4349
self._last_request_id: str | None = None
4450
self._last_message_id: str | None = None
45-
self._uipath = UiPath()
51+
self._uipath = uipath
4652
self._mcp_tracer = McpTracer(tracer, logger)
4753

4854
@property
@@ -284,8 +290,14 @@ def _get_message_id(self, message: JSONRPCMessage) -> str:
284290
class StdioSessionServer(BaseSessionServer):
285291
"""Manages a stdio server process for a specific session."""
286292

287-
def __init__(self, server_config: McpServer, server_slug: str, session_id: str):
288-
super().__init__(server_config, server_slug, session_id)
293+
def __init__(
294+
self,
295+
server_config: McpServer,
296+
server_slug: str,
297+
session_id: str,
298+
uipath: UiPath,
299+
):
300+
super().__init__(server_config, server_slug, session_id, uipath)
289301
self._server_stderr_output: str | None = None
290302

291303
@property

0 commit comments

Comments
 (0)