diff --git a/README.md b/README.md index f02b175..f739ac8 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,9 @@ run_task("./gradlew build", queue_name="android", ...) run_task("npm run build", queue_name="web", ...) ``` -Both agents block until their respective builds complete. The server handles sequencing automatically. +`run_task` returns the final result for short commands. For longer commands, it returns a task handle +after at most 30 seconds while the command continues in the queue. The agent can show progress inline, +accept steering, and call `task_status` for the next bounded update without rerunning the command. **Hierarchical queues** - Use `/`-delimited queue names plus `--queue-capacity` when you need parallelism with a shared cap: @@ -74,7 +76,15 @@ each exact `queue_name` is still a FIFO queue with capacity 1. command: "./gradlew assembleDebug" working_directory: "/path/to/android-project" - ⎿ "SUCCESS exit=0 192.6s output=/tmp/agent-task-queue/output/task_1.log" + ⎿ "RUNNING task_id=1 queue=global elapsed=30.0s process_alive=true ..." + +⏺ agent-task-queue - task_status (MCP) + task_id: 1 + output_offset: 4821 + + ... additional bounded task_status calls while the build runs ... + + ⎿ "SUCCESS task_id=1 exit=0 192.6s output=/tmp/agent-task-queue/output/task_1.log" ⏺ Build completed successfully in 192.6s. ``` @@ -87,7 +97,14 @@ each exact `queue_name` is still a FIFO queue with capacity 1. command: "./gradlew assembleDebug" working_directory: "/path/to/android-project" - ⎿ "SUCCESS exit=0 32.6s output=/tmp/agent-task-queue/output/task_2.log" + ⎿ "QUEUED task_id=2 queue=global elapsed=30.0s position=1 ..." + +⏺ agent-task-queue - task_status (MCP) + task_id: 2 + + ... additional bounded task_status calls while queued and running ... + + ⎿ "SUCCESS task_id=2 exit=0 32.6s output=/tmp/agent-task-queue/output/task_2.log" ⏺ Build completed successfully in 32.6s. ``` @@ -114,7 +131,9 @@ With the queue: ## Key Features - **FIFO Queuing**: Strict first-in-first-out ordering within each exact `queue_name` -- **No Queue Timeouts**: MCP keeps connection alive while waiting in queue. The `timeout_seconds` parameter only applies to execution time—tasks can wait in queue indefinitely without timing out. (see [Why MCP?](#why-mcp-instead-of-a-cli-tool)) +- **Bounded Inline Progress**: Long calls yield a task handle within 30 seconds; `task_status` returns on output, a state change, completion, or another 30-second heartbeat +- **No Queue-Wait Timeouts**: Commands can remain queued indefinitely. `timeout_seconds` starts only when execution begins. +- **Explicit Cancellation**: Interrupting `run_task` or `task_status` leaves the command running; only `cancel_task` stops it - **Environment Variables**: Pass `env_vars="ANDROID_SERIAL=emulator-5560"` - **Multiple Queues**: Isolate different workloads with `queue_name` - **Zombie Protection**: Detects dead processes, kills orphans, clears stale locks @@ -122,7 +141,7 @@ With the queue: ## Desktop Sidecar -The repo also includes a minimal Compose Multiplatform desktop app in [desktop-sidecar](desktop-sidecar/README.md) for watching the queue in real time. +The repo also includes an optional Compose Multiplatform desktop app in [desktop-sidecar](desktop-sidecar/README.md) for watching the queue in real time. It is not required for inline agent progress. It reads the same local SQLite database as `tq` and the IntelliJ plugin, then shows: @@ -294,6 +313,7 @@ Agents use the `run_task` MCP tool for expensive operations: | `queue_name` | No | Queue identifier (default: "global") | | `timeout_seconds` | No | Max **execution** time before kill (default: 1200). Queue wait time doesn't count. | | `env_vars` | No | Environment variables: `"KEY=val,KEY2=val2"` | +| `wait_seconds` | No | Initial wait before yielding a background handle (0–30, default: 30) | `queue_name` may be hierarchical, such as `gradle/emu-5557`, when the server is configured with `--queue-capacity` scopes. @@ -301,6 +321,31 @@ Agents use the `run_task` MCP tool for expensive operations: Sibling queues that share a parent scope compete for that parent capacity on a best-effort basis; FIFO ordering is guaranteed within each exact queue, not across sibling queues. +### Long-Running Tasks and Inline Progress + +A command that does not finish during the initial bounded wait returns its state, task ID, queue +position or process liveness, elapsed time, recent output, and `next_output_offset`: + +```text +RUNNING task_id=42 queue=gradle/build elapsed=30.0s process_alive=true last_output=2.1s_ago + +--- NEW OUTPUT --- +> Task :app:compileDebugKotlin + +Task continues in the background. Do not rerun it. Call task_status(task_id=42, +output_offset=1234) for the next update, or cancel_task(task_id=42) to stop it. +``` + +Call `task_status` with the returned offset. Each status call returns as soon as output or state +changes, when the command finishes, or after at most 30 seconds as a liveness heartbeat. This makes +progress visible in the normal agent TUI or GUI transcript and gives the agent regular boundaries at +which it can absorb steering. MCP progress notifications are also emitted when supported, but the +bounded tool results do not depend on clients rendering those notifications. + +Cancelling or steering away from a `run_task`/`task_status` wait does **not** kill the command. Use +`cancel_task(task_id=42)` when termination is intended. Terminal results remain queryable after the +active queue row is released. + ### Example ``` @@ -425,7 +470,7 @@ Run `uvx agent-task-queue@latest --help` to see all options. ## IntelliJ Plugin -An optional [IntelliJ plugin](intellij-plugin/) provides real-time IDE integration — status bar widget, tool window with live streaming output, and balloon notifications for queue events. See the [plugin README](intellij-plugin/README.md) for details. +An optional [IntelliJ plugin](intellij-plugin/) provides an additional status bar widget, tool window, and notifications. It is separate from the inline `run_task`/`task_status` output available in any MCP agent client. See the [plugin README](intellij-plugin/README.md) for details. ## Architecture @@ -436,7 +481,7 @@ flowchart TD B -->|Execute| D[Subprocess
gradle, docker, etc.] D -.->|stdout/stderr| B - B -.->|blocks until complete| A + B -.->|terminal result or bounded status handle| A ``` ### Data Directory @@ -444,12 +489,13 @@ flowchart TD All data is stored in `/tmp/agent-task-queue/` by default: - `queue.db` - SQLite database for queue state - `agent-task-queue-logs.json` - JSON metrics log (NDJSON format) +- `output/task_.log` and `.raw.log` - Full and incrementally readable task output To use a different location, pass `--data-dir=/path/to/data` or set the `TASK_QUEUE_DATA_DIR` environment variable. ### Database Schema -The queue state is stored in SQLite at `/tmp/agent-task-queue/queue.db`: +Active queue state is stored in the `queue` table at `/tmp/agent-task-queue/queue.db`: | Column | Type | Description | |--------|------|-------------| @@ -463,6 +509,10 @@ The queue state is stored in SQLite at `/tmp/agent-task-queue/queue.db`: | `created_at` | TIMESTAMP | When task was queued | | `updated_at` | TIMESTAMP | Last status change | +The `task_results` table stores each task's terminal result as JSON after its active `queue` row is +released. This lets `task_status` report completion without holding a queue slot. Terminal results +are retained with the configured output-file limit. + ### Zombie Protection If an agent crashes while a task is running: @@ -505,17 +555,17 @@ To reduce token usage, full command output is written to files instead of return ``` Each task produces two output files: -- **`task_.log`** — Formatted log with headers (`COMMAND:`, `WORKING DIR:`), section markers (`--- STDOUT ---`, `--- STDERR ---`, `--- SUMMARY ---`), and exit code. Used by the IntelliJ plugin notifier and the "View Output" action. -- **`task_.raw.log`** — Raw stdout+stderr only, no metadata. Used by the IntelliJ plugin for clean streaming output in tabs. Added in MCP server v0.4.0. +- **`task_.log`** — Formatted log with headers (`COMMAND:`, `WORKING DIR:`), compatibility section markers (`--- STDOUT ---`, `--- STDERR ---`, `--- SUMMARY ---`), and exit code. Stdout and stderr content may be interleaved because both pipes are drained concurrently to prevent subprocess deadlocks. +- **`task_.raw.log`** — Raw interleaved stdout+stderr with no metadata. Used for incremental `task_status` output and optional viewer streaming. **On success**, the tool returns a single line: ``` -SUCCESS exit=0 31.2s command=./gradlew build output=/tmp/agent-task-queue/output/task_8.log +SUCCESS task_id=8 exit=0 31.2s command=./gradlew build output=/tmp/agent-task-queue/output/task_8.log ``` **On failure**, the last 50 lines of output are included: ``` -FAILED exit=1 12.5s command=./gradlew build output=/tmp/agent-task-queue/output/task_9.log +FAILED task_id=9 exit=1 12.5s command=./gradlew build output=/tmp/agent-task-queue/output/task_9.log [error output here] ``` @@ -626,15 +676,15 @@ flowchart LR subgraph mcp [MCP Approach] A2[Agent] --> |MCP Protocol| B2[Server] B2 --> C2[Queue] - B2 -.-> |"✓ blocks until complete"| A2 + B2 -.-> |"✓ bounded inline updates"| A2 end ``` **Why MCP solves this:** -- The MCP server keeps the connection alive indefinitely -- The agent's tool call blocks until the task completes -- No timeout configuration needed—it "just works" -- The server manages the queue; the agent just waits +- The MCP server owns the queued command independently from any one tool request +- Long calls return a handle within 30 seconds instead of depending on a long client timeout +- The agent receives incremental inline output and can accept steering between status calls +- An interrupted status request does not cancel the command; cancellation is explicit | Aspect | CLI Wrapper | Agent Task Queue | |--------|-------------|----------------| diff --git a/pyproject.toml b/pyproject.toml index 7884bd9..7d52e7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,9 @@ dev = [ requires = ["hatchling"] build-backend = "hatchling.build" +[tool.ruff.lint] +select = ["E4", "E7", "E9", "F"] + [tool.hatch.build.targets.wheel] packages = ["."] only-include = ["task_queue.py", "tq.py", "queue_core.py"] diff --git a/queue_core.py b/queue_core.py index 821b354..e8e6a13 100644 --- a/queue_core.py +++ b/queue_core.py @@ -125,6 +125,14 @@ def insert_waiting_task( ) """ +TASK_RESULTS_SCHEMA = """ +CREATE TABLE IF NOT EXISTS task_results ( + task_id INTEGER PRIMARY KEY, + result_json TEXT NOT NULL, + completed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP +) +""" + # Migration to add server_id column to existing databases QUEUE_MIGRATION_SERVER_ID = """ ALTER TABLE queue ADD COLUMN server_id TEXT @@ -183,6 +191,7 @@ def init_db(paths: QueuePaths): paths.data_dir.mkdir(parents=True, exist_ok=True) with get_db(paths.db_path) as conn: conn.execute(QUEUE_SCHEMA) + conn.execute(TASK_RESULTS_SCHEMA) conn.execute(QUEUE_INDEX) # Run migrations for existing databases for migration in [ diff --git a/task_queue.py b/task_queue.py index 4590efa..a0b744b 100644 --- a/task_queue.py +++ b/task_queue.py @@ -8,6 +8,8 @@ import argparse import asyncio +import codecs +import json import os import resource import signal @@ -17,7 +19,7 @@ import threading import uuid from collections import deque -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from fastmcp import FastMCP @@ -56,6 +58,15 @@ _active_task_ids: set[int] = set() _active_task_ids_lock = threading.Lock() +# Background command tasks are deliberately independent from individual MCP requests. Cancelling +# a run_task/task_status request only stops that wait; cancel_task is the explicit command cancel. +_background_tasks: dict[int, asyncio.Task] = {} + +MAX_TOOL_WAIT_SECONDS = 30 +STATUS_POLL_INTERVAL_SECONDS = 0.25 +PROGRESS_INTERVAL_SECONDS = 5 +MAX_INLINE_OUTPUT_BYTES = 16 * 1024 + # --- Argument Parsing --- def parse_args(): @@ -117,7 +128,7 @@ def _should_parse_module_args(argv0: str | None = None, module_name: str | None # Parse args at module load (before MCP server starts) _args = parse_args() if _should_parse_module_args() else argparse.Namespace( - data_dir="/tmp/agent-task-queue", + data_dir=os.environ.get("TASK_QUEUE_DATA_DIR", "/tmp/agent-task-queue"), max_log_size=5, max_output_files=50, tail_lines=50, @@ -164,7 +175,7 @@ def _current_context(): """Best-effort FastMCP request context; unavailable in tests and background codepaths.""" try: return get_context() - except LookupError: + except (LookupError, RuntimeError): return None @@ -245,21 +256,50 @@ def cleanup_queue(conn, queue_name: str, queue_capacities: dict[str, int] | None # --- Output File Management --- def cleanup_output_files(): - """Remove oldest output files if over limit. Covers both .log and .raw.log files.""" - if not OUTPUT_DIR.exists(): - return - - # Group files by task ID so both .log and .raw.log are cleaned together - files = sorted(OUTPUT_DIR.glob("task_*"), key=lambda f: f.stat().st_mtime) - # Each task produces up to 2 files (.log + .raw.log), so scale the limit - max_files = MAX_OUTPUT_FILES * 2 - if len(files) > max_files: - for old_file in files[: len(files) - max_files]: + """Remove the oldest task output groups and terminal results if over the limit.""" + old_task_ids = [] + retained_task_count = max(0, MAX_OUTPUT_FILES) + if OUTPUT_DIR.exists(): + formatted_logs = sorted( + ( + path + for path in OUTPUT_DIR.glob("task_*.log") + if not path.name.endswith(".raw.log") + ), + key=lambda path: path.stat().st_mtime, + ) + if len(formatted_logs) > retained_task_count: + old_logs = ( + formatted_logs + if retained_task_count == 0 + else formatted_logs[:-retained_task_count] + ) + old_task_ids = [ + path.name.removeprefix("task_").removesuffix(".log") + for path in old_logs + ] + + for task_id in old_task_ids: + for old_file in ( + OUTPUT_DIR / f"task_{task_id}.log", + OUTPUT_DIR / f"task_{task_id}.raw.log", + ): try: old_file.unlink() except OSError: pass + with get_db() as conn: + conn.execute( + """DELETE FROM task_results + WHERE task_id NOT IN ( + SELECT task_id FROM task_results + ORDER BY completed_at DESC, task_id DESC + LIMIT ? + )""", + (retained_task_count,), + ) + def clear_output_files() -> int: """Delete all output files. Returns number of files deleted.""" @@ -286,36 +326,39 @@ def get_memory_mb() -> float: # --- Core Queue Logic --- -async def wait_for_turn( +class TaskNoLongerActive(Exception): + """Raised when an explicitly cancelled task disappears while waiting for its turn.""" + + +async def register_task( queue_name: str, - command: str | None = None, - task_origin: TaskOrigin | None = None, + command: str, + task_origin: TaskOrigin, ) -> int: - """Register task, wait for turn, return task ID when acquired.""" - queue_name = normalize_queue_name(queue_name) - - # Ensure database exists and is valid + """Register a task and return its durable ID without waiting for execution.""" ensure_db() - - # Run cleanup BEFORE inserting - this clears orphaned tasks that would otherwise - # block the queue forever (since cleanup only runs during polling) with get_db() as conn: cleanup_queue(conn, queue_name, QUEUE_CAPACITIES) - - my_pid = os.getpid() - ctx = _current_context() - - with get_db() as conn: task_id = insert_waiting_task( conn, queue_name, - my_pid, + os.getpid(), SERVER_INSTANCE_ID, command=command, task_origin=task_origin, ) - # Track this task as active for orphan detection + # SQLite may reuse low IDs after the database is recreated. Never expose stale output under + # a newly registered handle (most notably when this project tests its own queue server). + for output_file in ( + OUTPUT_DIR / f"task_{task_id}.log", + OUTPUT_DIR / f"task_{task_id}.raw.log", + ): + try: + output_file.unlink() + except OSError: + pass + with _active_task_ids_lock: _active_task_ids.add(task_id) @@ -323,101 +366,601 @@ async def wait_for_turn( "task_queued", task_id=task_id, queue_name=queue_name, - pid=my_pid, + pid=os.getpid(), **task_origin_kwargs(task_origin), ) - queued_at = time.time() + ctx = _current_context() if ctx: - await ctx.info( - log_fmt(f"Request #{task_id} received. Entering '{queue_name}' queue.") + await ctx.info(log_fmt(f"Request #{task_id} received. Entering '{queue_name}' queue.")) + return task_id + + +async def wait_for_turn( + task_id: int, + queue_name: str, + task_origin: TaskOrigin, +) -> None: + """Wait until an already-registered task acquires its queue slot.""" + queued_at = time.time() + while True: + try: + with get_db() as conn: + cleanup_queue(conn, queue_name, QUEUE_CAPACITIES) + row = conn.execute( + "SELECT status FROM queue WHERE id = ?", + (task_id,), + ).fetchone() + if row is None: + raise TaskNoLongerActive(f"Task #{task_id} is no longer queued") + + started, _ = attempt_task_start( + conn, + task_id, + queue_name, + QUEUE_CAPACITIES, + os.getpid(), + ) + if started: + wait_time = time.time() - queued_at + log_metric( + "task_started", + task_id=task_id, + queue_name=queue_name, + pid=os.getpid(), + wait_time_seconds=round(wait_time, 2), + **task_origin_kwargs(task_origin), + ) + return + except sqlite3.OperationalError as exc: + if "database is locked" not in str(exc).lower(): + raise + + await asyncio.sleep(POLL_INTERVAL_WAITING) + + +async def release_lock(task_id: int): + """Release a queue slot and remove the task from active tracking.""" + with _active_task_ids_lock: + _active_task_ids.discard(task_id) + + try: + with get_db() as conn: + conn.execute("DELETE FROM queue WHERE id = ?", (task_id,)) + except sqlite3.OperationalError: + pass + + +def _load_task_result(task_id: int) -> dict | None: + ensure_db() + with get_db() as conn: + row = conn.execute( + "SELECT result_json FROM task_results WHERE task_id = ?", + (task_id,), + ).fetchone() + return json.loads(row["result_json"]) if row else None + + +def _store_task_result(result: dict) -> dict: + """Persist a terminal result once; explicit cancellation wins completion races.""" + ensure_db() + with get_db() as conn: + conn.execute( + "INSERT OR IGNORE INTO task_results (task_id, result_json) VALUES (?, ?)", + (result["task_id"], json.dumps(result)), + ) + row = conn.execute( + "SELECT result_json FROM task_results WHERE task_id = ?", + (result["task_id"],), + ).fetchone() + return json.loads(row["result_json"]) + + +def _read_new_output(task_id: int, output_offset: int) -> tuple[str, int, float | None]: + raw_output_file = OUTPUT_DIR / f"task_{task_id}.raw.log" + if not raw_output_file.exists(): + return "", 0, None + + stat = raw_output_file.stat() + size = stat.st_size + if output_offset > size: + output_offset = 0 + start = output_offset + omitted = 0 + if size - start > MAX_INLINE_OUTPUT_BYTES: + start = size - MAX_INLINE_OUTPUT_BYTES + omitted = start - output_offset + + with open(raw_output_file, "rb") as output: + output.seek(start) + text = output.read(size - start).decode(errors="replace").rstrip() + + if omitted: + text = f"[... {omitted} earlier output bytes omitted ...]\n{text}" + return text, size, round(max(0.0, time.time() - stat.st_mtime), 1) + + +def _active_task_snapshot(task_id: int, output_offset: int, include_output: bool = True) -> dict: + with get_db() as conn: + row = conn.execute("SELECT * FROM queue WHERE id = ?", (task_id,)).fetchone() + if row is None: + return { + "task_id": task_id, + "status": "unknown", + "message": "Task is not active and no terminal result is available", + } + position = None + if row["status"] == "waiting": + position = conn.execute( + """SELECT COUNT(*) + 1 AS position FROM queue + WHERE queue_name = ? AND status = 'waiting' AND id < ?""", + (row["queue_name"], task_id), + ).fetchone()["position"] + + created_at = datetime.fromisoformat(str(row["created_at"]).replace(" ", "T")) + utc_now = datetime.now(timezone.utc).replace(tzinfo=None) + elapsed_seconds = max(0.0, (utc_now - created_at).total_seconds()) + new_output, next_output_offset, output_age = ( + _read_new_output(task_id, output_offset) if include_output else ("", output_offset, None) + ) + external_status = "queued" if row["status"] == "waiting" else "running" + return { + "task_id": task_id, + "status": external_status, + "queue_name": row["queue_name"], + "queue_position": position, + "elapsed_seconds": round(elapsed_seconds, 1), + "command": row["command"], + "working_directory": row["working_directory"], + "process_alive": ( + is_process_alive(row["child_pid"]) + if row["status"] == "running" and row["child_pid"] + else None + ), + "output_file": str(OUTPUT_DIR / f"task_{task_id}.log"), + "new_output": new_output, + "next_output_offset": next_output_offset, + "last_output_seconds_ago": output_age, + } + + +def _task_snapshot(task_id: int, output_offset: int, include_output: bool = True) -> dict: + result = _load_task_result(task_id) + if result is None: + return _active_task_snapshot(task_id, output_offset, include_output=include_output) + + new_output, next_output_offset, output_age = ( + _read_new_output(task_id, output_offset) if include_output else ("", output_offset, None) + ) + return { + **result, + "new_output": new_output, + "next_output_offset": next_output_offset, + "last_output_seconds_ago": output_age, + } + + +def _format_task_snapshot(snapshot: dict) -> ToolResult: + status = snapshot["status"] + task_id = snapshot["task_id"] + output_file = snapshot.get("output_file") + + if status == "success": + text = ( + f"SUCCESS task_id={task_id} exit=0 {snapshot['duration_seconds']:.1f}s " + f"command={snapshot['command']} output={output_file}" + ) + elif status == "failed": + text = ( + f"FAILED task_id={task_id} exit={snapshot['exit_code']} " + f"{snapshot['duration_seconds']:.1f}s command={snapshot['command']} " + f"output={output_file}\n{snapshot['tail']}" + ) + elif status == "timeout": + text = ( + f"TIMEOUT task_id={task_id} killed after {snapshot['timeout_seconds']}s " + f"command={snapshot['command']} output={output_file}\n{snapshot['tail']}" + ) + elif status == "cancelled": + text = ( + f"CANCELLED task_id={task_id} after {snapshot['duration_seconds']:.1f}s " + f"command={snapshot['command']} output={output_file}" ) + elif status == "error": + text = f"ERROR task_id={task_id}: {snapshot['error']} output={output_file}" + elif status == "unknown": + text = f"ERROR task_id={task_id}: {snapshot['message']}" + else: + details = [ + f"{status.upper()} task_id={task_id}", + f"queue={snapshot['queue_name']}", + f"elapsed={snapshot['elapsed_seconds']:.1f}s", + ] + if snapshot.get("queue_position") is not None: + details.append(f"position={snapshot['queue_position']}") + if snapshot.get("process_alive") is not None: + details.append(f"process_alive={str(snapshot['process_alive']).lower()}") + if snapshot.get("last_output_seconds_ago") is not None: + details.append(f"last_output={snapshot['last_output_seconds_ago']:.1f}s_ago") + text = " ".join(details) + if snapshot.get("new_output"): + text += f"\n\n--- NEW OUTPUT ---\n{snapshot['new_output']}" + else: + text += "\n\n(no new output)" + text += ( + "\n\nTask continues in the background. Do not rerun it. " + f"Call task_status(task_id={task_id}, " + f"output_offset={snapshot['next_output_offset']}) for the next update, " + f"or cancel_task(task_id={task_id}) to stop it." + ) + + if status in {"success", "cancelled", "error"} and snapshot.get("new_output"): + text += f"\n\n--- NEW OUTPUT ---\n{snapshot['new_output']}" + + return ToolResult( + content=[TextContent(type="text", text=text)], + structured_content={"result": snapshot}, + ) + - last_pos = -1 - wait_ticks = 0 +async def _report_progress(snapshot: dict) -> None: + ctx = _current_context() + if not ctx or snapshot["status"] not in {"queued", "running"}: + return + message = f"Task #{snapshot['task_id']} {snapshot['status']} for {snapshot['elapsed_seconds']:.0f}s" + if snapshot.get("queue_position") is not None: + message += f" (queue position {snapshot['queue_position']})" try: - while True: + await ctx.report_progress(snapshot["elapsed_seconds"], message=message) + except Exception: + # Progress notifications are optional and unsupported clients must not break execution. + pass + + +async def _wait_for_initial_result(task_id: int, wait_seconds: int) -> dict: + deadline = time.monotonic() + wait_seconds + next_progress = time.monotonic() + PROGRESS_INTERVAL_SECONDS + while True: + snapshot = _task_snapshot(task_id, 0, include_output=False) + if snapshot["status"] not in {"queued", "running"}: + return snapshot + if time.monotonic() >= deadline: + return _task_snapshot(task_id, 0) + if time.monotonic() >= next_progress: + await _report_progress(snapshot) + next_progress = time.monotonic() + PROGRESS_INTERVAL_SECONDS + remaining = max(0.0, deadline - time.monotonic()) + await asyncio.sleep(min(STATUS_POLL_INTERVAL_SECONDS, remaining)) + + +async def _wait_for_status(task_id: int, wait_seconds: int, output_offset: int) -> dict: + snapshot = _task_snapshot(task_id, output_offset) + if snapshot["status"] not in {"queued", "running"} or snapshot.get("new_output"): + return snapshot + + initial_status = snapshot["status"] + deadline = time.monotonic() + wait_seconds + next_progress = time.monotonic() + PROGRESS_INTERVAL_SECONDS + while time.monotonic() < deadline: + remaining = max(0.0, deadline - time.monotonic()) + await asyncio.sleep(min(STATUS_POLL_INTERVAL_SECONDS, remaining)) + snapshot = _task_snapshot(task_id, output_offset) + if ( + snapshot["status"] not in {"queued", "running"} + or snapshot["status"] != initial_status + or snapshot.get("new_output") + ): + return snapshot + if time.monotonic() >= next_progress: + await _report_progress(snapshot) + next_progress = time.monotonic() + PROGRESS_INTERVAL_SECONDS + return snapshot + + +async def _terminate_process(proc: asyncio.subprocess.Process) -> None: + if proc.returncode is not None: + return + kill_process_tree(proc.pid) + try: + await asyncio.wait_for(asyncio.shield(proc.wait()), timeout=5.0) + except asyncio.TimeoutError: + try: + os.killpg(proc.pid, signal.SIGKILL) + except OSError: try: - with get_db() as conn: - cleanup_queue(conn, queue_name, QUEUE_CAPACITIES) - - started, pos = attempt_task_start( - conn, - task_id, - queue_name, - QUEUE_CAPACITIES, - my_pid, - ) + os.kill(proc.pid, signal.SIGKILL) + except OSError: + pass + await proc.wait() + else: + # The shell can exit before a signal-resistant descendant in its process group. + try: + os.killpg(proc.pid, 0) + os.killpg(proc.pid, signal.SIGKILL) + except OSError: + pass + - if started: - wait_time = time.time() - queued_at - log_metric( - "task_started", - task_id=task_id, - queue_name=queue_name, - pid=my_pid, - wait_time_seconds=round(wait_time, 2), - **task_origin_kwargs(task_origin), - ) - if ctx: - await ctx.info(log_fmt("Lock ACQUIRED. Starting execution.")) - return task_id - - wait_ticks += 1 - - if pos != last_pos: - if ctx: - await ctx.info(log_fmt(f"Position #{pos} in queue. Waiting...")) - last_pos = pos - elif wait_ticks % 10 == 0 and ctx: # Update every ~10 polls - await ctx.info( - log_fmt( - f"Still waiting... Position #{pos} ({int(wait_ticks * POLL_INTERVAL_WAITING)}s elapsed)" - ) - ) - except sqlite3.OperationalError as exc: - if "database is locked" not in str(exc).lower(): - raise - - await asyncio.sleep(POLL_INTERVAL_WAITING) +async def _terminate_external_process_group(pid: int) -> None: + """Terminate a task owned by another MCP server, escalating after a grace period.""" + kill_process_tree(pid) + await asyncio.sleep(1.0) + try: + os.killpg(pid, 0) + except OSError: + return + try: + os.killpg(pid, signal.SIGKILL) + except OSError: + try: + os.kill(pid, signal.SIGKILL) + except OSError: + pass + + +async def _execute_command( + task_id: int, + queue_name: str, + command: str, + working_directory: str, + timeout_seconds: int, + env: dict[str, str], + task_origin: TaskOrigin, +) -> dict: + mem_before = get_memory_mb() + start = time.time() + stdout_tail: deque = deque(maxlen=TAIL_LINES_ON_FAILURE) + stderr_tail: deque = deque(maxlen=TAIL_LINES_ON_FAILURE) + stdout_count = 0 + stderr_count = 0 + + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + output_file = OUTPUT_DIR / f"task_{task_id}.log" + raw_output_file = OUTPUT_DIR / f"task_{task_id}.raw.log" + + # nosec B602: shell execution is intentional. Users explicitly provide build commands and + # shell features such as pipes, redirects, and globs are part of the tool contract. + proc = await asyncio.create_subprocess_shell( # nosec B602 + command, + cwd=working_directory, + env=env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + start_new_session=True, + ) + with get_db() as conn: + conn.execute("UPDATE queue SET child_pid = ? WHERE id = ?", (proc.pid, task_id)) + + with open(output_file, "w") as formatted, open(raw_output_file, "w") as raw: + formatted.write(f"COMMAND: {command}\n") + formatted.write(f"WORKING DIR: {working_directory}\n") + formatted.write(f"STARTED: {datetime.now().isoformat()}\n") + # Keep the historical markers for older output viewers. Output is interleaved because + # both pipes must be drained concurrently to prevent a full stderr pipe from deadlocking. + formatted.write("\n--- STDOUT ---\n--- STDERR ---\n") + formatted.flush() + + async def stream_to_files(stream, tail_buffer: deque, label: str): + """Drain one pipe concurrently so neither subprocess pipe can block the command.""" + nonlocal stdout_count, stderr_count + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + while True: + chunk = await stream.read(8192) + if not chunk: + decoded = decoder.decode(b"", final=True) + if decoded: + formatted.write(decoded) + raw.write(decoded) + break + decoded = decoder.decode(chunk) + formatted.write(decoded) + formatted.flush() + raw.write(decoded) + raw.flush() + tail_buffer.extend(decoded.splitlines()) + line_count = max(1, decoded.count("\n")) + if label == "stdout": + stdout_count += line_count + else: + stderr_count += line_count + + stdout_task = asyncio.create_task(stream_to_files(proc.stdout, stdout_tail, "stdout")) + stderr_task = asyncio.create_task(stream_to_files(proc.stderr, stderr_tail, "stderr")) + process_task = asyncio.create_task(proc.wait()) + execution = asyncio.gather(stdout_task, stderr_task, process_task) + + try: + await asyncio.wait_for(asyncio.shield(execution), timeout=timeout_seconds) + except asyncio.TimeoutError: + await _terminate_process(proc) + await asyncio.gather(execution, return_exceptions=True) + duration = time.time() - start + formatted.write("\n--- SUMMARY ---\n") + formatted.write(f"EXIT CODE: TIMEOUT (killed after {timeout_seconds}s)\n") + formatted.write(f"DURATION: {duration:.1f}s\n") + tail = list(stderr_tail) if stderr_tail else list(stdout_tail) + tail_text = "\n".join(tail) if tail else "(no output)" + log_metric( + "task_timeout", + task_id=task_id, + queue_name=queue_name, + pid=os.getpid(), + command=command, + timeout_seconds=timeout_seconds, + memory_mb=round(get_memory_mb(), 1), + **task_origin_kwargs(task_origin), + ) + return { + "task_id": task_id, + "status": "timeout", + "exit_code": None, + "duration_seconds": round(duration, 1), + "timeout_seconds": timeout_seconds, + "command": command, + "output_file": str(output_file), + "tail": tail_text, + } + except asyncio.CancelledError: + await _terminate_process(proc) + await asyncio.gather(execution, return_exceptions=True) + duration = time.time() - start + formatted.write("\n--- SUMMARY ---\n") + formatted.write("EXIT CODE: CANCELLED\n") + formatted.write(f"DURATION: {duration:.1f}s\n") + raise + except Exception: + await _terminate_process(proc) + await asyncio.gather(execution, return_exceptions=True) + raise + + duration = time.time() - start + formatted.write("\n--- SUMMARY ---\n") + formatted.write(f"EXIT CODE: {proc.returncode}\n") + formatted.write(f"DURATION: {duration:.1f}s\n") + + mem_after = get_memory_mb() + log_metric( + "task_completed", + task_id=task_id, + queue_name=queue_name, + pid=os.getpid(), + command=command, + exit_code=proc.returncode, + duration_seconds=round(duration, 2), + stdout_lines=stdout_count, + stderr_lines=stderr_count, + memory_before_mb=round(mem_before, 1), + memory_after_mb=round(mem_after, 1), + **task_origin_kwargs(task_origin), + ) + tail = list(stderr_tail) if stderr_tail else list(stdout_tail) + tail_text = "\n".join(tail) if tail else "(no output)" + return { + "task_id": task_id, + "status": "success" if proc.returncode == 0 else "failed", + "exit_code": proc.returncode, + "duration_seconds": round(duration, 1), + "command": command, + "output_file": str(output_file), + "tail": None if proc.returncode == 0 else tail_text, + } + + +async def _run_registered_task( + task_id: int, + queue_name: str, + command: str, + working_directory: str, + timeout_seconds: int, + env: dict[str, str], + task_origin: TaskOrigin, +) -> None: + try: + await wait_for_turn(task_id, queue_name, task_origin) + result = await _execute_command( + task_id, + queue_name, + command, + working_directory, + timeout_seconds, + env, + task_origin, + ) + _store_task_result(result) + except TaskNoLongerActive: + if _load_task_result(task_id) is None: + _store_task_result({ + "task_id": task_id, + "status": "cancelled", + "duration_seconds": 0.0, + "command": command, + "output_file": str(OUTPUT_DIR / f"task_{task_id}.log"), + }) except asyncio.CancelledError: - # Client disconnected (e.g., sub-agent cancelled) - clean up our queue entry - with _active_task_ids_lock: - _active_task_ids.discard(task_id) + if _load_task_result(task_id) is None: + log_metric( + "task_cancelled", + task_id=task_id, + queue_name=queue_name, + pid=os.getpid(), + command=command, + reason="background_task_cancelled", + **task_origin_kwargs(task_origin), + ) + _store_task_result({ + "task_id": task_id, + "status": "cancelled", + "duration_seconds": 0.0, + "command": command, + "output_file": str(OUTPUT_DIR / f"task_{task_id}.log"), + }) + raise + except Exception as exc: log_metric( - "task_cancelled", + "task_error", task_id=task_id, queue_name=queue_name, - pid=my_pid, - reason="client_disconnected", + pid=os.getpid(), + command=command, + error=str(exc), **task_origin_kwargs(task_origin), ) - with get_db() as conn: - conn.execute("DELETE FROM queue WHERE id = ?", (task_id,)) - raise # Re-raise to propagate cancellation - + _store_task_result({ + "task_id": task_id, + "status": "error", + "duration_seconds": 0.0, + "command": command, + "output_file": str(OUTPUT_DIR / f"task_{task_id}.log"), + "error": str(exc), + }) + finally: + await release_lock(task_id) + cleanup_output_files() -async def release_lock(task_id: int): - """Release a queue lock.""" - # Remove from active tracking - with _active_task_ids_lock: - _active_task_ids.discard(task_id) - - ctx = _current_context() +def _forget_background_task(task_id: int, completed: asyncio.Task) -> None: + _background_tasks.pop(task_id, None) try: - with get_db() as conn: - conn.execute("DELETE FROM queue WHERE id = ?", (task_id,)) - except sqlite3.OperationalError: - # Database was deleted (e.g., by tests) - nothing to release + completed.exception() + except asyncio.CancelledError: pass - if ctx: - await ctx.info(log_fmt("Task complete. Queue slot released.")) + +def _start_background_task( + task_id: int, + queue_name: str, + command: str, + working_directory: str, + timeout_seconds: int, + env: dict[str, str], + task_origin: TaskOrigin, +) -> None: + background_task = asyncio.create_task( + _run_registered_task( + task_id, + queue_name, + command, + working_directory, + timeout_seconds, + env, + task_origin, + ), + name=f"queued-task-{task_id}", + ) + _background_tasks[task_id] = background_task + background_task.add_done_callback( + lambda completed: _forget_background_task(task_id, completed) + ) -# --- The Tool --- +def _validate_wait_seconds(wait_seconds: int) -> str | None: + if not 0 <= wait_seconds <= MAX_TOOL_WAIT_SECONDS: + return f"wait_seconds must be between 0 and {MAX_TOOL_WAIT_SECONDS}" + return None + + +# --- Tools --- @mcp.tool( title="Run Queued Task", annotations={ @@ -433,13 +976,19 @@ async def run_task( timeout_seconds: int = 1200, env_vars: str = "", agent_name: str = "", + wait_seconds: int = MAX_TOOL_WAIT_SECONDS, ): """ Execute a command through the task queue for sequential processing. IMPORTANT: Before calling this tool, tell the user the exact command you are - about to run (e.g., "Running `./gradlew :app:compileDebugKotlin`"). - This provides visibility since the tool execution may take a while. + about to run. The tool returns the final result for short commands. After at + most 30 seconds it returns a queued/running task handle with inline output; + the command continues in the background. Use task_status with the returned + task_id and next_output_offset for further updates. Do not rerun the command. + + Cancelling this tool's wait does not cancel the command. Use cancel_task for + explicit command cancellation. When a command fails, analyze the output tail to identify the root cause and show the user the specific error with the responsible file/line if available. @@ -448,60 +997,51 @@ async def run_task( command involves ANY of the following: BUILD TOOLS (always use this tool): - - gradle, gradlew, ./gradlew (any Gradle command) - - bazel, bazelisk (any Bazel command) - - make, cmake, ninja - - mvn, maven - - cargo build, cargo test - - go build, go test + - gradle, gradlew, ./gradlew, bazel, bazelisk, make, cmake, ninja + - mvn, maven, cargo build, cargo test, go build, go test - npm run build, npm test, yarn build, pnpm build - dotnet build, dotnet test, msbuild CONTAINER/VM OPERATIONS (always use this tool): - docker build, docker-compose up, docker compose - - podman build, podman-compose - - kubectl apply, helm install + - podman build, podman-compose, kubectl apply, helm install PACKAGE OPERATIONS (always use this tool): - - pip install (with compilation) - - npm install, yarn install, pnpm install - - bundle install - - composer install + - pip install, npm install, yarn install, pnpm install + - bundle install, composer install TEST SUITES (always use this tool): - - pytest, jest, mocha, rspec - - Any command running a full test suite + - pytest, jest, mocha, rspec, or any full test suite WHY: Running multiple builds simultaneously causes system freeze and race - conditions. This tool ensures only one heavy task runs at a time using a - FIFO queue. + conditions. This tool ensures only the configured number of heavy tasks run. Args: command: The full shell command to run. working_directory: ABSOLUTE path to the execution root. - queue_name: Queue identifier for grouping tasks (default: "global"). - Queue names may be hierarchical (for example `gradle/emu-5557`) when the server - is configured with `--queue-capacity` scopes. - timeout_seconds: Max **execution** time before killing the task (default: 1200 = 20 mins). - Queue wait time does NOT count against this timeout. - env_vars: Environment variables to set, format: "KEY1=value1,KEY2=value2" - agent_name: Optional friendly caller label (for example `amp` or `claude-code`). + queue_name: Queue identifier (default: "global"). + timeout_seconds: Max execution time. Queue wait time does not count. + env_vars: Environment variables as "KEY1=value1,KEY2=value2". + agent_name: Optional caller label such as "amp". + wait_seconds: Initial bounded wait, from 0 through 30 seconds (default: 30). Returns: - Command output including stdout, stderr, and exit code. + A terminal result, or a queued/running handle for task_status. """ if not command or not command.strip(): return "ERROR: Command cannot be empty" - if not os.path.exists(working_directory): return f"ERROR: Working directory does not exist: {working_directory}" + if timeout_seconds < 1: + return "ERROR: timeout_seconds must be at least 1" + if wait_error := _validate_wait_seconds(wait_seconds): + return f"ERROR: {wait_error}" try: queue_name = normalize_queue_name(queue_name) except ValueError as exc: return f"ERROR: {str(exc)}" - # Parse environment variables env = os.environ.copy() if env_vars: for pair in env_vars.split(","): @@ -511,213 +1051,109 @@ async def run_task( caller_name = agent_name.strip() or _current_client_id() task_origin = collect_task_origin(working_directory, caller_name) + task_id = await register_task(queue_name, command, task_origin) + _start_background_task( + task_id, + queue_name, + command, + working_directory, + timeout_seconds, + env, + task_origin, + ) + snapshot = await _wait_for_initial_result(task_id, wait_seconds) + return _format_task_snapshot(snapshot) - task_id = await wait_for_turn(queue_name, command, task_origin=task_origin) - mem_before = get_memory_mb() - - start = time.time() - # Use bounded deques - only keep last N lines in memory for error messages - stdout_tail: deque = deque(maxlen=TAIL_LINES_ON_FAILURE) - stderr_tail: deque = deque(maxlen=TAIL_LINES_ON_FAILURE) - stdout_count = 0 - stderr_count = 0 - - # Two output files are written per task: - # task_.log — formatted log with metadata headers, section markers (--- STDOUT ---, - # --- STDERR ---, --- SUMMARY ---), and exit code. Written by all MCP - # server versions. Used by the IntelliJ plugin notifier to read exit - # codes, and by "View Output" to open full logs. - # task_.raw.log — raw stdout+stderr only, no markers or metadata. Added in MCP server - # v0.4.0 (not present in v0.3.x and earlier). Used by the IntelliJ - # plugin OutputStreamer for clean tailing in output tabs. - OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - output_file = OUTPUT_DIR / f"task_{task_id}.log" - - try: - # nosec B602: shell execution is intentional - this MCP tool executes user-provided - # build commands (gradle, docker, pytest, etc.). Shell features (pipes, redirects, - # globs) are required. Input comes from AI agents which users explicitly invoke. - proc = await asyncio.create_subprocess_shell( # nosec B602 - command, - cwd=working_directory, - env=env, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=True, # Run in own process group for clean kill - ) - - # Record child PID for zombie protection - with get_db() as conn: - conn.execute( - "UPDATE queue SET child_pid = ? WHERE id = ?", (proc.pid, task_id) - ) - - # Open files for streaming output - formatted log + raw log for plugin tailing - raw_output_file = OUTPUT_DIR / f"task_{task_id}.raw.log" - with open(output_file, "w") as f, open(raw_output_file, "w") as raw_f: - # Header to formatted log only - f.write(f"COMMAND: {command}\n") - f.write(f"WORKING DIR: {working_directory}\n") - f.write(f"STARTED: {datetime.now().isoformat()}\n") - f.write("\n--- STDOUT ---\n") - - async def stream_to_file(stream, tail_buffer: deque, label: str): - """Stream output directly to both files, keeping only tail in memory.""" - nonlocal stdout_count, stderr_count - while True: - line = await stream.readline() - if not line: - break - decoded = line.decode().rstrip() - f.write(decoded + "\n") - f.flush() - raw_f.write(decoded + "\n") - raw_f.flush() - tail_buffer.append(decoded) - if label == "stdout": - stdout_count += 1 - else: - stderr_count += 1 - try: - # Stream stdout first, then stderr (written sequentially to file) - await asyncio.wait_for( - stream_to_file(proc.stdout, stdout_tail, "stdout"), - timeout=timeout_seconds, - ) - f.write("\n--- STDERR ---\n") - await asyncio.wait_for( - stream_to_file(proc.stderr, stderr_tail, "stderr"), - timeout=timeout_seconds, - ) - await proc.wait() - duration = time.time() - start - - # Append summary to formatted log only - f.write("\n--- SUMMARY ---\n") - f.write(f"EXIT CODE: {proc.returncode}\n") - f.write(f"DURATION: {duration:.1f}s\n") - - except asyncio.TimeoutError: - try: - os.killpg(proc.pid, signal.SIGKILL) - await proc.wait() - except Exception: - pass - f.write("\n--- SUMMARY ---\n") - f.write(f"EXIT CODE: TIMEOUT (killed after {timeout_seconds}s)\n") +@mcp.tool( + title="Check Queued Task", + annotations={ + "readOnlyHint": True, + "openWorldHint": False, + "idempotentHint": True, + }, +) +async def task_status( + task_id: int, + output_offset: int = 0, + wait_seconds: int = MAX_TOOL_WAIT_SECONDS, +): + """ + Wait for and return the next inline update from a background queued task. - log_metric( - "task_timeout", - task_id=task_id, - queue_name=queue_name, - pid=os.getpid(), - command=command, - timeout_seconds=timeout_seconds, - memory_mb=round(get_memory_mb(), 1), - **task_origin_kwargs(task_origin), - ) - cleanup_output_files() - - tail = list(stderr_tail) if stderr_tail else list(stdout_tail) - tail_text = "\n".join(tail) if tail else "(no output)" - text = f"TIMEOUT killed after {timeout_seconds}s command={command} output={output_file}\n{tail_text}" - return ToolResult( - content=[TextContent(type="text", text=text)], - structured_content={"result": { - "status": "timeout", - "exit_code": None, - "duration_seconds": timeout_seconds, - "command": command, - "output_file": str(output_file), - "tail": tail_text, - }}, - ) + Returns when new output appears, task state changes, the task completes, or + after at most 30 seconds as a liveness heartbeat. Pass next_output_offset from + the previous result to avoid repeating output. Cancelling this status wait + does not cancel the command; use cancel_task to stop it explicitly. + """ + if task_id < 1: + return "ERROR: task_id must be a positive integer" + if output_offset < 0: + return "ERROR: output_offset cannot be negative" + if wait_error := _validate_wait_seconds(wait_seconds): + return f"ERROR: {wait_error}" - # File is now closed, log metrics - mem_after = get_memory_mb() - log_metric( - "task_completed", - task_id=task_id, - queue_name=queue_name, - pid=os.getpid(), - command=command, - exit_code=proc.returncode, - duration_seconds=round(duration, 2), - stdout_lines=stdout_count, - stderr_lines=stderr_count, - memory_before_mb=round(mem_before, 1), - memory_after_mb=round(mem_after, 1), - **task_origin_kwargs(task_origin), - ) - cleanup_output_files() + snapshot = await _wait_for_status(task_id, wait_seconds, output_offset) + return _format_task_snapshot(snapshot) - # Return concise summary for agents - if proc.returncode == 0: - text = f"SUCCESS exit=0 {duration:.1f}s command={command} output={output_file}" - return ToolResult( - content=[TextContent(type="text", text=text)], - structured_content={"result": { - "status": "success", - "exit_code": 0, - "duration_seconds": round(duration, 1), - "command": command, - "output_file": str(output_file), - "tail": None, - }}, - ) - else: - # On failure, include tail of output for context - tail = list(stderr_tail) if stderr_tail else list(stdout_tail) - tail_text = "\n".join(tail) if tail else "(no output)" - text = f"FAILED exit={proc.returncode} {duration:.1f}s command={command} output={output_file}\n{tail_text}" - return ToolResult( - content=[TextContent(type="text", text=text)], - structured_content={"result": { - "status": "failed", - "exit_code": proc.returncode, - "duration_seconds": round(duration, 1), - "command": command, - "output_file": str(output_file), - "tail": tail_text, - }}, - ) - except asyncio.CancelledError: - # Client disconnected while task was running - kill the subprocess - log_metric( - "task_cancelled", - task_id=task_id, - queue_name=queue_name, - pid=os.getpid(), - command=command, - reason="client_disconnected_during_execution", - **task_origin_kwargs(task_origin), - ) - try: - os.killpg(proc.pid, signal.SIGTERM) - await asyncio.wait_for(proc.wait(), timeout=5.0) - except Exception: - try: - os.killpg(proc.pid, signal.SIGKILL) - except Exception: - pass - raise # Re-raise to propagate cancellation +@mcp.tool( + title="Cancel Queued Task", + annotations={ + "destructiveHint": True, + "openWorldHint": False, + "idempotentHint": True, + }, +) +async def cancel_task(task_id: int): + """Explicitly cancel one queued or running background task by task ID.""" + if task_id < 1: + return "ERROR: task_id must be a positive integer" + + existing_result = _load_task_result(task_id) + if existing_result is not None: + return _format_task_snapshot(_task_snapshot(task_id, 0)) + + snapshot = _active_task_snapshot(task_id, 0) + if snapshot["status"] == "unknown": + return _format_task_snapshot(snapshot) + + result = _store_task_result({ + "task_id": task_id, + "status": "cancelled", + "duration_seconds": snapshot["elapsed_seconds"], + "command": snapshot["command"], + "output_file": snapshot["output_file"], + }) + log_metric( + "task_cancelled", + task_id=task_id, + queue_name=snapshot["queue_name"], + pid=os.getpid(), + command=snapshot["command"], + reason="explicit_cancel", + ) - except Exception as e: - log_metric( - "task_error", - task_id=task_id, - queue_name=queue_name, - pid=os.getpid(), - command=command, - error=str(e), - **task_origin_kwargs(task_origin), - ) - return f"ERROR: {str(e)}" + background_task = _background_tasks.get(task_id) + if background_task is not None: + background_task.cancel() + await asyncio.gather(background_task, return_exceptions=True) + else: + with get_db() as conn: + row = conn.execute( + "SELECT child_pid FROM queue WHERE id = ?", + (task_id,), + ).fetchone() + if row and row["child_pid"]: + await _terminate_external_process_group(row["child_pid"]) + conn.execute("DELETE FROM queue WHERE id = ?", (task_id,)) - finally: - await release_lock(task_id) + return _format_task_snapshot({ + **result, + "new_output": "", + "next_output_offset": snapshot["next_output_offset"], + "last_output_seconds_ago": snapshot["last_output_seconds_ago"], + }) @mcp.tool() diff --git a/tests/test_queue.py b/tests/test_queue.py index 9ca0a83..c46d4aa 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -7,7 +7,9 @@ import asyncio import json import os +import shlex import subprocess +import sys import time from pathlib import Path from types import SimpleNamespace @@ -41,7 +43,7 @@ @pytest.fixture(autouse=True) -def clean_db(): +async def clean_db(): """Clean database before each test.""" if DB_PATH.exists(): DB_PATH.unlink() @@ -57,6 +59,14 @@ def clean_db(): init_db() yield # Cleanup after test + background_tasks = list(task_queue._background_tasks.values()) + for background_task in background_tasks: + background_task.cancel() + if background_tasks: + await asyncio.gather(*background_tasks, return_exceptions=True) + task_queue._background_tasks.clear() + with task_queue._active_task_ids_lock: + task_queue._active_task_ids.clear() if DB_PATH.exists(): DB_PATH.unlink() @@ -436,7 +446,7 @@ async def run_task_b(): @pytest.mark.asyncio -async def test_parent_capacity_preserves_fifo_within_child_queue(monkeypatch): +async def test_parent_capacity_preserves_fifo_within_child_queue(monkeypatch, tmp_path): """A tighter parent scope should not let a younger child task jump the queue.""" monkeypatch.setattr( task_queue, @@ -445,7 +455,7 @@ async def test_parent_capacity_preserves_fifo_within_child_queue(monkeypatch): ) results = {} - end_times = {} + execution_order = tmp_path / "execution-order.txt" async def run_parent_blocker(): client = Client(mcp) @@ -467,12 +477,14 @@ async def run_older_child(): result = await client.call_tool( "run_task", { - "command": "sleep 1 && echo 'older child done'", + "command": ( + "sleep 1 && echo older >> " + f"{shlex.quote(str(execution_order))} && echo 'older child done'" + ), "working_directory": "/tmp", "queue_name": "gradle/emu-5557", }, ) - end_times["older"] = time.time() results["older"] = str(result) async def run_younger_child(): @@ -482,12 +494,14 @@ async def run_younger_child(): result = await client.call_tool( "run_task", { - "command": "echo 'younger child done'", + "command": ( + f"echo younger >> {shlex.quote(str(execution_order))} " + "&& echo 'younger child done'" + ), "working_directory": "/tmp", "queue_name": "gradle/emu-5557", }, ) - end_times["younger"] = time.time() results["younger"] = str(result) await asyncio.gather(run_parent_blocker(), run_older_child(), run_younger_child()) @@ -496,30 +510,57 @@ async def run_younger_child(): assert "SUCCESS" in results["younger"] assert "older child done" in read_output_file(results["older"]) assert "younger child done" in read_output_file(results["younger"]) - assert end_times["older"] <= end_times["younger"] + assert execution_order.read_text().splitlines() == ["older", "younger"] @pytest.mark.asyncio -async def test_parent_capacity_allows_parallel_child_queues(monkeypatch): +async def test_parent_capacity_allows_parallel_child_queues(monkeypatch, tmp_path): """A parent scope with capacity 2 should allow two child queues to run together.""" monkeypatch.setattr(task_queue, "QUEUE_CAPACITIES", parse_queue_capacities(["gradle=2"])) results = {} - end_times = {} - overall_start = time.time() + markers = { + "A": tmp_path / "child-a-started", + "B": tmp_path / "child-b-started", + } + barrier_script = tmp_path / "wait-for-sibling.py" + barrier_script.write_text( + "import sys\n" + "import time\n" + "from pathlib import Path\n" + "mine = Path(sys.argv[1])\n" + "other = Path(sys.argv[2])\n" + "mine.touch()\n" + "deadline = time.monotonic() + 5\n" + "while not other.exists():\n" + " if time.monotonic() >= deadline:\n" + " raise TimeoutError('sibling child queue did not start')\n" + " time.sleep(0.01)\n" + "print(f'{sys.argv[3]} done')\n" + ) async def run_child(queue_name: str, result_key: str): + other_key = "B" if result_key == "A" else "A" + command = " ".join( + shlex.quote(part) + for part in ( + sys.executable, + str(barrier_script), + str(markers[result_key]), + str(markers[other_key]), + queue_name, + ) + ) client = Client(mcp) async with client: result = await client.call_tool( "run_task", { - "command": f"sleep 2 && echo '{queue_name} done'", + "command": command, "working_directory": "/tmp", "queue_name": queue_name, }, ) - end_times[result_key] = time.time() results[result_key] = str(result) await asyncio.gather( @@ -527,20 +568,227 @@ async def run_child(queue_name: str, result_key: str): run_child("gradle/emu-5559", "B"), ) - total_elapsed = time.time() - overall_start assert "SUCCESS" in results["A"] assert "SUCCESS" in results["B"] - assert total_elapsed < 3.5 - assert abs(end_times["A"] - end_times["B"]) < 1.0 + assert markers["A"].exists() + assert markers["B"].exists() @pytest.mark.asyncio async def test_tool_available(client): - """Test that the run_task tool is available.""" + """Test that background task lifecycle tools are available.""" async with client: tools = await client.list_tools() tool_names = [t.name for t in tools] - assert "run_task" in tool_names + assert {"run_task", "task_status", "cancel_task"} <= set(tool_names) + + +@pytest.mark.asyncio +async def test_long_task_returns_handle_and_streams_inline_output(client): + """A long task yields control, then exposes incremental output and its final result.""" + async with client: + initial = await client.call_tool( + "run_task", + { + "command": "echo started; sleep 1; echo finished", + "working_directory": "/tmp", + "queue_name": "background_status_test", + "wait_seconds": 0, + }, + ) + initial_result = initial.structured_content["result"] + task_id = initial_result["task_id"] + offset = initial_result["next_output_offset"] + + assert initial_result["status"] in {"queued", "running"} + assert "Task continues in the background" in str(initial) + + collected_output = "" + final_result = initial_result + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + update = await client.call_tool( + "task_status", + { + "task_id": task_id, + "output_offset": offset, + "wait_seconds": 2, + }, + ) + final_result = update.structured_content["result"] + collected_output += final_result.get("new_output", "") + offset = final_result["next_output_offset"] + if final_result["status"] not in {"queued", "running"}: + break + + assert final_result["status"] == "success" + assert "started" in collected_output + assert "finished" in collected_output + with get_db() as conn: + assert conn.execute( + "SELECT COUNT(*) AS c FROM queue WHERE id = ?", + (task_id,), + ).fetchone()["c"] == 0 + assert conn.execute( + "SELECT COUNT(*) AS c FROM task_results WHERE task_id = ?", + (task_id,), + ).fetchone()["c"] == 1 + + +@pytest.mark.asyncio +async def test_cancelling_run_task_wait_does_not_cancel_command(client): + """Steering away from run_task only interrupts its bounded wait.""" + async with client: + request = asyncio.create_task( + client.call_tool( + "run_task", + { + "command": "sleep 1; echo survived", + "working_directory": "/tmp", + "queue_name": "cancelled_run_wait_test", + "wait_seconds": 30, + }, + ) + ) + + task_id = None + deadline = time.monotonic() + 3 + while time.monotonic() < deadline and task_id is None: + with get_db() as conn: + row = conn.execute( + "SELECT id FROM queue WHERE queue_name = ?", + ("cancelled_run_wait_test",), + ).fetchone() + task_id = row["id"] if row else None + await asyncio.sleep(0.05) + + assert task_id is not None + request.cancel() + with pytest.raises(asyncio.CancelledError): + await request + + final = await client.call_tool( + "task_status", + {"task_id": task_id, "wait_seconds": 3}, + ) + assert final.structured_content["result"]["status"] == "success" + assert "survived" in final.structured_content["result"]["new_output"] + + +@pytest.mark.asyncio +async def test_cancelling_status_wait_does_not_cancel_command(client): + """Steering away from a status heartbeat leaves the background command running.""" + async with client: + initial = await client.call_tool( + "run_task", + { + "command": "sleep 1; echo survived status cancellation", + "working_directory": "/tmp", + "queue_name": "cancelled_status_wait_test", + "wait_seconds": 0, + }, + ) + initial_result = initial.structured_content["result"] + task_id = initial_result["task_id"] + + status_wait = asyncio.create_task( + client.call_tool( + "task_status", + { + "task_id": task_id, + "output_offset": initial_result["next_output_offset"], + "wait_seconds": 30, + }, + ) + ) + await asyncio.sleep(0.2) + status_wait.cancel() + with pytest.raises(asyncio.CancelledError): + await status_wait + + final = await client.call_tool( + "task_status", + {"task_id": task_id, "wait_seconds": 3}, + ) + assert final.structured_content["result"]["status"] == "success" + assert "survived status cancellation" in final.structured_content["result"]["new_output"] + + +@pytest.mark.asyncio +async def test_cancel_task_explicitly_terminates_command(client, tmp_path): + """Only cancel_task stops the subprocess and records a cancelled result.""" + marker = tmp_path / "should-not-exist" + async with client: + initial = await client.call_tool( + "run_task", + { + "command": f"sleep 5; touch {shlex.quote(str(marker))}", + "working_directory": "/tmp", + "queue_name": "explicit_cancel_test", + "wait_seconds": 0, + }, + ) + task_id = initial.structured_content["result"]["task_id"] + + child_pid = None + deadline = time.monotonic() + 3 + while time.monotonic() < deadline and child_pid is None: + with get_db() as conn: + row = conn.execute( + "SELECT child_pid FROM queue WHERE id = ?", + (task_id,), + ).fetchone() + child_pid = row["child_pid"] if row else None + await asyncio.sleep(0.05) + + assert child_pid is not None + cancelled = await client.call_tool("cancel_task", {"task_id": task_id}) + + assert cancelled.structured_content["result"]["status"] == "cancelled" + assert not marker.exists() + assert not queue_core.is_process_alive(child_pid) + with get_db() as conn: + assert conn.execute( + "SELECT COUNT(*) AS c FROM queue WHERE id = ?", + (task_id,), + ).fetchone()["c"] == 0 + + +@pytest.mark.asyncio +async def test_large_stderr_does_not_deadlock(client): + """stdout and stderr are drained concurrently so a full stderr pipe cannot block.""" + script = "import sys; sys.stderr.write('x' * 2_000_000); sys.stderr.flush(); print('done')" + command = f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}" + + async with client: + result = await client.call_tool( + "run_task", + { + "command": command, + "working_directory": "/tmp", + "queue_name": "large_stderr_test", + "timeout_seconds": 5, + }, + ) + + assert result.structured_content["result"]["status"] == "success" + assert "done" in read_output_file(str(result)) + + +@pytest.mark.asyncio +async def test_background_tool_argument_validation(client): + async with client: + invalid_wait = await client.call_tool( + "task_status", + {"task_id": 1, "wait_seconds": 31}, + ) + invalid_offset = await client.call_tool( + "task_status", + {"task_id": 1, "output_offset": -1}, + ) + + assert "wait_seconds must be between 0 and 30" in str(invalid_wait) + assert "output_offset cannot be negative" in str(invalid_offset) @pytest.mark.asyncio @@ -1263,11 +1511,13 @@ def test_stale_server_instance_cleanup(): # --- Configuration Tests --- -def test_parse_args_defaults(): +def test_parse_args_defaults(monkeypatch): """Test that parse_args returns correct defaults.""" import sys from task_queue import parse_args + monkeypatch.delenv("TASK_QUEUE_DATA_DIR", raising=False) + # Save original argv and replace with empty args original_argv = sys.argv sys.argv = ["task_queue.py"]