feat(pipeline-run): add task-status and task-wait commands

packages/tangle-cli/src/tangle_cli/pipeline_run_manager.py

@@ -12,6 +12,7 @@
 import copy
 import inspect
 import json
+import math
 import re
 import time
 import uuid
@@ -21,6 +22,7 @@
 from pathlib import Path
 from typing import Any, Mapping
 
+import requests
 import yaml
 
 from .handler import TangleCliHandler
@@ -32,9 +34,42 @@
 from .pipeline_run_search import PipelineRunSearch
 from .utils import dump_yaml
 
-_TERMINAL_STATUSES = ("FAILED", "SYSTEM_ERROR", "CANCELLED", "CANCELED", "SKIPPED", "SUCCEEDED", "INVALID")
-_ACTIVE_STATUSES = ("RUNNING", "CANCELLING", "CANCELING", "PENDING", "QUEUED")
+# Precedence order for reducing a mixed terminal aggregate to a single status:
+# failure terminals must precede the non-failure terminals SKIPPED and SUCCEEDED
+# so a mixed aggregate (e.g. {SKIPPED, INVALID}) reduces to the failure instead
+# of masking it (a fully-SKIPPED map still reduces to SKIPPED).
+_TERMINAL_PRECEDENCE = ("FAILED", "SYSTEM_ERROR", "CANCELLED", "CANCELED", "INVALID", "SKIPPED", "SUCCEEDED")
+# Membership set for is_terminal_status.
+_TERMINAL_STATUSES = frozenset(_TERMINAL_PRECEDENCE)
+# Nonterminal statuses outrank any terminal count when reducing an aggregate;
+# WAITING_FOR_UPSTREAM/UNINITIALIZED are backend enum members and nonterminal.
+_ACTIVE_STATUSES = (
+    "RUNNING",
+    "CANCELLING",
+    "CANCELING",
+    "PENDING",
+    "QUEUED",
+    "WAITING_FOR_UPSTREAM",
+    "UNINITIALIZED",
+)
 _FAILURE_EARLY_EXIT_STATUSES = ("FAILED", "SYSTEM_ERROR")
+# The only terminals that are not per-task failures; every other terminal makes
+# a root-execution `task-wait` exit non-zero. A single SKIPPED task (e.g. a
+# conditional branch) is not a failure, unlike a run whose reduced status is
+# SKIPPED at the run level.
+_NON_FAILURE_TERMINAL_STATUSES = frozenset({"SKIPPED", "SUCCEEDED"})
+_TASK_FAILURE_STATUSES = _TERMINAL_STATUSES - _NON_FAILURE_TERMINAL_STATUSES
+# A child execution's status is unknown until its state row is written; 404
+# means "wrong endpoint type" (container vs graph) and transient 5xx can be
+# returned while the orchestrator is still settling. Both fall through to the
+# next endpoint / report UNKNOWN so the wait loop polls again.
+_RETRYABLE_EXEC_STATE_STATUSES = frozenset({404, 500, 502, 503, 504})
+# An unbounded task-wait (max_wait=None) has no deadline to stop retrying a
+# failing root details fetch, so give up after this many consecutive server
+# errors instead of hiding a persistent outage behind endless retries.
+_MAX_UNBOUNDED_ROOT_FETCH_FAILURES = 3
+_UNKNOWN_TASK_STATUS = "UNKNOWN"
+_ROOT_TASK_NAME = "root"
 _EXECUTION_STATE_TIMINGS_METADATA_KEY = "execution_state_timings"
 _EXECUTION_STATE_TIMING_MONOTONIC_METADATA_KEY = "_execution_state_timing_monotonic"
 _SUBMISSION_ID_ANNOTATION_KEY = "tangle-cli/submission-id"
@@ -54,6 +89,51 @@ class AmbiguousPipelineRunRecoveryError(PipelineRunError):
     """Raised when submit recovery finds multiple runs for one submission id."""
 
 
