Skip to content

Commit 9e3a3b1

Browse files
edis-uipathclaude
andcommitted
feat(mcp): suspend/resume the agent on long-running UiPath jobs via _meta
Implements Option 6 of the long-running-jobs design (mcp-longrunning-jobs-design.md): a UiPath MCP server that backs a tool with an Orchestrator job advertises `uipath.com/job` on initialize; the client then opts in per call and suspends the LangGraph agent while the job runs, resuming with the result — all over MCP `_meta`, so deployed agents gain the behavior on package upgrade with no agent.json change. Client core (C1): - `McpClient`: read the `InitializeResult._meta["uipath.com/job"]` advertisement (`_apply_job_advertisement` → `is_job_aware`/`job_version`), and thread request `_meta` through `call_tool(..., meta=)`. - `mcp_tool._invoke_job_aware`: on a job-aware session, send the START `_meta` (`{version}`) on every call; if the server returns a `{key, folderKey}` handle, delegate to the injected `McpJobExecutor` with neutral `start`/`fetch` closures (FETCH re-calls the tool with the handle `_meta`). Non-advertising/old servers keep today's plain blocking path (back-compat). LangGraph executor (C2): - `LangGraphJobExecutor` (default for `create_mcp_tools_and_clients`): starts the job inside `@durable_interrupt`, interrupts with `WaitJobRaw` (JOB/JOB_RAW trigger), and on resume re-derives the handle from the terminal `Job` and FETCHes the server-formatted result. `interrupt` is confined here; the neutral core lives in `uipath.platform.mcp_jobs` (uipath-python). Non-job results resolve via `SkipInterruptValue` to keep the durable index aligned without suspending. Drops the fragile Option-10 notification-string bridge: advertisement-gating alone gives back-compat (old server → no advertisement → plain path) and the free gain on central server upgrade. Tests: 9 new (advertisement, START/FETCH `_meta`, executor suspend/resume/faulted); full test_mcp suite green (232 incl. circular-import guard); ruff + mypy clean. Note: stacked on uipath-python's unreleased `uipath-platform>=0.1.65` (the `mcp_jobs` neutral core). Floor bumped; uv.lock to be regenerated and CI goes green once that release publishes. Developed against an editable local link (reverted). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent ccb95d3 commit 9e3a3b1

9 files changed

