Coverage for gco_mcp/tools/_long_task.py: 96.53%
244 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Shared async subprocess runner for long-running MCP tools.
3Streams progress through FastMCP's Progress dependency, emits a periodic
4heartbeat when the underlying process goes quiet, captures a bounded tail of
5stderr for failure surfacing, and raises ``ToolError`` on non-zero exit.
7In parallel with the MCP wire, every invocation writes a JSON status file plus
8a size-bounded raw log under ``~/.gco/tasks/`` via
9``_task_status.TaskStatusWriter``. This gives operators an out-of-band view of
10work even when the MCP client drops streamed notifications.
11"""
13from __future__ import annotations
15import asyncio
16import contextlib
17import json
18import re
19import time
20from collections import deque
21from collections.abc import AsyncIterator, Sequence
22from typing import Any
24import cli_runner
25from audit import _try_get_task_id
26from fastmcp.exceptions import ToolError
28from tools._task_status import TaskStatusWriter, is_valid_task_id, make_task_id
30# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
31# Generated at (UTC): 2026-07-18T01:03:40Z
32# Flowchart(s) generated from this file:
33# * ``_run_long_task`` -> ``diagrams/code_diagrams/gco_mcp/tools/_long_task._run_long_task.html``
34# (PNG: ``diagrams/code_diagrams/gco_mcp/tools/_long_task._run_long_task.png``)
35# Regenerate with ``python diagrams/code_diagrams/generate.py``.
36# <pyflowchart-code-diagram> END
39_CFN_FAILED_RE = re.compile(r"(CREATE|UPDATE|DELETE)_FAILED")
40_CDK_STACK_LINE_RE = re.compile(r"\b(gco-[a-z0-9-]+)\b")
41_CDK_STACK_DONE_RE = re.compile(r"[✅✨]\s+(gco-[a-z0-9-]+)\b")
42_CANCEL_GRACE_SECONDS = 10
43_HEARTBEAT_INTERVAL_SECONDS = 30
44_STDERR_TAIL_LINES = 80
45_FAILED_EVENT_LINES = 10
46_STREAM_READ_BYTES = 16 * 1024
47_STREAM_LINE_MAX_BYTES = 64 * 1024
48_DIAGNOSTIC_LINE_MAX_BYTES = 4 * 1024
49_CLIENT_MESSAGE_MAX_CHARS = 200
50_TRUNCATED_TEXT = "...[truncated]"
51_PARTIAL_STATE_DISCLAIMER = (
52 "Partial CloudFormation state may remain — inspect via stack_status or the AWS console."
53)
56def _argv_has_traversal(argv: Sequence[str]) -> tuple[int, str] | None:
57 """Return the first non-flag argv element containing a ``..`` segment."""
58 for index, value in enumerate(argv):
59 if value.startswith("-"):
60 continue
61 if ".." in value.split("/") or ".." in value.split("\\"):
62 return index, value[:100]
63 return None
66def _bounded_text(value: str, max_bytes: int) -> str:
67 """Return UTF-8 text within ``max_bytes`` with an explicit marker."""
68 encoded = value.encode("utf-8", errors="replace")
69 if len(encoded) <= max_bytes:
70 return value
71 marker = _TRUNCATED_TEXT.encode()
72 if max_bytes <= len(marker):
73 return marker[:max_bytes].decode("utf-8", errors="ignore")
74 prefix = encoded[: max_bytes - len(marker)]
75 return prefix.decode("utf-8", errors="ignore") + _TRUNCATED_TEXT
78async def _bounded_stream_lines(stream: asyncio.StreamReader) -> AsyncIterator[str]:
79 """Yield newline-delimited output without ever accumulating a giant line.
81 ``StreamReader`` async iteration delegates to ``readline()``, which can
82 raise once its internal limit is exceeded before application-level bounds
83 run. Fixed-size ``read()`` calls avoid that failure mode. Bytes beyond the
84 per-line budget are discarded until the next newline and represented by an
85 explicit truncation marker.
86 """
87 marker_bytes = _TRUNCATED_TEXT.encode()
88 payload_budget = max(0, _STREAM_LINE_MAX_BYTES - len(marker_bytes))
89 buffered = bytearray()
90 truncated = False
92 while True:
93 chunk = await stream.read(_STREAM_READ_BYTES)
94 if not chunk:
95 break
96 offset = 0
97 while offset < len(chunk):
98 newline = chunk.find(b"\n", offset)
99 end = len(chunk) if newline < 0 else newline
100 segment = chunk[offset:end]
101 if not truncated:
102 remaining = max(0, payload_budget - len(buffered))
103 if remaining: 103 ↛ 105line 103 didn't jump to line 105 because the condition on line 103 was always true
104 buffered.extend(segment[:remaining])
105 if len(segment) > remaining:
106 truncated = True
108 if newline < 0:
109 break
111 raw = bytes(buffered)
112 if raw.endswith(b"\r"):
113 raw = raw[:-1]
114 text = raw.decode("utf-8", errors="replace")
115 if truncated:
116 text += _TRUNCATED_TEXT
117 yield _bounded_text(text, _STREAM_LINE_MAX_BYTES)
118 buffered.clear()
119 truncated = False
120 offset = newline + 1
122 if buffered or truncated:
123 raw = bytes(buffered)
124 if raw.endswith(b"\r"):
125 raw = raw[:-1]
126 text = raw.decode("utf-8", errors="replace")
127 if truncated:
128 text += _TRUNCATED_TEXT
129 yield _bounded_text(text, _STREAM_LINE_MAX_BYTES)
132async def _best_effort_client_call(target: Any, method_name: str, *args: object) -> None:
133 """Call one optional async client notification without breaking work.
135 Client disconnects and version-skewed Progress implementations must not
136 terminate the underlying infrastructure operation. ``CancelledError`` is a
137 ``BaseException`` and intentionally still propagates.
138 """
139 try:
140 method = getattr(target, method_name)
141 await method(*args)
142 except Exception:
143 return
146async def _terminate_and_reap(
147 process: asyncio.subprocess.Process,
148 wait_task: asyncio.Task[int],
149) -> None:
150 """Terminate a subprocess, escalate after grace, and always reap it."""
151 if process.returncode is None:
152 with contextlib.suppress(OSError):
153 process.terminate()
154 try:
155 await asyncio.wait_for(asyncio.shield(wait_task), timeout=_CANCEL_GRACE_SECONDS)
156 return
157 except TimeoutError:
158 pass
159 except Exception:
160 # A failed wait task should not prevent the kill/reap fallback.
161 pass
163 if process.returncode is None: 163 ↛ 166line 163 didn't jump to line 166 because the condition on line 163 was always true
164 with contextlib.suppress(OSError):
165 process.kill()
166 if wait_task.cancelled():
167 wait_task = asyncio.create_task(process.wait())
168 with contextlib.suppress(Exception):
169 await asyncio.shield(wait_task)
170 if process.returncode is None:
171 # Defensive final wait if a mocked or unusual Process did not update
172 # returncode through the original wait task.
173 with contextlib.suppress(Exception):
174 await process.wait()
177async def _run_long_task(
178 argv: Sequence[str],
179 *,
180 ctx: Any,
181 progress: Any,
182 is_stack_op: bool = True,
183 total_units: int | None = None,
184) -> str:
185 """Run a long-lived command with bounded output and durable status.
187 Logical ``gco`` argv is resolved to the executable installed beside this
188 MCP environment and runs from the project root. Process wait and both pipe
189 drains are coordinated concurrently, so a drain error is surfaced promptly
190 instead of allowing an unread pipe to deadlock the child. Every exception
191 after spawn terminates/reaps the process and stamps terminal disk status.
192 """
193 logical_argv = list(argv)
194 hit = _argv_has_traversal(logical_argv)
195 if hit is not None:
196 index, value = hit
197 return json.dumps({"error": "path_traversal_detected", "argv_index": index, "value": value})
199 protocol_task_id = _try_get_task_id(ctx)
200 if len(logical_argv) >= 3 and logical_argv[0] == "gco":
201 tool_name = f"{logical_argv[1]}_{logical_argv[2].replace('-', '_')}"
202 elif len(logical_argv) >= 2 and logical_argv[0] == "gco":
203 tool_name = logical_argv[1]
204 else:
205 tool_name = logical_argv[0] if logical_argv else "task"
206 task_id = (
207 protocol_task_id
208 if isinstance(protocol_task_id, str) and is_valid_task_id(protocol_task_id)
209 else make_task_id(tool_name)
210 )
212 spawn_argv = list(logical_argv)
213 if spawn_argv and spawn_argv[0] == "gco":
214 spawn_argv[0] = cli_runner._gco_executable()
216 started = time.monotonic()
217 # Initialize observability before spawning. The writer degrades to an
218 # in-memory no-op when its directory is unavailable, so construction cannot
219 # strand a child process after spawn.
220 status_writer = TaskStatusWriter(
221 task_id=task_id,
222 tool=tool_name,
223 argv=logical_argv,
224 pid=None,
225 total_units=total_units,
226 )
228 process: asyncio.subprocess.Process | None = None
229 wait_task: asyncio.Task[int] | None = None
230 drains: list[asyncio.Task[None]] = []
231 coordination: asyncio.Future[list[Any]] | None = None
232 heartbeat: asyncio.Task[None] | None = None
233 stacks_completed = 0
234 failed_lines: deque[str] = deque(maxlen=_FAILED_EVENT_LINES)
235 stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES)
236 last_activity = time.monotonic()
237 last_stack: str | None = None
238 completed_stacks: set[str] = set()
239 return_code: int | None = None
241 async def _drain(stream: asyncio.StreamReader | None, label: str) -> None:
242 nonlocal stacks_completed, last_activity, last_stack
243 if stream is None:
244 raise RuntimeError(f"{label} pipe was not created")
245 async for line in _bounded_stream_lines(stream):
246 if not line:
247 continue
248 last_activity = time.monotonic()
249 status_writer.record_line(line, stream=label)
251 increment_progress = False
252 stack_done = _CDK_STACK_DONE_RE.search(line)
253 if stack_done is not None:
254 name = stack_done.group(1)
255 if name not in completed_stacks:
256 completed_stacks.add(name)
257 stacks_completed += 1
258 status_writer.increment_stacks(name)
259 increment_progress = True
261 if _CFN_FAILED_RE.search(line):
262 failed_lines.append(_bounded_text(line, _DIAGNOSTIC_LINE_MAX_BYTES))
263 stack_match = _CDK_STACK_LINE_RE.search(line)
264 if stack_match:
265 last_stack = stack_match.group(1)
266 status_writer.set_last_stack(last_stack)
267 if label == "stderr":
268 stderr_tail.append(_bounded_text(line, _DIAGNOSTIC_LINE_MAX_BYTES))
270 await _best_effort_client_call(
271 progress,
272 "set_message",
273 line[:_CLIENT_MESSAGE_MAX_CHARS],
274 )
275 if increment_progress:
276 await _best_effort_client_call(progress, "increment")
277 if label == "stderr":
278 await _best_effort_client_call(
279 ctx,
280 "info",
281 f"stderr: {line[:_CLIENT_MESSAGE_MAX_CHARS]}",
282 )
284 async def _heartbeat() -> None:
285 while True:
286 await asyncio.sleep(_HEARTBEAT_INTERVAL_SECONDS)
287 if time.monotonic() - last_activity < _HEARTBEAT_INTERVAL_SECONDS:
288 continue
289 elapsed = int(time.monotonic() - started)
290 stack_part = f" (last: {last_stack})" if last_stack else ""
291 message = f"still running … {_format_duration(elapsed)} elapsed{stack_part}"
292 await _best_effort_client_call(progress, "set_message", message)
293 await _best_effort_client_call(ctx, "info", message)
295 async def _clean_process_tasks() -> None:
296 """Terminate/reap the child and consume every auxiliary task result."""
297 nonlocal wait_task
298 if process is None:
299 return
300 if wait_task is None or wait_task.cancelled():
301 wait_task = asyncio.create_task(process.wait())
302 await _terminate_and_reap(process, wait_task)
303 for drain in drains:
304 if not drain.done(): 304 ↛ 305line 304 didn't jump to line 305 because the condition on line 304 was never true
305 drain.cancel()
306 if drains:
307 await asyncio.gather(*drains, return_exceptions=True)
308 if coordination is not None:
309 if not coordination.done(): 309 ↛ 310line 309 didn't jump to line 310 because the condition on line 309 was never true
310 coordination.cancel()
311 with contextlib.suppress(asyncio.CancelledError, Exception):
312 await coordination
314 try:
315 if total_units is not None and total_units > 0:
316 await _best_effort_client_call(progress, "set_total", int(total_units))
318 process = await asyncio.create_subprocess_exec(
319 *spawn_argv,
320 stdout=asyncio.subprocess.PIPE,
321 stderr=asyncio.subprocess.PIPE,
322 cwd=str(cli_runner.PROJECT_ROOT),
323 )
324 status_writer.set_pid(process.pid)
326 wait_task = asyncio.create_task(process.wait())
327 drains = [
328 asyncio.create_task(_drain(process.stdout, "stdout")),
329 asyncio.create_task(_drain(process.stderr, "stderr")),
330 ]
331 heartbeat = asyncio.create_task(_heartbeat())
332 # Shield keeps caller cancellation from cancelling wait/drain tasks
333 # before the process cleanup path can terminate and reap the child.
334 coordination = asyncio.gather(wait_task, *drains)
335 results = await asyncio.shield(coordination)
336 return_code = int(results[0])
337 except asyncio.CancelledError:
338 with contextlib.suppress(Exception):
339 await _clean_process_tasks()
340 status_writer.finish(
341 state="cancelled",
342 exit_code=process.returncode if process is not None else None,
343 error=_PARTIAL_STATE_DISCLAIMER if is_stack_op else "cancelled",
344 )
345 if is_stack_op:
346 raise asyncio.CancelledError(_PARTIAL_STATE_DISCLAIMER) from None
347 raise
348 except Exception as exc:
349 with contextlib.suppress(Exception):
350 await _clean_process_tasks()
351 status_writer.finish(
352 state="failed",
353 exit_code=process.returncode if process is not None else None,
354 error=_bounded_text(
355 f"{type(exc).__name__}: {exc}",
356 _DIAGNOSTIC_LINE_MAX_BYTES,
357 ),
358 )
359 raise
360 finally:
361 if heartbeat is not None: 361 ↛ 366line 361 didn't jump to line 366 because the condition on line 361 was always true
362 heartbeat.cancel()
363 with contextlib.suppress(asyncio.CancelledError, Exception):
364 await heartbeat
366 duration = int(time.monotonic() - started)
367 if return_code is None: 367 ↛ 370line 367 didn't jump to line 370 because the condition on line 367 was never true
368 # Coordination only completes normally after process.wait(). Keep the
369 # guard explicit for unusual Process implementations and static safety.
370 return_code = process.returncode if process is not None else None
371 if return_code is None: 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 status_writer.finish(state="failed", error="subprocess return code unavailable")
373 raise RuntimeError("subprocess return code unavailable")
375 if return_code != 0:
376 payload: dict[str, Any] = {
377 "error": f"exit_code={return_code}",
378 "exit_code": return_code,
379 "task_id": task_id,
380 "stacks_completed": stacks_completed,
381 "duration_seconds": duration,
382 "last_stack": last_stack,
383 "failed_events": list(failed_lines),
384 "stderr_tail": list(stderr_tail),
385 }
386 if is_stack_op:
387 payload["disclaimer"] = _PARTIAL_STATE_DISCLAIMER
388 status_writer.finish(
389 state="failed", exit_code=return_code, error=f"exit_code={return_code}"
390 )
391 raise ToolError(json.dumps(payload))
393 status_writer.finish(state="succeeded", exit_code=return_code)
394 return json.dumps(
395 {
396 "status": "ok",
397 "task_id": task_id,
398 "stacks_completed": stacks_completed,
399 "duration_seconds": duration,
400 "last_stack": last_stack,
401 }
402 )
405def _format_duration(seconds: int) -> str:
406 """Render an integer second count as ``HhMmSs`` / ``MmSs`` / ``Ss``."""
407 if seconds < 60:
408 return f"{seconds}s"
409 minutes, sec = divmod(seconds, 60)
410 if minutes < 60:
411 return f"{minutes}m{sec:02d}s"
412 hours, mins = divmod(minutes, 60)
413 return f"{hours}h{mins:02d}m{sec:02d}s"