diff --git a/CHANGELOG.md b/CHANGELOG.md index 08c930b..85c53c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,12 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); ver ### Fixed +- `s1a-mcp` long runs: `run_agent` reports progress per episode or browser step when the request has a + `progressToken`; a call made while another runs fails at once with a tool error instead of queueing past the + host's timeout; a cancelled run finishes its browser and model teardown (it was cut at the first `await`) and + records the cancellation in the job `result.json` or `answer.json` even when no episode finished; cancelled + seeds count in `n_total_trials`. The server instructions put serialization and the timeout advice in their + first 512 characters, and docs/skills.md gives the Codex and Claude Code settings. - `decide` rejects repeated `--option` keys before model setup instead of silently replacing an earlier choice description. The usage error names the duplicate key without printing either description. - Malformed ticket-router JSONL now names the dataset, physical line and JSON column without printing ticket contents. diff --git a/README.md b/README.md index b465718..afa6fe4 100644 --- a/README.md +++ b/README.md @@ -132,7 +132,9 @@ folder with the score. Without a key, `--model cua` answers in process after `uv codex mcp add s1a -- uv run --project /path/to/system1-agents s1a-mcp ``` -It serves three tools, `list_agents`, `run_agent` and `decide`, to any host that speaks MCP. +It serves three tools, `list_agents`, `run_agent` and `decide`, to any host that speaks MCP. Run `uv sync` in the +checkout once before registering it, and raise the host's tool timeout: a browser run takes minutes, beyond Codex's +60 s default. The settings for Codex and Claude Code: [docs/skills.md](docs/skills.md#host-timeouts). ## Build your own diff --git a/docs/skills.md b/docs/skills.md index 58c12fb..bebb4a9 100644 --- a/docs/skills.md +++ b/docs/skills.md @@ -80,8 +80,37 @@ or `cua`; callers that omit it keep using Jev. `clm` needs `CLM_URL` (a running `LAYA_SERVED_URL`. For local decisions, install the matching extra in the server's checkout (`uv sync --extra laya` or `uv sync --extra cua`) and pass `model="laya"` or `model="cua"`; no Jev API key is needed. The first local call may download the checkpoint. Each call loads and closes its model; `ms` measures the -decision, not model loading. Agent runs and decisions are serialized, and model output stays off the stdio -protocol stream. +decision, not model loading. Model output stays off the stdio protocol stream. + +The server runs one call at a time. A `run_agent` or `decide` that arrives while another call runs fails at once +with a tool error naming the running call and how long it has run; it does not queue. When the request carries a +`progressToken`, `run_agent` sends `notifications/progress` once per finished episode (`episode 3/20`) or decided +browser step (`step 7/100`), at most one a second. A cancelled call (`notifications/cancelled`) releases the slot, +closes the browser and the decision model it started, and records the cancellation: a tool agent's job +`result.json` has `"cancelled": true` and counts the unfinished episodes in `n_cancelled_trials`; a browser agent's +`answer.json` has `"cancelled": true`. + +#### Host timeouts + +`flights` budgets 180 s and `allrecipes` 300 s, and a tool agent's run grows with `--episodes`. Run `uv sync` in the +checkout once before registering the server, so the first start does not spend minutes resolving the `openjiuwen` +git pin inside the host's startup wait. + +[Codex](https://developers.openai.com/codex/mcp) stops a tool call after 60 s and a server start after 10 s by +default. Raise both for this server in `~/.codex/config.toml` +([config reference](https://developers.openai.com/codex/config-reference)): + +```toml +[mcp_servers.s1a] +command = "uv" +args = ["run", "--project", "/path/to/system1-agents", "s1a-mcp"] +startup_timeout_sec = 60 +tool_timeout_sec = 600 +``` + +[Claude Code](https://code.claude.com/docs/en/mcp) has no short tool limit by default. To set one, add +`"timeout": 600000` (milliseconds) to the server's `.mcp.json` entry or export `MCP_TOOL_TIMEOUT=600000`; give the +first start more room with `MCP_TIMEOUT=60000`. Each tool carries a `title` and MCP behavior hints, which hosts such as Codex (`default_tools_approval_mode = "writes"`) and Claude use to decide what to auto-approve. `list_agents` is read-only, idempotent and closed-world: it diff --git a/pyproject.toml b/pyproject.toml index de9e176..eaf3091 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,12 +10,13 @@ requires-python = ">=3.11" dependencies = [ # Tag on the jiuwen-jev branch of the ThinkFlowLab fork; CONTRIBUTING.md § The openjiuwen pin lists what it adds over upstream develop. "openjiuwen @ git+https://github.com/ThinkFlowLab/agent-core@jj-0.1.0", - # Imported directly; openjiuwen depends on all three, so they add nothing to the install. + # Imported directly; openjiuwen (and mcp, for anyio) already depends on each, so they add nothing to the install. "httpx>=0.28", "python-dotenv>=1.0", # MCP Python SDK 1.x: openjiuwen 0.1.18 needs fastmcp, whose 3.x line requires mcp<2 (fastmcp 4.x needs mcp>=2). # Lift the bound with the 2.x port; docs/architecture.md § MCP SDK 2.x lists the call sites and the unblock condition. "mcp>=1.26,<2", + "anyio>=4", "pillow>=10", ] diff --git a/s1a/browser/browse.py b/s1a/browser/browse.py index cd81551..df162cb 100644 --- a/s1a/browser/browse.py +++ b/s1a/browser/browse.py @@ -16,6 +16,7 @@ from urllib.parse import urlparse from uuid import uuid4 +import anyio from openjiuwen.core.common.logging import logger from openjiuwen.core.foundation.llm import Model from openjiuwen.core.runner import Runner @@ -211,8 +212,9 @@ async def run_task(agent: DeepAgent, goal: str, *, timeout_s: float, logs_dir: P try: await save_final_screenshot(agent, answer, logs_dir) # before the cleanup closes the page finally: # a cancellation during the screenshot still releases the browser and the session - await agent.cleanup_task_resources() - await Runner.release(conversation_id) + with anyio.CancelScope(shield=True): # a cancelled task (an MCP host's cancel) must not cut the teardown + await agent.cleanup_task_resources() + await Runner.release(conversation_id) answer["final"] = str(result.get("output") or "") answer["ok"] = result.get("result_type") == "answer" and bool(answer["final"]) if not answer["ok"]: @@ -253,7 +255,9 @@ async def browse( case "jev" | "clm" | "laya" | "laya-served" | "cua" | "omnijev": if decision_model is None: raise RuntimeError(f"--model {model_name} needs a decision model") - slot_model = BrowserDecisionModel(spec, policy, counted, decision_model=decision_model, value_model=None) + slot_model = BrowserDecisionModel( + spec, policy, counted, decision_model=decision_model, value_model=None, max_steps=max_steps + ) agent = create_browser_agent( slot_model, language=spec.language, @@ -396,7 +400,9 @@ def policy_from_args(args: argparse.Namespace) -> BrowserPolicy: async def play(spec: BrowserAgentSpec, args: argparse.Namespace) -> Answer: - """One task inside an already started Runner: the answer without the ticks, which ``decision_ticks.json`` holds.""" + """One task inside an already started Runner: the answer without the ticks, which ``decision_ticks.json`` holds. + + A cancelled task writes ``answer.json`` with ``"cancelled": true`` and re-raises.""" if args.profile_out is not None and args.model == "llm": raise RuntimeError("--profile-out needs a decision model: the profiler times the slot model's turns") chat = chat_model_from_env() @@ -415,9 +421,15 @@ async def play(spec: BrowserAgentSpec, args: argparse.Namespace) -> Answer: chat=chat, decision_model=decision_model, ) + except asyncio.CancelledError: # a host cancelled the call: answer.json says so instead of keeping a stale answer + args.logs_dir.mkdir(parents=True, exist_ok=True) + cancelled = {"ok": False, "cancelled": True, "error": "cancelled", "agent": spec.name, "model": args.model} + (args.logs_dir / "answer.json").write_text(json.dumps(cancelled, indent=1), encoding="utf-8") + raise finally: - if decision_model is not None: - await decision_model.close() + with anyio.CancelScope(shield=True): # a cancelled task still closes its decision model + if decision_model is not None: + await decision_model.close() if profiler is not None: profiler.detach() result = {key: value for key, value in answer.items() if key != "ticks"} diff --git a/s1a/browser/decision_model.py b/s1a/browser/decision_model.py index 67f681b..bd119a2 100644 --- a/s1a/browser/decision_model.py +++ b/s1a/browser/decision_model.py @@ -32,6 +32,7 @@ from openjiuwen.core.foundation.llm import AssistantMessage, AssistantMessageChunk, ToolCall from openjiuwen.core.foundation.llm.model import Model from openjiuwen.harness.tools.browser_move.utils.parsing import extract_json_object +from s1a import progress from s1a.browser import prompts from s1a.browser.action_space import ( NONE_VALUE, @@ -172,6 +173,7 @@ def __init__( *, decision_model: DecisionModel, value_model: Model | None, + max_steps: int | None = None, ) -> None: """``policy.goal_value_cache`` offers goal-extracted values to Jev as a choice head; off means every typed value comes from the chat model given the field and page context. ``policy.prefetch_values`` makes that @@ -180,7 +182,8 @@ def __init__( sends each action and the next probe as one ``browser_run_code_unsafe`` call when that tool is offered, skipping the runtime's target validation and one transport round trip per step. ``value_model`` answers the typed-value calls when given; the fallback chat model otherwise. - ``decision_model`` decides every browser step; its ``name`` (``jev``, ``laya``) is the ``--model`` value.""" + ``decision_model`` decides every browser step; its ``name`` (``jev``, ``laya``) is the ``--model`` value. + Each decided step is reported to ``s1a.progress``, out of ``max_steps`` when given.""" super().__init__(fallback.model_client_config, fallback.model_config) self._spec = spec self._fallback = fallback @@ -190,6 +193,7 @@ def __init__( self._batch_actions = policy.batch_actions self._rethink_on = policy.rethink_on self._recovery_limits = policy.recovery_limits + self._max_steps = max_steps self._language = spec.language self._decision_model = decision_model # the chat model writes the bounded recovery's plan, as the tool front's RethinkRail does @@ -365,6 +369,8 @@ async def _decide_message(self, messages: Any) -> AssistantMessage: if screenshot_ms is not None: # an image-reading model: the capture is part of the step's cost record["screenshot_ms"] = screenshot_ms run.ticks.append(record) + step = len(run.ticks) + await progress.report(step, self._max_steps, f"step {step}/{self._max_steps or '?'}") if move.operation == "WAIT": run.consecutive_waits += 1 wait_entry = {"action": "wait", "kind": "wait", "text": None, "page_changed": None} diff --git a/s1a/jobs.py b/s1a/jobs.py index 5c86afc..3f91677 100644 --- a/s1a/jobs.py +++ b/s1a/jobs.py @@ -89,12 +89,28 @@ def summarize(episodes: list[Episode]) -> dict[str, Any]: } -def write_job(eval_name: str, episodes: list[Episode], *, results_dir: Path) -> Path: +def write_job( + eval_name: str, + episodes: list[Episode], + *, + results_dir: Path, + cancelled: int = 0, + policy: str | None = None, +) -> Path: """One Harbor job folder per run: config.json, result.json, summary.json, one trial folder per episode. + ``cancelled`` counts the planned trials a cancellation stopped before they finished; a job with any has + ``"cancelled": true`` in its result.json. ``n_total_trials`` (and the dataset ``n_tasks``) include those + cancelled trials so Harbor-style readers see a consistent total. A cancellation that finished no episode + still writes the job when ``cancelled`` is positive; pass ``policy`` so the folder can name the agent. + The trial shape follows agent-eval's runner so Nullius reads both roots alike. """ - policy = episodes[0].policy + if not episodes and cancelled <= 0: + raise ValueError("write_job needs at least one episode, or a positive cancelled count") + if not episodes and policy is None: + raise ValueError("write_job with no episodes needs policy") + policy = policy if policy is not None else episodes[0].policy agent_name = f"s1a-evals/{policy}" job_dir = ( results_dir / eval_name / f"{datetime.now():%Y-%m-%d__%H-%M-%S-%f}__{policy}" @@ -152,28 +168,32 @@ def write_job(eval_name: str, episodes: list[Episode], *, results_dir: Path) -> } (trial_dir / "result.json").write_text(json.dumps(trial, ensure_ascii=False, indent=2), encoding="utf-8") finished_at = now_iso() + started_at = episodes[0].started_at if episodes else finished_at errored = sum(episode.error is not None for episode in episodes) + n_total = len(episodes) + cancelled job = { "id": str(uuid.uuid4()), - "started_at": episodes[0].started_at, + "started_at": started_at, "updated_at": finished_at, "finished_at": finished_at, - "n_total_trials": len(episodes), + "n_total_trials": n_total, "stats": { "n_completed_trials": len(episodes) - errored, "n_errored_trials": errored, "n_running_trials": 0, "n_pending_trials": 0, - "n_cancelled_trials": 0, + "n_cancelled_trials": cancelled, "n_retries": 0, }, + "cancelled": cancelled > 0, } (job_dir / "result.json").write_text(json.dumps(job, indent=2), encoding="utf-8") config = { "job_name": job_dir.name, "agents": [{"name": agent_name, "model_name": policy}], - "datasets": [{"name": eval_name, "n_tasks": len(episodes)}], # the dict shape Nullius's parse_job reads + "datasets": [{"name": eval_name, "n_tasks": n_total}], # the dict shape Nullius's parse_job reads } (job_dir / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") - (job_dir / "summary.json").write_text(json.dumps(summarize(episodes), indent=2), encoding="utf-8") + summary = summarize(episodes) if episodes else {"episodes": 0, "scored": 0, "errors": 0} + (job_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") return job_dir diff --git a/s1a/mcp_server.py b/s1a/mcp_server.py index a5be356..53013ff 100644 --- a/s1a/mcp_server.py +++ b/s1a/mcp_server.py @@ -8,27 +8,32 @@ import contextlib import logging import sys +import time from collections.abc import AsyncIterator from typing import Any, Literal import s1a.console # first: routes the harness logs to files before openjiuwen loads and logs to the console -from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp import Context, FastMCP +from mcp.server.fastmcp.exceptions import ToolError from mcp.types import ToolAnnotations from typing_extensions import NotRequired, TypedDict # pydantic, under mcp, needs this TypedDict before Python 3.12 -from s1a import config, probe +from s1a import config, probe, progress from s1a.decision_models import build_model from s1a import run as agents from s1a.run import started_runner INSTRUCTIONS = ( "S1A agents: System 1 decision models (TypeSafe Jev, Laya, Cua-S1) in the model slot of openJiuwen agents. Use run_agent for a page task with enumerable " - "controls, a game or a quiz, and decide for one selection over options you enumerate. Not for arithmetic, " + "controls, a game or a quiz, and decide for one selection over options you enumerate. One run or decision at a " + "time per server; a second call fails at once, naming the first. A run can take minutes: raise the host's tool " + "timeout to 600 s (Codex tool_timeout_sec, Claude Code MCP_TOOL_TIMEOUT). Not for arithmetic, " "constraint puzzles or free-text generation. list_agents gives every agent's flags: a browser agent takes " "--model jev --goal '...' and needs a chat-model key (OPENAI_API_KEY or LLM_API_KEY plus MODEL_NAME), a Jev key " "(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. Runs and decisions go one at a time per server." + "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." ) @@ -46,6 +51,51 @@ async def lifespan(_server: FastMCP) -> AsyncIterator[None]: _ONE_RUN = ( asyncio.Lock() ) # ponytail: runs are serialized; per-run browser and workspace isolation before parallel tasks +_RUNNING: list[tuple[str, float]] = [] # the call holding _ONE_RUN and its monotonic start, for the busy error +PROGRESS_INTERVAL_S = 1.0 # at most one progress notification per second, the last point of a run always sent +log = logging.getLogger(__name__) + + +@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. + + 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" + ) + async with _ONE_RUN: + _RUNNING.append((label, time.monotonic())) + try: + yield + except asyncio.CancelledError: + log.warning("%s cancelled after %.0f s", label, time.monotonic() - _RUNNING[-1][1]) + raise + finally: + _RUNNING.pop() + + +def _progress_reporter(ctx: Context) -> progress.Reporter | None: + """Forward progress to the host when its request carries a progressToken: increasing values, rate-bounded.""" + meta = ctx.request_context.meta + if meta is None or meta.progressToken is None: + return None + last: list[float] = [float("-inf"), float("-inf")] # the last value sent and when + + async def send(done: float, total: float | None, message: str) -> None: + now = time.monotonic() + final = total is not None and done >= total + 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) + + return send class AgentRow(TypedDict): @@ -106,15 +156,21 @@ def _agent_rows() -> list[AgentRow]: @server.tool(title="Run an S1A agent", annotations=RUN_AGENT_HINTS) -async def run_agent(name: str, flags: list[str]) -> dict[str, Any]: +async def run_agent(name: str, flags: list[str], ctx: Context) -> dict[str, Any]: """Run one agent with the flags `s1a run ` takes, e.g. ["--model", "jev", "--goal", "..."]. A browser agent returns its answer, a tool agent its series summary with the job folder it wrote, a rail its - evaluation summary. The result stays an untyped object because its shape depends on the agent's front. + 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. """ + reporter = _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_RUN: - with contextlib.redirect_stdout(sys.stderr): + async with _one_at_a_time(f"run_agent {name}"): + with ( + contextlib.redirect_stdout(sys.stderr), + progress.reporting(reporter) if reporter is not None else contextlib.nullcontext(), + ): try: return await agents.run_by_name(name, flags) except SystemExit as exc: @@ -134,7 +190,7 @@ async def decide( running Laya server (LAYA_SERVED_URL). Local model loading is excluded from the reported latency. """ - async with _ONE_RUN: + async with _one_at_a_time(f"decide ({model})"): with contextlib.redirect_stdout(sys.stderr): # local SDK loading must not write to the stdio protocol decision_model = await asyncio.to_thread(build_model, model) try: diff --git a/s1a/progress.py b/s1a/progress.py new file mode 100644 index 0000000..8e5739d --- /dev/null +++ b/s1a/progress.py @@ -0,0 +1,41 @@ +# coding: utf-8 +"""Progress for a long run: a host that asked for it hears about each episode or browser step, nobody else does. + +The tool series and the browser decision model call ``report``; it is a no-op unless an entry point installed a +reporter with ``reporting``. ``s1a-mcp`` installs one when the request carries a ``progressToken``. +""" + +from __future__ import annotations + +import logging +from collections.abc import Awaitable, Callable, Iterator +from contextlib import contextmanager +from contextvars import ContextVar + +Reporter = Callable[[float, float | None, str], Awaitable[None]] + +log = logging.getLogger(__name__) +_REPORTER: ContextVar[Reporter | None] = ContextVar("s1a_progress_reporter", default=None) + + +@contextmanager +def reporting(reporter: Reporter) -> Iterator[None]: + """Route ``report`` calls made in this context (and the tasks it starts) to ``reporter``.""" + token = _REPORTER.set(reporter) + try: + yield + finally: + _REPORTER.reset(token) + + +async def report(done: float, total: float | None, message: str) -> None: + """One progress point: ``done`` of ``total`` (None when unknown) with a short message such as ``episode 3/20``. + + A reporter that fails is logged and dropped for this point; progress never ends a run.""" + reporter = _REPORTER.get() + if reporter is None: + return + try: + await reporter(done, total, message) + except Exception as exc: # a closed transport or a host that went away must not fail the run itself + log.warning("progress report %s failed: %s", message, exc) diff --git a/s1a/tool/series.py b/s1a/tool/series.py index 18f1e05..42e59f2 100644 --- a/s1a/tool/series.py +++ b/s1a/tool/series.py @@ -7,6 +7,9 @@ import asyncio from pathlib import Path +import anyio + +from s1a import progress from s1a.decision_models import DecisionModel, build_model from s1a.config import first_env, optional_chat_model from s1a.jobs import Episode, summarize, write_job @@ -76,8 +79,9 @@ async def play(spec: ToolAgentSpec, args: argparse.Namespace, *, results_dir: Pa Jev, Laya, Cua and the rule share one model for the series (warmed once, closed at the end); chance gets a fresh stream per episode from that episode's seed. The keys and the price variables are checked before the series is - built. Whatever ends the seed loop, the finished episodes are written; a series whose every episode errored writes - its job folder, then raises.""" + built. Whatever ends the seed loop, the finished episodes are written, a cancellation included (its job records the + unfinished seeds as cancelled); a series whose every episode errored writes its job folder, then raises. Each + finished episode is reported to ``s1a.progress``.""" env_prices() chat = optional_chat_model() if args.model == "llm" and chat is None: @@ -105,6 +109,8 @@ def model_for(seed: int) -> DecisionModel | None: return shared episodes: list[Episode] = [] + total = len(run.seeds) + cancelled = 0 try: if shared is not None: await shared.warm() @@ -127,12 +133,22 @@ def model_for(seed: int) -> DecisionModel | None: ) run.annotate(env, episode) episodes.append(episode) + await progress.report(len(episodes), total, f"episode {len(episodes)}/{total}") + except asyncio.CancelledError: # a host cancelled the call: the seeds not finished are recorded as cancelled + cancelled = total - len(episodes) + raise finally: - if shared is not None: - await shared.close() - if episodes: - await asyncio.to_thread(price_episodes, episodes) # the price catalogue fetch is a blocking HTTP call - job_dir = write_job(spec.name, episodes, results_dir=results_dir) + with anyio.CancelScope(shield=True): # a cancelled series still closes its model and writes what it finished + if shared is not None: + await shared.close() + if episodes or cancelled: # cancel during warm-up / first episode still writes a cancelled job + if ( + episodes and not cancelled + ): # the price catalogue fetch is a blocking HTTP call; a cancelled run leaves cost unknown + await asyncio.to_thread(price_episodes, episodes) + job_dir = write_job( + spec.name, episodes, results_dir=results_dir, cancelled=cancelled, policy=args.model + ) if not episodes: raise RuntimeError("the series selected no episodes") if all(episode.error is not None for episode in episodes): diff --git a/tests/test_jobs.py b/tests/test_jobs.py index b120e7e..cc49c35 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -141,3 +141,22 @@ def test_job_folder_takes_the_episode_frames(self) -> None: trial = next(p for p in job_dir.iterdir() if p.is_dir()) self.assertEqual(sorted(p.name for p in (trial / "agent" / "frames").iterdir()), ["000.png", "001.png"]) self.assertFalse(frames.exists()) + + def test_cancelled_trials_count_in_the_job_total(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + job_dir = write_job("counter", [_episode(0, 1.0, 3)], results_dir=Path(tmp), cancelled=2) + job = json.loads((job_dir / "result.json").read_text()) + self.assertEqual(job["n_total_trials"], 3) + self.assertEqual(job["stats"]["n_cancelled_trials"], 2) + self.assertTrue(job["cancelled"]) + config = json.loads((job_dir / "config.json").read_text()) + self.assertEqual(config["datasets"], [{"name": "counter", "n_tasks": 3}]) + + def test_a_cancel_with_no_finished_episode_still_writes_the_job(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + job_dir = write_job("counter", [], results_dir=Path(tmp), cancelled=3, policy="random") + job = json.loads((job_dir / "result.json").read_text()) + self.assertEqual((job["n_total_trials"], job["stats"]["n_cancelled_trials"]), (3, 3)) + self.assertTrue(job["cancelled"]) + self.assertTrue(job_dir.name.endswith("__random")) + self.assertEqual(json.loads((job_dir / "summary.json").read_text())["episodes"], 0) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index ccb9fbd..024e29a 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,8 +1,10 @@ # coding: utf-8 -"""The MCP server over an in-memory session: the three tools, decide against a fake client, and the agent listing.""" +"""The MCP server over an in-memory session: the three tools, decide against a fake client, the agent listing, and +long runs (progress, the busy refusal, cancellation).""" from __future__ import annotations +import asyncio import contextlib import importlib.util import io @@ -19,16 +21,21 @@ from pathlib import Path from typing import Any from unittest import IsolatedAsyncioTestCase, TestCase, skipUnless +from types import SimpleNamespace from unittest.mock import AsyncMock, patch from mcp import ClientSession, StdioServerParameters 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 s1a.browser import browse from s1a import run as agents from s1a.decision_models import JevModel, ScriptedModel, ScriptedTransport from s1a.tool import loop, series +from s1a.agents.flights import SPEC as FLIGHTS from support import COUNTER RANDOM_RUN = ["--model", "random", "--rethink", "off", "--episodes", "1", "--showcase"] @@ -287,6 +294,205 @@ def load(name: str) -> agents.Spec: self.assertEqual(rows["game2048"]["front"], "tool") +async def _cancel(session: ClientSession, request_id: int) -> None: + params = CancelledNotificationParams(requestId=request_id, reason="host gave up") + await session.send_notification(ClientNotification(CancelledNotification(params=params))) + + +class TestLongRuns(IsolatedAsyncioTestCase): + """What a host with a tool timeout needs: progress, a fast refusal while busy, and a clean cancellation.""" + + 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)) + + def test_the_first_512_characters_say_runs_are_serialized_and_need_a_raised_timeout(self) -> None: + head = mcp_server.INSTRUCTIONS[:512] + self.assertIn("One run or decision at a time per server", head) + for setting in ("tool timeout", "tool_timeout_sec", "MCP_TOOL_TIMEOUT"): + self.assertIn(setting, head) + + async def test_a_progress_token_gets_increasing_progress_per_episode(self) -> None: + points: list[tuple[float, float | None, str | None]] = [] + + async def heard(done: float, total: float | None, message: str | None) -> None: + points.append((done, total, message)) + + flags = [*RANDOM_RUN[:5], "3", "--showcase"] + with patch.object(mcp_server, "PROGRESS_INTERVAL_S", 0.0): + async with create_connected_server_and_client_session(mcp_server.server) as session: + result = await session.call_tool( + "run_agent", {"name": "counter", "flags": flags}, progress_callback=heard + ) + self.assertFalse(result.isError, result.content) + self.assertEqual(_rows(result)["episodes"], 3) + self.assertEqual(points, [(1, 3, "episode 1/3"), (2, 3, "episode 2/3"), (3, 3, "episode 3/3")]) + + async def test_without_a_token_no_progress_is_sent(self) -> None: + flags = [*RANDOM_RUN[:5], "3", "--showcase"] + with patch.object(mcp_server.Context, "report_progress", new_callable=AsyncMock) as sent: + async with create_connected_server_and_client_session(mcp_server.server) as session: + result = await session.call_tool("run_agent", {"name": "counter", "flags": flags}) + self.assertFalse(result.isError, result.content) + sent.assert_not_awaited() + + async def test_the_rate_bound_keeps_the_last_point_and_drops_repeats(self) -> None: + ctx = SimpleNamespace( + request_context=SimpleNamespace(meta=SimpleNamespace(progressToken="t")), report_progress=AsyncMock() + ) + send = mcp_server._progress_reporter(ctx) # ty: ignore[invalid-argument-type] + assert send is not None + for done in (1, 1, 2, 3, 4): + await send(done, 4, f"step {done}/4") + sent = [call.args for call in ctx.report_progress.await_args_list] + self.assertEqual(sent, [(1, 4, "step 1/4"), (4, 4, "step 4/4")]) # 2 and 3 fell inside the interval + + async def test_a_call_while_a_run_holds_the_slot_is_refused_at_once(self) -> None: + started = asyncio.Event() + + async def slow(name: str, flags: list[str]) -> dict[str, Any]: + started.set() + await asyncio.Event().wait() + return {} + + with ( + patch.object(agents, "run_by_name", slow), + patch.object(mcp_server, "build_model", return_value=ScriptedModel(choose="a")), + ): + async with create_connected_server_and_client_session(mcp_server.server) as session: + request_id = session._request_id # the id the next request takes + first = asyncio.create_task(session.call_tool("run_agent", {"name": "counter", "flags": []})) + await asyncio.wait_for(started.wait(), 5) + busy = await asyncio.wait_for( + session.call_tool("decide", {"state": {}, "options": {"a": "a"}, "rules": "pick"}), 1 + ) + again = await asyncio.wait_for( + session.call_tool("run_agent", {"name": "counter", "flags": RANDOM_RUN}), 1 + ) + await _cancel(session, request_id) + with contextlib.suppress(McpError): + await asyncio.wait_for(first, 5) + for refused in (busy, again): + self.assertTrue(refused.isError) + self.assertIn("run_agent counter has been running for", refused.content[0].text) + self.assertIn("one call at a time", refused.content[0].text) + + async def test_a_cancelled_browser_run_tears_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) # a real teardown awaits; the cancelled task must not cut it short + 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(mcp_server, "build_model", return_value=ScriptedModel(choose="a")), + ): + async with create_connected_server_and_client_session(mcp_server.server) as session: + with patch.object(agents, "run_by_name", browser_run): + request_id = session._request_id + first = asyncio.create_task(session.call_tool("run_agent", {"name": "flights", "flags": []})) + await asyncio.wait_for(started.wait(), 5) + await _cancel(session, request_id) + with contextlib.suppress(McpError): + await asyncio.wait_for(first, 5) + 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 asyncio.wait_for( + session.call_tool("decide", {"state": {}, "options": {"a": "a"}, "rules": "pick"}), 5 + ) + self.assertEqual(torn_down, ["browser", "session"]) + self.assertFalse(after.isError, after.content) + self.assertEqual(_rows(after)["choice"], "a") + + async def test_a_cancelled_series_writes_its_job_as_cancelled(self) -> None: + real, blocked = loop.run_episode, asyncio.Event() + + async def second_blocks(*args: Any, **kwargs: Any) -> Any: + if kwargs["seed"] > 0: + blocked.set() + await asyncio.Event().wait() + return await real(*args, **kwargs) + + parsed = series.parser(COUNTER).parse_args([*RANDOM_RUN[:5], "3"]) + with patch.object(series, "run_episode", second_blocks): + task = asyncio.create_task(series.play(COUNTER, parsed, results_dir=self.tmp / "results")) + await asyncio.wait_for(blocked.wait(), 5) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + (job_dir,) = (self.tmp / "results" / COUNTER.name).iterdir() + job = json.loads((job_dir / "result.json").read_text(encoding="utf-8")) + self.assertTrue(job["cancelled"]) + self.assertEqual((job["n_total_trials"], job["stats"]["n_cancelled_trials"]), (3, 2)) + + async def test_a_cancelled_browser_play_writes_a_cancelled_answer(self) -> None: + async def blocks(*args: Any, **kwargs: Any) -> Any: + await asyncio.Event().wait() + + args = SimpleNamespace( + profile_out=None, model="llm", goal="g", timeout=60.0, max_steps=5, logs_dir=self.tmp / "logs", headed=False + ) + with ( + patch.object(browse, "chat_model_from_env", lambda: None), + patch.object(browse, "policy_from_args", lambda args: None), + patch.object(browse, "browse", blocks), + ): + task = asyncio.create_task(browse.play(FLIGHTS, args)) # ty: ignore[invalid-argument-type] + await asyncio.sleep(0.01) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + answer = json.loads((self.tmp / "logs" / "answer.json").read_text(encoding="utf-8")) + self.assertEqual((answer["ok"], answer["cancelled"], answer["model"]), (False, True, "llm")) + + async def test_a_cancel_before_any_episode_still_writes_a_cancelled_job(self) -> None: + """Cancel during the first episode: no finished trial, but every planned seed is recorded.""" + real, blocked = series.run_episode, asyncio.Event() + + async def blocks_first(*args: Any, **kwargs: Any) -> Any: + blocked.set() + await asyncio.Event().wait() + return await real(*args, **kwargs) + + parsed = series.parser(COUNTER).parse_args([*RANDOM_RUN[:5], "3"]) + with patch.object(series, "run_episode", blocks_first): + task = asyncio.create_task(series.play(COUNTER, parsed, results_dir=self.tmp / "results")) + await asyncio.wait_for(blocked.wait(), 5) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + (job_dir,) = (self.tmp / "results" / COUNTER.name).iterdir() + job = json.loads((job_dir / "result.json").read_text(encoding="utf-8")) + self.assertTrue(job["cancelled"]) + self.assertEqual((job["n_total_trials"], job["stats"]["n_cancelled_trials"]), (3, 3)) + self.assertEqual(job["stats"]["n_completed_trials"], 0) + config = json.loads((job_dir / "config.json").read_text(encoding="utf-8")) + self.assertEqual(config["datasets"], [{"name": COUNTER.name, "n_tasks": 3}]) + + class TestSdkBound(TestCase): """mcp stays on 1.x until openjiuwen's fastmcp allows 2.x; docs/architecture.md § MCP SDK 2.x tracks the port.""" diff --git a/uv.lock b/uv.lock index 7d99036..e18420b 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -5221,6 +5221,7 @@ name = "system1-agents" version = "0.1.0" source = { editable = "." } dependencies = [ + { name = "anyio" }, { name = "httpx" }, { name = "mcp" }, { name = "openjiuwen" }, @@ -5289,6 +5290,7 @@ requires-dist = [ { name = "accelerate", marker = "extra == 'omnijev'" }, { name = "ai2thor", marker = "extra == 'alfworld-visual'", specifier = "==2.1.0" }, { name = "alfworld", marker = "extra == 'alfworld'", specifier = ">=0.4" }, + { name = "anyio", specifier = ">=4" }, { name = "cua-s1", marker = "extra == 'cua'", git = "https://github.com/trycua/cua.git?subdirectory=libs%2Fcua-s1%2Fpython&rev=aea61b6eb97e2d8c0f6f71eb804e5769fe910af4" }, { name = "httpx", specifier = ">=0.28" }, { name = "huggingface-hub", marker = "extra == 'cua'", specifier = ">=0.24" },