Lines changed: 542 additions & 17 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ requires-python = ">=3.11"
77
dependencies = [
88
"uipath>=2.10.79, <2.11.0",
99
"uipath-core>=0.5.17, <0.6.0",
10-
"uipath-platform>=0.1.61, <0.2.0",
10+
"uipath-platform>=0.1.65, <0.2.0",
1111
"uipath-runtime>=0.11.0, <0.12.0",
1212
"langgraph>=1.1.8, <2.0.0",
1313
"langchain-core>=1.2.11, <2.0.0",

src/uipath_langchain/agent/tools/mcp/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""MCP (Model Context Protocol) tools."""
22

3+
from .job_executor import LangGraphJobExecutor
34
from .mcp_client import McpClient, SessionInfoFactory
45
from .mcp_tool import (
56
create_mcp_tools,
@@ -9,6 +10,7 @@
910
from .streamable_http import SessionInfo
1011

1112
__all__ = [
13+
"LangGraphJobExecutor",
1214
"McpClient",
1315
"SessionInfo",
1416
"SessionInfoFactory",

src/uipath_langchain/agent/tools/mcp/claude.md

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,14 @@ src/uipath_langchain/agent/tools/mcp/
2424
├── __init__.py # Public exports
2525
├── mcp_client.py # SessionInfoFactory, McpClient
2626
├── mcp_tool.py # Tool factory functions
27+
├── job_executor.py # LangGraphJobExecutor (uipath.com/job suspend/resume)
2728
└── streamable_http.py # SessionInfo, StreamableHTTPTransport (copied from MCP SDK)
2829
```
2930

3031
### Public Exports (`__init__.py`)
3132

3233
```python
34+
from .job_executor import LangGraphJobExecutor
3335
from .mcp_client import McpClient, SessionInfoFactory
3436
from .mcp_tool import (
3537
create_mcp_tools_and_clients,
@@ -502,6 +504,57 @@ endpoint (`GET/PUT agenthub_/design/debugstate/{agentId}/{key}`). It lives
502504
in `uipath-agents` because it depends on execution-type logic that belongs
503505
in the agent layer, not in the langchain tools layer.
504506

507+
## `uipath.com/job` — long-running job support
508+
509+
When a UiPath MCP server backs a tool with an Orchestrator **job**, the agent
510+
should *suspend* while the job runs and *resume* with its result instead of
511+
blocking. This is negotiated entirely through MCP `_meta` (no `agent.json`
512+
change), so a deployed agent gains the behavior on package upgrade.
513+
514+
**Flow (single key `uipath.com/job`, all in `_meta`):**
515+
516+
1. **Advertise** — on `initialize`, a job-capable server returns
517+
`InitializeResult._meta["uipath.com/job"] = {"version": 1}`.
518+
`McpClient._apply_job_advertisement()` reads it and sets `is_job_aware` /
519+
`job_version`. Re-read on every fresh `initialize`, so the flag is stable
520+
across a production suspend/resume.
521+
2. **START** — for a job-aware session, `mcp_tool._invoke_job_aware` sends
522+
`params._meta["uipath.com/job"] = {"version": N}` (no `key`) on `tools/call`
523+
via `McpClient.call_tool(..., meta=...)`. A job-backed tool returns
524+
`result._meta["uipath.com/job"] = {"key", "folderKey"}` immediately.
525+
3. **Suspend** — the `McpJobExecutor` (default `LangGraphJobExecutor`) interrupts
526+
with `WaitJobRaw(Job(id=0, key, folder_key), process_folder_key=folder_key)`
527+
inside `@durable_interrupt` (same mechanism as `process_tool`). The runtime
528+
persists a `JOB` resume trigger (`JOB_RAW` name → no output extraction);
529+
Orchestrator resumes the agent when the child job is terminal.
530+
4. **FETCH** — on resume the body is skipped; `interrupt(None)` returns the
531+
terminal `Job`. The executor re-derives `{key, folderKey}` from it and calls
532+
the neutral `fetch`, which re-issues `tools/call` (no args) with
533+
`params._meta["uipath.com/job"] = {"key", "folderKey"}`. The server returns
534+
the formatted result — that `CallToolResult` is the tool's output.
535+
536+
**Where the pieces live:**
537+
538+
| Piece | Location | Notes |
539+
|-------|----------|-------|
540+
| `_meta` build/parse, `UiPathJobHandle`, `JobStart`, `McpJobExecutor`, `BlockingJobExecutor` | `uipath.platform.mcp_jobs` (uipath-python) | Framework- and MCP-SDK-neutral (plain dicts) |
541+
| `is_job_aware`, `job_version`, `job_executor`, `call_tool(meta=)`, `_apply_job_advertisement` | `mcp_client.py` | Reads the advertisement; threads `_meta` |
542+
| `_invoke_job_aware`, `_normalize_tool_result`, `create_mcp_tools_and_clients(job_executor=)` | `mcp_tool.py` | Builds the neutral `start`/`fetch`; default executor = `LangGraphJobExecutor` |
543+
| `LangGraphJobExecutor` (the only place `interrupt` is called) | `job_executor.py` | Suspend → resume → fetch |
544+
545+
**Key invariants:**
546+
547+
- `interrupt` is confined to `LangGraphJobExecutor`. The neutral core never
548+
imports langgraph; the executor never imports `mcp`.
549+
- The job handle survives the suspend **only** via the `WaitJobRaw` payload
550+
(re-derived from the resumed `Job`), never via closures — the
551+
`@durable_interrupt` body does not re-run on resume.
552+
- Non-job tools on a job-aware server: the server returns a normal result (no
553+
handle); `LangGraphJobExecutor` resolves it via `SkipInterruptValue`
554+
(`_NonJobStartValue`) so the durable index stays aligned with no real suspend.
555+
- Old/non-UiPath servers never advertise → `is_job_aware` stays `False` → today's
556+
plain blocking path (back-compat, zero behavior change).
557+
505558
## Tests
506559

507560
Tests are in `tests/agent/tools/test_mcp/`.
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
"""LangGraph executor for MCP-backed UiPath jobs.
2+
3+
When an MCP ``tools/call`` starts a UiPath job (the server returns a
4+
``uipath.com/job`` handle), :class:`LangGraphJobExecutor` suspends the LangGraph
5+
agent on that job and resumes when it finishes — the same durable suspend/resume
6+
mechanism :mod:`uipath_langchain.agent.tools.process_tool` uses.
7+
8+
It mirrors ``process_tool`` exactly:
9+
10+
* the START ``tools/call`` runs **inside** ``@durable_interrupt`` so it executes
11+
exactly once (a resume re-runs the node but skips the body);
12+
* it interrupts with ``WaitJobRaw`` so the runtime persists a ``JOB`` resume
13+
trigger (with the ``JOB_RAW`` name → resume without output extraction) and
14+
Orchestrator resumes the agent when the child job reaches a terminal state;
15+
* on resume the body is skipped and ``interrupt(None)`` returns the terminal
16+
``Job``; we re-derive the ``{key, folderKey}`` handle from it and FETCH the
17+
result with a follow-up ``tools/call`` (the server formats the output).
18+
19+
The neutral wire work (building the START/FETCH ``_meta``, parsing the handle)
20+
lives in :mod:`uipath.platform.mcp_jobs`; this class only owns the *suspend*
21+
policy, so ``interrupt`` is confined here.
22+
"""
23+
24+
from __future__ import annotations
25+
26+
from typing import Any
27+
28+
from uipath.platform.common import WaitJobRaw
29+
from uipath.platform.mcp_jobs import (
30+
FetchFn,
31+
JobStart,
32+
StartFn,
33+
UiPathJobHandle,
34+
)
35+
from uipath.platform.orchestrator import Job, JobState
36+
37+
from uipath_langchain._utils.durable_interrupt import (
38+
SkipInterruptValue,
39+
durable_interrupt,
40+
)
41+
42+
43+
class _NonJobStartValue(SkipInterruptValue):
44+
"""Carries a non-job :class:`JobStart` back through ``@durable_interrupt``.
45+
46+
A job-aware client sends the START ``_meta`` on every call, but the server
47+
only returns a handle for job-backed tools. For a normal (non-job) result we
48+
must NOT suspend — yet the ``@durable_interrupt`` body still has to run on
49+
every pass to keep the durable index aligned. Returning this value injects the
50+
result into the scratchpad and resumes immediately, without a real suspend.
51+
"""
52+
53+
def __init__(self, outcome: JobStart) -> None:
54+
self._outcome = outcome
55+
56+
@property
57+
def resume_value(self) -> Any:
58+
"""The :class:`JobStart` to return to the executor without suspending."""
59+
return self._outcome
60+
61+
62+
class LangGraphJobExecutor:
63+
"""``McpJobExecutor`` that suspends the LangGraph agent on the started job."""
64+
65+
async def run(self, *, start: StartFn, fetch: FetchFn, tool_name: str) -> Any:
66+
"""Start the job, suspend until it finishes, then FETCH its result.
67+
68+
Args:
69+
start: Issues the START ``tools/call`` once; returns a :class:`JobStart`.
70+
fetch: Re-calls the tool with the FETCH ``_meta`` for a handle.
71+
tool_name: The MCP tool name (for diagnostics).
72+
73+
Returns:
74+
The FETCH result for a job-backed call, or the normal tool result when
75+
the call did not start a job.
76+
"""
77+
78+
@durable_interrupt
79+
async def _suspend_on_job() -> Any:
80+
outcome = await start()
81+
if outcome.handle is None:
82+
# Non-job tool / no opt-in: do not suspend; carry the result.
83+
return _NonJobStartValue(outcome)
84+
# The whole handle round-trips the suspend via this WaitJobRaw payload.
85+
return WaitJobRaw(
86+
job=Job(
87+
id=0,
88+
key=outcome.handle.job_key,
89+
folder_key=outcome.handle.folder_key,
90+
),
91+
process_folder_key=outcome.handle.folder_key,
92+
)
93+
94+
resumed = await _suspend_on_job()
95+
96+
if isinstance(resumed, JobStart):
97+
# Non-job path: the START result is the tool's output.
98+
return resumed.result
99+
100+
# Resume path: `resumed` is the runtime-materialized terminal raw Job.
101+
# WaitJobRaw skips state validation, so re-check for failure (as process_tool does).
102+
job = resumed
103+
if (job.state or "").lower() == JobState.FAULTED:
104+
return str(
105+
getattr(job, "info", None) or f"Job for tool '{tool_name}' faulted"
106+
)
107+
if not job.key or not job.folder_key:
108+
return str(
109+
getattr(job, "info", None)
110+
or f"Job for tool '{tool_name}' returned no key to fetch its result"
111+
)
112+
113+
# Re-derive the handle from the resumed Job and let the server format the result.
114+
handle = UiPathJobHandle(job_key=job.key, folder_key=job.folder_key)
115+
return await fetch(handle)

src/uipath_langchain/agent/tools/mcp/mcp_client.py

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@
1515
from mcp import ClientSession
1616
from mcp.shared.exceptions import McpError
1717
from mcp.shared.message import SessionMessage
18-
from mcp.types import CallToolResult, ListToolsResult
18+
from mcp.types import CallToolResult, InitializeResult, ListToolsResult
1919
from uipath._utils._ssl_context import get_httpx_client_kwargs
20+
from uipath.platform.mcp_jobs import read_job_version
2021
from uipath.runtime.base import UiPathDisposableProtocol
2122

2223
from uipath_langchain._utils import get_execution_folder_path
@@ -25,6 +26,7 @@
2526

2627
if TYPE_CHECKING:
2728
from uipath.agent.models.agent import AgentMcpResourceConfig
29+
from uipath.platform.mcp_jobs import McpJobExecutor
2830
from uipath.platform.orchestrator.mcp import McpServer
2931

3032
logger = logging.getLogger(__name__)
@@ -78,6 +80,7 @@ def __init__(
7880
max_retries: int = 1,
7981
session_info_factory: SessionInfoFactory | None = None,
8082
terminate_on_close: bool = True,
83+
job_executor: "McpJobExecutor | None" = None,
8184
) -> None:
8285
"""Initialize the MCP tool session.
8386
@@ -90,12 +93,21 @@ def __init__(
9093
max_retries: Maximum number of retries on session disconnect errors.
9194
session_info_factory: Factory for creating SessionInfo instances.
9295
Defaults to ``SessionInfoFactory`` which returns a plain SessionInfo.
96+
terminate_on_close: Whether to terminate the MCP session on close.
97+
job_executor: Executor that awaits a UiPath job the server starts behind
98+
a ``tools/call`` (see ``uipath.com/job``). Only used when the server
99+
advertised support on ``initialize``. ``None`` disables job handling.
93100
"""
94101
self._config = config
95102
self._timeout = timeout or httpx.Timeout(600)
96103
self._max_retries = max_retries
97104
self._session_info_factory = session_info_factory or SessionInfoFactory()
98105
self._terminate_on_close = terminate_on_close
106+
self._job_executor = job_executor
107+
108+
# uipath.com/job negotiation state (set from the initialize advertisement).
109+
self._job_aware: bool = False
110+
self._job_version: int | None = None
99111

100112
# URL and headers are resolved lazily from SDK
101113
self._url: str | None = None
@@ -229,7 +241,8 @@ async def _initialize_session(self) -> None:
229241
)
230242

231243
if existing_session_id is None:
232-
await self._session.initialize()
244+
init_result = await self._session.initialize()
245+
self._apply_job_advertisement(init_result)
233246

234247
# The transport calls set_session_id during initialize,
235248
# so we just read the current value here.
@@ -240,6 +253,35 @@ async def _initialize_session(self) -> None:
240253
)
241254
logger.info(f"MCP session initialized with session ID: {new_session_id}")
242255

256+
def _apply_job_advertisement(self, init_result: InitializeResult) -> None:
257+
"""Read the ``uipath.com/job`` advertisement from the initialize result.
258+
259+
When the server advertises support (it has job-backed tools), mark the
260+
session job-aware so the client opts in (sends the START ``_meta``) on
261+
every ``tools/call``. Re-evaluated on every fresh ``initialize`` so the
262+
flag is stable across a suspend/resume in production.
263+
"""
264+
version = read_job_version(init_result.meta)
265+
if version is not None:
266+
self._job_aware = True
267+
self._job_version = version
268+
logger.info(f"MCP server advertised uipath.com/job v{version}")
269+
270+
@property
271+
def is_job_aware(self) -> bool:
272+
"""Whether the server advertised ``uipath.com/job`` support on initialize."""
273+
return self._job_aware
274+
275+
@property
276+
def job_version(self) -> int | None:
277+
"""The advertised ``uipath.com/job`` contract version, if any."""
278+
return self._job_version
279+
280+
@property
281+
def job_executor(self) -> "McpJobExecutor | None":
282+
"""The executor used to await job-backed tool calls, if configured."""
283+
return self._job_executor
284+
243285
async def _ensure_session(self) -> ClientSession:
244286
"""Ensure client and session are initialized, return the session.
245287
@@ -351,18 +393,22 @@ async def call_tool(
351393
self,
352394
name: str,
353395
arguments: dict[str, Any] | None = None,
396+
*,
397+
meta: dict[str, Any] | None = None,
354398
) -> CallToolResult:
355399
"""Call an MCP tool by name.
356400
357401
Args:
358402
name: The name of the tool to call.
359403
arguments: Optional arguments to pass to the tool.
404+
meta: Optional request ``_meta`` (e.g. the ``uipath.com/job`` START or
405+
FETCH marker). Mapped onto ``CallToolRequest.params._meta``.
360406
361407
Returns:
362408
The tool call result.
363409
"""
364410
return await self._execute_with_retry(
365-
lambda session: session.call_tool(name, arguments=arguments),
411+
lambda session: session.call_tool(name, arguments=arguments, meta=meta),
366412
f"call_tool({name})",
367413
)
368414

0 commit comments

Comments
 (0)