diff --git a/CHANGELOG.md b/CHANGELOG.md index 224de25..8b0234a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,11 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); ver ### Added +- `s1a-mcp`: `run_agent` can be called as an MCP task (opt-in by the client, 2025-11-25 experimental tasks in the + 1.x SDK). The call returns a task id at once; `tasks/get` reports the latest episode or step, `tasks/result` the + usual result, and `tasks/cancel` stops the run through the existing cancellation path. Plain calls are unchanged; + `decide` and `list_agents` refuse task calls. The SDK marks this API deprecated (tasks become an extension in the + 2026-07-28 spec), and no host has been verified against it yet. See #102. - `--model onejev`: [OneJev](https://github.com/OmniJev/OneJev) (Apache-2.0) as a **served** multimodal System One client against `qev`/`flide` (`ONEJEV_URL`, optional `ONEJEV_MODEL` / `ONEJEV_API_KEY` / `ONEJEV_TIMEOUT_S`). Posts choice/noul with TypeSafe-compatible `media` and `` state markers; refuses an observation without a diff --git a/docs/architecture.md b/docs/architecture.md index 09e67a9..ae64bba 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -121,6 +121,11 @@ them to: | `tests/test_desktop_driver.py` | `mcp.client.stdio.StdioServerParameters` | the 2.x stdio transport parameters | | `scripts/smoke.sh` | `mcp.shared.memory.create_connected_server_and_client_session` | `Client` in memory | +`s1a/mcp_server.py` also uses the 1.x experimental task API (`Server.experimental.enable_tasks`, +`InMemoryTaskStore`, `request_context.experimental.run_task`) for task-augmented `run_agent` calls (#102). That API +is deprecated in 1.x and removed in 2.x, where tasks are expected as a separate MCP extension; the port moves the +task path to that extension or drops it, and the `TestRunAgentAsTask` cases say what it must keep. + `mcp.types` (`CallToolResult`, `ImageContent`, `TextContent` in `s1a/desktop/driver.py` and the desktop tests) is checked against 2.x during the port. `tests/test_mcp_server.py` holds what a port must preserve: the three-tool listing, `decide` against a scripted model, `run_agent` returning its summary, and a clean stdio stream; a test there diff --git a/docs/skills.md b/docs/skills.md index bebb4a9..07bc978 100644 --- a/docs/skills.md +++ b/docs/skills.md @@ -90,6 +90,25 @@ closes the browser and the decision model it started, and records the cancellati `result.json` has `"cancelled": true` and counts the unfinished episodes in `n_cancelled_trials`; a browser agent's `answer.json` has `"cancelled": true`. +#### Long runs as MCP tasks + +`run_agent` (and only it) also accepts a task-augmented call: `tools/list` marks it `execution.taskSupport: +"optional"` and the server declares the `tasks` capability. When a host sends `tools/call` with a `task` field, the +call returns a task id at once and the run continues on the server. `tasks/get` reports `working` with the latest +`episode N/M` or `step N/M` as the `statusMessage` (also sent as `notifications/tasks/status`, at most once a second), +`tasks/result` returns the same object a plain call returns, and `tasks/cancel` stops the run through the cancellation +path above: the slot is freed, the browser is closed and the cancellation is recorded. The busy refusal and bad flags +complete the task with an `isError` result, as they are the result of a plain call. A task-augmented call of +`decide` or `list_agents` is a protocol error (`-32601`). Without a `task` field nothing changes. + +This is the experimental task API of the 2025-11-25 specification (SEP-1686) as the 1.x Python SDK implements it. +The SDK marks that API deprecated: tasks left the core specification in the 2026-07-28 revision and are expected +back as a separate extension, which 1.x does not provide. A host that speaks only the extension form will not see +task support here and keeps the plain call. Tasks are kept in memory for the server process, so they outlive the +request but not a server restart; a task's `ttl` is raised to at least one hour, because the store counts it from +creation and a shorter one could drop a run part-way. No host has been verified against this path yet, so keep the +raised tool timeout below either way. + #### Host timeouts `flights` budgets 180 s and `allrecipes` 300 s, and a tool agent's run grows with `--episodes`. Run `uv sync` in the diff --git a/s1a/mcp_server.py b/s1a/mcp_server.py index 53013ff..f7bf2c4 100644 --- a/s1a/mcp_server.py +++ b/s1a/mcp_server.py @@ -6,16 +6,40 @@ import argparse import asyncio import contextlib +import contextvars +import json import logging import sys import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Sequence from typing import Any, Literal +import anyio + import s1a.console # first: routes the harness logs to files before openjiuwen loads and logs to the console +from mcp.server.experimental.task_context import ServerTaskContext from mcp.server.fastmcp import Context, FastMCP from mcp.server.fastmcp.exceptions import ToolError -from mcp.types import ToolAnnotations +from mcp.shared.exceptions import McpError +from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore +from mcp.types import ( + METHOD_NOT_FOUND, + TASK_OPTIONAL, + TASK_STATUS_CANCELLED, + CallToolRequest, + CallToolResult, + CancelTaskRequest, + ContentBlock, + ErrorData, + Result, + ServerResult, + Task, + TaskStatus, + TextContent, + Tool, + ToolAnnotations, + ToolExecution, +) from typing_extensions import NotRequired, TypedDict # pydantic, under mcp, needs this TypedDict before Python 3.12 from s1a import config, probe, progress @@ -33,7 +57,9 @@ "(TYPESAFE_API_KEY or OPENROUTER_API_KEY) and Node for @playwright/mcp; a tool agent takes --model, --rethink and " "--episodes; a rail takes --labelled-set. decide accepts model jev (default), clm, laya, cua or laya-served; local models " "need their optional extra and no Jev key, laya-served needs LAYA_SERVED_URL. run_agent reports progress per " - "episode or browser step when the request carries a progressToken." + "episode or browser step when the request carries a progressToken. A client that supports MCP tasks may call " + "run_agent as a task: the call returns a task id at once, tasks/get reports the latest step and tasks/cancel " + "stops the run." ) @@ -44,7 +70,123 @@ async def lifespan(_server: FastMCP) -> AsyncIterator[None]: yield -server = FastMCP("s1a", instructions=INSTRUCTIONS, lifespan=lifespan) +TASK_TOOLS = frozenset({"run_agent"}) # the one long call; list_agents and decide stay plain requests (#102) +TASK_TTL_MS = 3_600_000 # the shortest task ttl granted: longer than any run's budget, reported back to the client +_TASK_REPORTER: contextvars.ContextVar[progress.Reporter | None] = contextvars.ContextVar( + "s1a_task_reporter", default=None +) +_TASK_SCOPES: dict[str, anyio.CancelScope] = {} # task id: the scope tasks/cancel cancels + + +class _TaskServer(FastMCP): + """FastMCP with ``run_agent`` callable as an MCP task (2025-11-25 experimental tasks in the 1.x SDK, #102). + + A request without task metadata takes the unchanged synchronous path. A task-augmented ``run_agent`` returns a + ``CreateTaskResult`` at once; the run continues in the session's task group, its progress messages become the + task's ``statusMessage``, and ``tasks/cancel`` cancels it through the same path as ``notifications/cancelled``. + A tool error, the busy refusal included, completes the task with an ``isError`` result, as it is the result of + the synchronous call. + Tasks live in memory: they outlive the request, not the server process or the session.""" + + async def list_tools(self) -> list[Tool]: + tools = await super().list_tools() + return [ + tool.model_copy(update={"execution": ToolExecution(taskSupport=TASK_OPTIONAL)}) + if tool.name in TASK_TOOLS + else tool + for tool in tools + ] + + async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock] | dict[str, Any]: + experimental = self.get_context().request_context.experimental + if not experimental.is_task: + return await super().call_tool(name, arguments) + metadata = experimental.task_metadata + if metadata is not None: # the store counts ttl from creation, so a shorter one would drop a run mid-way + metadata.ttl = max(metadata.ttl or 0, TASK_TTL_MS) + created = await experimental.run_task(lambda task: self._run_as_task(task, name, arguments)) + return created # the low-level handler passes a CreateTaskResult through + + async def _run_as_task(self, task: ServerTaskContext, name: str, arguments: dict[str, Any]) -> CallToolResult: + async def status(done: float, total: float | None, message: str) -> None: + await task.update_status(message) # the task's statusMessage, and notifications/tasks/status + + token = _TASK_REPORTER.set(_rate_bounded(status)) + with anyio.CancelScope() as scope: + _TASK_SCOPES[task.task_id] = scope + try: + converted = await FastMCP.call_tool(self, name, arguments) + except Exception as exc: # a tool error is the task's result, as it is the synchronous call's + return CallToolResult(content=[TextContent(type="text", text=str(exc))], isError=True) + finally: + _TASK_SCOPES.pop(task.task_id, None) + _TASK_REPORTER.reset(token) + if scope.cancelled_caught: # tasks/cancel already set the status; this result is never delivered + return CallToolResult(content=[TextContent(type="text", text=f"{name} cancelled")], isError=True) + return _call_tool_result(converted) + + +def _call_tool_result(converted: Any) -> CallToolResult: + """What the low-level handler would have built from FastMCP's converted result.""" + if isinstance(converted, CallToolResult): + return converted + if isinstance(converted, tuple): + content, structured = converted + return CallToolResult(content=list(content), structuredContent=structured) + if isinstance(converted, dict): + return CallToolResult( + content=[TextContent(type="text", text=json.dumps(converted, indent=2))], structuredContent=converted + ) + return CallToolResult(content=list(converted)) + + +class _TaskStore(InMemoryTaskStore): + """A cancelled task stays cancelled: a status or result the run writes while it unwinds is dropped, not raised.""" + + async def update_task( + self, task_id: str, status: TaskStatus | None = None, status_message: str | None = None + ) -> Task: + current = await self.get_task(task_id) + if current is not None and current.status == TASK_STATUS_CANCELLED: + return current + return await super().update_task(task_id, status=status, status_message=status_message) + + async def store_result(self, task_id: str, result: Result) -> None: + current = await self.get_task(task_id) + if current is None or current.status != TASK_STATUS_CANCELLED: + await super().store_result(task_id, result) + + +server = _TaskServer("s1a", instructions=INSTRUCTIONS, lifespan=lifespan) +# In memory, one store per server process; registers tasks/get, tasks/result, tasks/list and tasks/cancel. +# Deprecated with SEP-1686 leaving the spec for an extension; the only task API 1.x has (docs/architecture.md). +_TASKS = server._mcp_server.experimental.enable_tasks(store=_TaskStore()) # ty: ignore[deprecated] +_default_cancel = server._mcp_server.request_handlers[CancelTaskRequest] + + +async def _cancel_task(req: CancelTaskRequest) -> ServerResult: + """The SDK's tasks/cancel (scope check, terminal-state check, status) and then stop the run itself.""" + result = await _default_cancel(req) + scope = _TASK_SCOPES.get(req.params.taskId) + if scope is not None: + scope.cancel() + return result + + +server._mcp_server.request_handlers[CancelTaskRequest] = _cancel_task +_call_tool_handler = server._mcp_server.request_handlers[CallToolRequest] + + +async def _call_tool_request(req: CallToolRequest) -> ServerResult: + """A task-augmented call of a tool that takes none is a protocol error (-32601), not a tool result.""" + if req.params.task is not None and req.params.name not in TASK_TOOLS: + raise McpError( + ErrorData(code=METHOD_NOT_FOUND, message=f"{req.params.name} does not support task-augmented calls") + ) + return await _call_tool_handler(req) + + +server._mcp_server.request_handlers[CallToolRequest] = _call_tool_request # FastMCP puts a RichHandler on the root logger and the harness loggers propagate to it; only warnings pass. for _handler in logging.getLogger().handlers: _handler.setLevel(logging.WARNING) @@ -56,6 +198,16 @@ async def lifespan(_server: FastMCP) -> AsyncIterator[None]: log = logging.getLogger(__name__) +def _refuse_if_busy(label: str) -> None: + """A tool error naming the call that holds the run slot, if one does.""" + if _ONE_RUN.locked(): + running, started = _RUNNING[-1] if _RUNNING else ("another call", time.monotonic()) + raise ToolError( + f"{label} refused: {running} has been running for {time.monotonic() - started:.0f} s and this server " + "runs one call at a time; call again when it finishes, or cancel it" + ) + + @contextlib.asynccontextmanager async def _one_at_a_time(label: str) -> AsyncIterator[None]: """Hold the server's one run slot, or fail at once with a tool error naming the call that holds it. @@ -63,12 +215,7 @@ async def _one_at_a_time(label: str) -> AsyncIterator[None]: Queueing would outlast a host's tool timeout (Codex's default is 60 s) while the earlier run kept going, so a second call is refused instead of waiting. Checking and taking the lock happen without a suspension point between them, so two calls cannot both pass the check.""" - if _ONE_RUN.locked(): - running, started = _RUNNING[-1] if _RUNNING else ("another call", time.monotonic()) - raise ToolError( - f"{label} refused: {running} has been running for {time.monotonic() - started:.0f} s and this server " - "runs one call at a time; call again when it finishes, or cancel it" - ) + _refuse_if_busy(label) async with _ONE_RUN: _RUNNING.append((label, time.monotonic())) try: @@ -85,6 +232,11 @@ def _progress_reporter(ctx: Context) -> progress.Reporter | None: meta = ctx.request_context.meta if meta is None or meta.progressToken is None: return None + return _rate_bounded(ctx.report_progress) + + +def _rate_bounded(report: progress.Reporter) -> progress.Reporter: + """Increasing values only, at most one per PROGRESS_INTERVAL_S, the final point always passed on.""" last: list[float] = [float("-inf"), float("-inf")] # the last value sent and when async def send(done: float, total: float | None, message: str) -> None: @@ -93,7 +245,7 @@ async def send(done: float, total: float | None, message: str) -> None: if done <= last[0] or (not final and now - last[1] < PROGRESS_INTERVAL_S): return # the spec wants increasing values; a burst of fast steps sends one notification per second last[0], last[1] = done, now - await ctx.report_progress(done, total, message) + await report(done, total, message) return send @@ -163,8 +315,9 @@ async def run_agent(name: str, flags: list[str], ctx: Context) -> dict[str, Any] evaluation summary. The result stays an untyped object because its shape depends on the agent's front. A run can take minutes: progress goes out per episode or browser step when the request has a progressToken, a call made while another runs is refused at once, and a cancelled run releases its browser and records the cancellation. + A client with MCP task support may call it as a task and poll tasks/get instead of holding the request open. """ - reporter = _progress_reporter(ctx) + reporter = _TASK_REPORTER.get() or _progress_reporter(ctx) # argparse writes help to stdout and any stray print would too; under the stdio transport that stream is the protocol async with _one_at_a_time(f"run_agent {name}"): with ( diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 024e29a..a4c37a6 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -28,9 +28,15 @@ from mcp.client.stdio import stdio_client from mcp.shared.exceptions import McpError from mcp.shared.memory import create_connected_server_and_client_session -from mcp.types import CancelledNotification, CancelledNotificationParams, ClientNotification - -from s1a import mcp_server +from mcp.types import ( + METHOD_NOT_FOUND, + CallToolResult, + CancelledNotification, + CancelledNotificationParams, + ClientNotification, +) + +from s1a import mcp_server, progress from s1a.browser import browse from s1a import run as agents from s1a.decision_models import JevModel, ScriptedModel, ScriptedTransport @@ -583,3 +589,149 @@ async def test_run_agent_stays_an_untyped_object(self) -> None: schema = (await self._tools())["run_agent"].outputSchema self.assertEqual((schema or {}).get("type"), "object") self.assertNotIn("required", schema or {}) + + +class TestRunAgentAsTask(IsolatedAsyncioTestCase): + """``run_agent`` as an MCP task (#102): a task id at once, status while it runs, the result after, and a cancel + that stops the run the way ``notifications/cancelled`` does. Without task metadata nothing changes (above).""" + + def setUp(self) -> None: + self.tmp = Path(self.enterContext(tempfile.TemporaryDirectory())) + self.enterContext(patch.object(agents, "load", lambda name: COUNTER)) + self.enterContext(patch.object(agents, "SHOWCASE_DIR", self.tmp / "showcase")) + self.enterContext(patch.object(loop, "WORKSPACE", self.tmp / "ws")) + self.enterContext(patch.object(series, "optional_chat_model", lambda: None)) + + async def _until(self, session: ClientSession, task_id: str, *statuses: str) -> Any: + for _ in range(500): + got = await session.experimental.get_task(task_id) + if got.status in statuses: + return got + await asyncio.sleep(0.01) + self.fail(f"task {task_id} never reached {statuses}: {got.status}") + + async def test_only_run_agent_advertises_optional_task_support_and_the_server_takes_task_calls(self) -> None: + async with create_connected_server_and_client_session(mcp_server.server) as session: + listed = {tool.name: tool for tool in (await session.list_tools()).tools} + capabilities = session.get_server_capabilities() + self.assertEqual(listed["run_agent"].execution.taskSupport, "optional") + for name in ("decide", "list_agents"): + self.assertIsNone(listed[name].execution) + assert capabilities is not None and capabilities.tasks is not None + self.assertIsNotNone(capabilities.tasks.requests.tools.call) + self.assertIsNotNone(capabilities.tasks.cancel) + + async def test_a_task_returns_at_once_reports_status_and_then_the_summary(self) -> None: + started, release = asyncio.Event(), asyncio.Event() + + async def slow(name: str, flags: list[str]) -> dict[str, Any]: + await progress.report(1, 2, "episode 1/2") + started.set() + await release.wait() + return {"env": name, "episodes": 2} + + with patch.object(agents, "run_by_name", slow), patch.object(mcp_server, "PROGRESS_INTERVAL_S", 0.0): + async with create_connected_server_and_client_session(mcp_server.server) as session: + created = await asyncio.wait_for( + session.experimental.call_tool_as_task("run_agent", {"name": "counter", "flags": []}, ttl=1000), 1 + ) + task_id = created.task.taskId + await asyncio.wait_for(started.wait(), 5) + working = await session.experimental.get_task(task_id) + self.assertEqual((working.status, working.statusMessage), ("working", "episode 1/2")) + self.assertGreaterEqual(working.ttl, mcp_server.TASK_TTL_MS) # a 1 s ttl would drop a run mid-way + self.assertTrue(mcp_server._ONE_RUN.locked()) # the task holds the slot the synchronous call does + release.set() + await self._until(session, task_id, "completed") + result = await session.experimental.get_task_result(task_id, CallToolResult) + self.assertFalse(result.isError, result.content) + self.assertEqual(_rows(result), {"env": "counter", "episodes": 2}) + self.assertFalse(mcp_server._ONE_RUN.locked()) + + async def test_a_real_tool_agent_run_as_a_task_returns_the_same_summary(self) -> None: + flags = [*RANDOM_RUN[:5], "2", "--showcase"] + async with create_connected_server_and_client_session(mcp_server.server) as session: + created = await session.experimental.call_tool_as_task("run_agent", {"name": "counter", "flags": flags}) + await self._until(session, created.task.taskId, "completed") + result = await session.experimental.get_task_result(created.task.taskId, CallToolResult) + self.assertFalse(result.isError, result.content) + summary = _rows(result) + self.assertEqual((summary["env"], summary["policy"], summary["episodes"]), ("counter", "random", 2)) + self.assertTrue(Path(summary["job_dir"]).is_dir()) + + async def test_bad_flags_and_a_busy_server_complete_the_task_with_a_tool_error(self) -> None: + started = asyncio.Event() + + async def slow(name: str, flags: list[str]) -> dict[str, Any]: + started.set() + await asyncio.Event().wait() + return {} + + async with create_connected_server_and_client_session(mcp_server.server) as session: + bad = await session.experimental.call_tool_as_task("run_agent", {"name": "counter", "flags": ["--help"]}) + await self._until(session, bad.task.taskId, "completed") + bad_result = await session.experimental.get_task_result(bad.task.taskId, CallToolResult) + with patch.object(agents, "run_by_name", slow): + first = await session.experimental.call_tool_as_task("run_agent", {"name": "counter", "flags": []}) + await asyncio.wait_for(started.wait(), 5) + second = await session.experimental.call_tool_as_task("run_agent", {"name": "counter", "flags": []}) + await self._until(session, second.task.taskId, "completed") + busy = await session.experimental.get_task_result(second.task.taskId, CallToolResult) + await session.experimental.cancel_task(first.task.taskId) + self.assertTrue(bad_result.isError) + self.assertIn("bad flags", bad_result.content[0].text) + self.assertTrue(busy.isError) + self.assertIn("run_agent counter has been running for", busy.content[0].text) + + async def test_tasks_cancel_tears_the_browser_down_and_frees_the_slot(self) -> None: + started, torn_down = asyncio.Event(), [] + + async def run_agent(agent: Any, inputs: dict[str, Any]) -> dict[str, Any]: + started.set() + await asyncio.Event().wait() + return {} + + async def cleanup() -> None: + await asyncio.sleep(0.01) + torn_down.append("browser") + + async def release(conversation_id: str) -> None: + await asyncio.sleep(0) + torn_down.append("session") + + agent = SimpleNamespace(ensure_initialized=AsyncMock(), cleanup_task_resources=cleanup) + runner = SimpleNamespace(run_agent=run_agent, release=release) + + async def browser_run(name: str, flags: list[str]) -> dict[str, Any]: + return dict(await browse.run_task(agent, "goal", timeout_s=60, logs_dir=self.tmp)) # ty: ignore[invalid-argument-type] + + with ( + patch.object(browse, "Runner", runner), + patch.object(browse, "save_final_screenshot", AsyncMock()), + patch.object(agents, "run_by_name", browser_run), + ): + async with create_connected_server_and_client_session(mcp_server.server) as session: + created = await session.experimental.call_tool_as_task("run_agent", {"name": "flights", "flags": []}) + await asyncio.wait_for(started.wait(), 5) + cancelled = await session.experimental.cancel_task(created.task.taskId) + for _ in range(100): + if not mcp_server._ONE_RUN.locked(): + break + await asyncio.sleep(0.01) + self.assertFalse(mcp_server._ONE_RUN.locked()) + after = await session.experimental.get_task(created.task.taskId) + with self.assertRaises(McpError): # a cancelled task cannot be cancelled again + await session.experimental.cancel_task(created.task.taskId) + self.assertEqual(cancelled.status, "cancelled") + self.assertEqual(after.status, "cancelled") + self.assertEqual(torn_down, ["browser", "session"]) + self.assertEqual(mcp_server._TASK_SCOPES, {}) + + async def test_a_task_call_of_decide_is_a_protocol_error(self) -> None: + async with create_connected_server_and_client_session(mcp_server.server) as session: + with self.assertRaises(McpError) as caught: + await session.experimental.call_tool_as_task( + "decide", {"state": {}, "options": {"a": "a"}, "rules": "pick"} + ) + self.assertEqual(caught.exception.error.code, METHOD_NOT_FOUND) + self.assertIn("does not support task-augmented calls", caught.exception.error.message)