+class TransientServerError(PipelineRunError):
+    """An HTTP 5xx failure that a deadline-bounded polling loop may retry.
+
+    One-shot callers (e.g. ``task-status``) still fail fast on it like any
+    other PipelineRunError.
+    """
+
+
+def _transport_error(context: str, exc: requests.RequestException) -> PipelineRunError:
+    """Wrap a transport/HTTP failure as a clean CLI-surfaced PipelineRunError.
+
+    Root-execution status resolution talks to the API directly, so connection
+    resets/timeouts and non-retryable 4xx/5xx (including 429) would otherwise
+    escape task-status/task-wait as raw ``requests`` tracebacks. The CLI maps
+    PipelineRunError to a non-zero exit with the message instead. Server errors
+    are typed TransientServerError so polling loops can retry them.
+    """
+
+    response = getattr(exc, "response", None)
+    status_code = getattr(response, "status_code", None)
+    if status_code is not None:
+        message = f"{context}: request failed with HTTP {status_code}"
+        if status_code >= 500:
+            return TransientServerError(message)
+        return PipelineRunError(message)
+    return PipelineRunError(f"{context}: {exc}")
+
+
+class TaskStatusesFailed(PipelineRunError):
+    """Raised when a root execution reaches terminal state with failed tasks.
+
+    Carries the full final ``{task_name: status}`` map alongside the ``failures``
+    subset so callers can still print the complete map while reporting failures.
+    """
+
+    def __init__(self, root_execution_id: str, statuses: dict[str, str], failures: dict[str, str]):
+        self.root_execution_id = root_execution_id
+        self.statuses = statuses
+        self.failures = failures
+        summary = ", ".join(f"{name}={status}" for name, status in sorted(failures.items()))
+        super().__init__(
+            f"Root execution {root_execution_id} had {len(failures)} failed task(s): {summary}"
+        )
+
+
 @dataclass
 class PipelineSubmitPayload:
     """Prepared submit payload state before calling ``pipeline_runs_create``.
@@ -817,11 +897,30 @@ def status_from_counts(status_counts: Mapping[str, int]) -> str | None:
         for status in _ACTIVE_STATUSES:
             if int(status_counts.get(status, 0) or 0) > 0:
                 return status
-        for status in _TERMINAL_STATUSES:
+        for status in _TERMINAL_PRECEDENCE:
             if int(status_counts.get(status, 0) or 0) > 0:
                 return status
         return None
 
+    @staticmethod
+    def _status_from_container_state(state: Mapping[str, Any] | Any) -> str | None:
+        plain = PipelineRunManager.to_plain(state)
+        if isinstance(plain, Mapping):
+            status = plain.get("status")
+            if isinstance(status, str) and status:
+                return status
+        return None
+
+    @classmethod
+    def _status_from_graph_state(cls, state: Mapping[str, Any] | Any) -> str | None:
+        # A nested graph/subpipeline execution carries no direct ``status``; its
+        # status is derived from child execution counts the same way run-level
+        # graph polling does.
+        direct = cls._status_from_container_state(state)
+        if direct is not None:
+            return direct
+        return cls.status_from_counts(cls.status_counts_from_graph_state(cls.to_plain(state)))
+
     @staticmethod
     def status_from_run(run: Mapping[str, Any]) -> str | None:
         summary = run.get("execution_summary")
@@ -839,7 +938,7 @@ def status_from_run(run: Mapping[str, Any]) -> str | None:
             for status in _ACTIVE_STATUSES:
                 if int(stats.get(status, 0) or 0) > 0:
                     return status
-            for status in _TERMINAL_STATUSES:
+            for status in _TERMINAL_PRECEDENCE:
                 if int(stats.get(status, 0) or 0) > 0:
                     return status
         return None
@@ -1261,6 +1360,174 @@ def graph_state_output(self, run_ids: list[str], *, timeout: float = 30.0) -> di
     def logs(self, execution_id: str) -> dict[str, Any]:
         return self.to_plain(self.hooks.fetch_logs(self.client, execution_id))
 
+    def _exec_status(self, execution_id: str, *, deadline: float | None = None) -> str | None:
+        """Resolve a single execution's status, trying container then graph state.
+
+        404 means the execution is the other kind (container vs graph); a
+        transient 5xx means the state row is still being written. Both cases
+        fall through to the next endpoint and, if neither resolves, return
+        ``None`` so the caller reports UNKNOWN. Once past ``deadline`` (a
+        ``time.monotonic()`` timestamp) no further requests are issued, so one
+        poll cannot overrun the caller's wait budget.
+        """
+
+        endpoints = (
+            (self.client.executions_container_state, self._status_from_container_state),
+            (self.client.executions_graph_execution_state, self._status_from_graph_state),
+        )
+        context = f"Resolving status for execution {execution_id}"
+        last_server_error: int | None = None
+        for endpoint, extract_status in endpoints:
+            if deadline is not None and time.monotonic() >= deadline:
+                return None
+            try:
+                state = endpoint(execution_id)
+            except requests.HTTPError as exc:
+                response = exc.response
+                if response is not None and response.status_code in _RETRYABLE_EXEC_STATE_STATUSES:
+                    if response.status_code >= 500:
+                        last_server_error = response.status_code
+                    continue
+                raise _transport_error(context, exc) from exc
+            except requests.RequestException as exc:
+                raise _transport_error(context, exc) from exc
+            status = extract_status(state)
+            if status is not None:
+                return status
+        if last_server_error is not None:
+            # Without this a wait that times out reports UNKNOWN with no hint
+            # that the state API was failing the whole time.
+            self.logger.warn(
+                f"Status for execution {execution_id} is UNKNOWN: state endpoints "
+                f"kept failing with HTTP {last_server_error}"
+            )
+        return None
+
+    def task_statuses(
+        self, root_execution_id: str, *, deadline: float | None = None
+    ) -> dict[str, str]:
+        """Return a ``{task_name: status}`` map for a root execution.
+
+        Walks ``child_task_execution_ids`` from the root execution details and
+        resolves each child's status. A leaf/root-only execution (no children)
+        reports a single ``{"root": status}`` entry; unresolved child statuses
+        are reported as ``UNKNOWN``. A missing root id (404) is a clean
+        not-found error; other transport/HTTP failures surface cleanly too.
+        Past ``deadline`` (a ``time.monotonic()`` timestamp) remaining statuses
+        are reported as ``UNKNOWN`` without further requests.
+        """
+
+        try:
+            details = self.to_plain(self.client.executions_details(root_execution_id))
+        except requests.HTTPError as exc:
+            response = exc.response
+            if response is not None and response.status_code == 404:
+                raise PipelineRunError(f"Root execution {root_execution_id} not found") from exc
+            raise _transport_error(f"Fetching root execution {root_execution_id}", exc) from exc
+        except requests.RequestException as exc:
+            raise _transport_error(f"Fetching root execution {root_execution_id}", exc) from exc
+        children = details.get("child_task_execution_ids") if isinstance(details, Mapping) else None
+        if not isinstance(children, Mapping) or not children:
+            return {
+                _ROOT_TASK_NAME: self._exec_status(root_execution_id, deadline=deadline)
+                or _UNKNOWN_TASK_STATUS
+            }
+        return {
+            str(task_name): self._child_status(execution_id, deadline=deadline)
+            for task_name, execution_id in children.items()
+        }
+
+    def _child_status(self, execution_id: Any, *, deadline: float | None = None) -> str:
+        # A None/empty/whitespace-only child id has no execution to query;
+        # report UNKNOWN instead of polling a literal blank id forever.
+        child_id = str(execution_id).strip() if execution_id else ""
+        if not child_id:
+            return _UNKNOWN_TASK_STATUS
+        return self._exec_status(child_id, deadline=deadline) or _UNKNOWN_TASK_STATUS
+
+    def wait_for_task_statuses(
+        self,
+        root_execution_id: str,
+        *,
+        max_wait: float | None = 1800.0,
+        poll_interval: float = 5.0,
+        allow_failure: bool = False,
+    ) -> dict[str, str]:
+        """Poll ``task_statuses`` until every task is terminal.
+
+        Returns the final ``{task_name: status}`` map. Raises
+        ``TaskStatusesFailed`` when any task ends in a failure status unless
+        ``allow_failure`` is set, and ``PipelineRunError`` on timeout.
+        ``max_wait=None`` waits indefinitely.
+        """
+
+        self._validate_wait_params(
+            max_wait=max_wait,
+            poll_interval=poll_interval,
+            timeout_clock="monotonic",
+            allow_zero_poll_interval=False,
+        )
+        deadline = None if max_wait is None else time.monotonic() + max_wait
+        root_fetch_failures = 0
+        while True:
+            try:
+                statuses = self.task_statuses(root_execution_id, deadline=deadline)
+            except TransientServerError as exc:
+                # A 5xx on the root details fetch is tolerated like child-state
+                # 5xx (retried next poll) so one blip cannot abort a long wait.
+                # The deadline still bounds the retries; an unbounded wait gives
+                # up after a few consecutive failures instead of masking a
+                # persistent outage forever.
+                root_fetch_failures += 1
+                if deadline is not None and time.monotonic() >= deadline:
+                    raise
+                if deadline is None and root_fetch_failures >= _MAX_UNBOUNDED_ROOT_FETCH_FAILURES:
+                    raise
+                if root_fetch_failures == 1:
+                    # Warn once per outage, not once per poll, so a long outage
+                    # does not flood the logs.
+                    self.logger.warn(f"{exc}; retrying on the next poll")
+                self._sleep_before_next_poll(poll_interval, deadline)
+                continue
+            root_fetch_failures = 0
+            if statuses and all(self.is_terminal_status(status) for status in statuses.values()):
+                return self._handle_terminal_task_statuses(root_execution_id, statuses, allow_failure)
+            if deadline is not None and time.monotonic() >= deadline:
+                pending = ", ".join(
+                    f"{name}={status}"
+                    for name, status in sorted(statuses.items())
+                    if not self.is_terminal_status(status)
+                )
+                raise PipelineRunError(
+                    f"Root execution {root_execution_id} still has non-terminal tasks "
+                    f"after {max_wait:g}s: {pending}"
+                )
+            self._sleep_before_next_poll(poll_interval, deadline)
+
+    @staticmethod
+    def _sleep_before_next_poll(poll_interval: float, deadline: float | None) -> None:
+        if deadline is None:
+            time.sleep(poll_interval)
+            return
+        # Clamp the final sleep to the remaining budget so the loop times out
+        # at max_wait instead of overshooting by up to a poll interval.
+        time.sleep(min(poll_interval, max(0.0, deadline - time.monotonic())))
+
+    def _handle_terminal_task_statuses(
+        self,
+        root_execution_id: str,
+        statuses: dict[str, str],
+        allow_failure: bool,
+    ) -> dict[str, str]:
+        failures = {
+            name: status
+            for name, status in statuses.items()
+            if status.upper() in _TASK_FAILURE_STATUSES
+        }
+        if failures and not allow_failure:
+            raise TaskStatusesFailed(root_execution_id, statuses, failures)
+        return statuses
+
     def search_runs(
         self,
         *,
@@ -1452,6 +1719,32 @@ def _poll_run_status(
             execution_state_timings=execution_state_timings,
         )
 
+    @staticmethod
+    def _validate_wait_params(
+        *,
+        max_wait: float | None,
+        poll_interval: float,
+        timeout_clock: str,
+        allow_zero_poll_interval: bool,
+    ) -> None:
+        """Validate wait/poll parameters, raising before any run is submitted."""
+
+        # NaN/inf pass the sign checks below (NaN compares False, inf is
+        # "positive") and would become a never-firing deadline or a raw
+        # ValueError from time.sleep; reject them up front. An unbounded wait is
+        # spelled max_wait=None, not an explicit inf.
+        if max_wait is not None:
+            if not math.isfinite(max_wait):
+                raise PipelineRunError("--max-wait must be a finite number")
+            if max_wait < 0:
+                raise PipelineRunError("--max-wait must be non-negative")
+        if not math.isfinite(poll_interval):
+            raise PipelineRunError("--poll-interval must be a finite number")
+        if poll_interval < 0 or (poll_interval == 0 and not allow_zero_poll_interval):
+            raise PipelineRunError("--poll-interval must be positive")
+        if timeout_clock not in {"monotonic", "wall"}:
+            raise PipelineRunError("timeout_clock must be 'monotonic' or 'wall'")
+
     def wait_for_completion(
         self,
         run_id: str,
@@ -1467,12 +1760,12 @@ def wait_for_completion(
         wait_context = context or PipelineRunContext(run_id=run_id, start_time=time.time())
         if exit_on_first_failure:
             wait_context.metadata["exit_on_first_failure"] = True
-        if max_wait is not None and max_wait < 0:
-            raise PipelineRunError("--max-wait must be non-negative")
-        if poll_interval < 0 or (poll_interval == 0 and not allow_zero_poll_interval):
-            raise PipelineRunError("--poll-interval must be positive")
-        if timeout_clock not in {"monotonic", "wall"}:
-            raise PipelineRunError("timeout_clock must be 'monotonic' or 'wall'")
+        self._validate_wait_params(
+            max_wait=max_wait,
+            poll_interval=poll_interval,
+            timeout_clock=timeout_clock,
+            allow_zero_poll_interval=allow_zero_poll_interval,
+        )
         enforce_max_wait = max_wait is not None and self.hooks.should_enforce_max_wait(wait_context)
         poll_started_at = time.monotonic()
         deadline_now: Callable[[], float] = time.time if timeout_clock == "wall" else time.monotonic