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