Unify Compose/Helm lifecycle behavior in cds up/down/state

cli/k8s_runner.py

@@ -7,13 +7,20 @@
 import re
 import subprocess  # nosec B404
 import tempfile
+import time
+from collections.abc import Callable
 from pathlib import Path
 from typing import IO, Any
 
 import yaml
 
-from .state import parse_k8s_workloads_json
-from .up_runner import _validate_command, run_streamed
+from .state import group_services_by_health, parse_k8s_workloads_json
+from .up_runner import (
+    DEFAULT_POLL_INTERVAL_SECONDS,
+    _poll_until_settled,
+    _validate_command,
+    run_streamed,
+)
 
 # RFC 1123 DNS label: lowercase alphanumeric and `-`, not leading/trailing
 # with `-`, 1-63 chars. Both Kubernetes namespaces and Helm release names
@@ -90,16 +97,37 @@ def helm_up(
     timeout: float,
     detach: bool,
     log_file: IO[str],
+    poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
+    use_color: bool = False,
+    sleep_fn: Callable[[float], None] = time.sleep,
+    now_fn: Callable[[], float] = time.monotonic,
+    redraw_fn: Callable[[str], None] | None = None,
 ) -> int:
-    """Install or upgrade a rendered chart without persisting secret values."""
+    """
+    Install or upgrade a rendered chart without persisting secret values.
+
+    `timeout` is one overall budget shared across the whole call, the way
+    Compose's `--timeout` already is: it bounds `helm upgrade --install`
+    itself, then whatever remains bounds the readiness poll below, rather
+    than being reused in full for each step (which could let total
+    wall-clock time exceed `timeout` several times over on a
+    multi-workload release).
+
+    Readiness is decided by polling `get_k8s_state()` /
+    `group_services_by_health()` — the exact same functions `cds state
+    --target helm` uses — instead of a separate `helm --wait` +
+    `kubectl rollout status`/`wait` implementation, so "ready" here and
+    "healthy" in `cds state` can never disagree.
+    """
     namespace = _validate_k8s_name(namespace, "namespace")
     release = _validate_k8s_name(release, "release")
     kube_context = _validate_kube_context(kube_context)
     timeout = _validate_timeout(timeout)
+    deadline = now_fn() + timeout
     secret_values = _secret_values(plan)
     secret_path = _write_secret_values(secret_values)
     context_args = ["--kube-context", kube_context] if kube_context else []
-    timeout_arg = f"{timeout}s"
+    apply_timeout_arg = f"{timeout}s"
     command = [
         "helm",
         *context_args,
@@ -113,10 +141,8 @@ def helm_up(
         "--values",
         str(secret_path),
         "--timeout",
-        timeout_arg,
+        apply_timeout_arg,
     ]
-    if not detach:
-        command.append("--wait")
 
     try:
         result = run_streamed(command, log_file, timeout=timeout + 30)
@@ -126,31 +152,76 @@ def helm_up(
         return result
 
     workloads = get_k8s_workloads(namespace, release, kube_context)
-    for workload in workloads:
-        kind = str(workload.get("kind", "")).lower()
-        name = str((workload.get("metadata") or {}).get("name", ""))
-        if not kind or not name:
-            continue
-        if kind == "job":
-            wait_command = _kubectl_command(kube_context, namespace) + [
-                "wait",
-                "--for=condition=complete",
-                f"job/{name}",
-                f"--timeout={timeout_arg}",
-            ]
-        else:
-            wait_command = _kubectl_command(kube_context, namespace) + [
-                "rollout",
-                "status",
-                f"{kind}/{name}",
-                f"--timeout={timeout_arg}",
-            ]
-        result = run_streamed(wait_command, log_file, timeout=timeout + 30)
-        if result != 0:
-            return result
+    remaining = max(0.0, deadline - now_fn())
+    settled, grouped = poll_k8s_state_until_settled(
+        namespace,
+        release,
+        kube_context,
+        expected_service_count=len(workloads),
+        poll_interval=poll_interval,
+        timeout=remaining,
+        use_color=use_color,
+        sleep_fn=sleep_fn,
+        now_fn=now_fn,
+        redraw_fn=redraw_fn,
+    )
+    if not settled:
+        unhealthy = [name for bucket in ("UNHEALTHY", "UNHEALTHY EXIT") for name in grouped.get(bucket, [])]
+        log_file.write(
+            f"helm release {release} did not settle within {timeout:.0f}s"
+            + (f"; unhealthy: {', '.join(unhealthy)}" if unhealthy else "")
+            + "\n"
+        )
+        log_file.flush()
+        return 1
     return 0
 
 
+def poll_k8s_state_until_settled(
+    namespace: str,
+    release: str,
+    kube_context: str | None,
+    *,
+    expected_service_count: int | None = None,
+    poll_interval: float = DEFAULT_POLL_INTERVAL_SECONDS,
+    timeout: float = 180.0,
+    use_color: bool = False,
+    sleep_fn: Callable[[float], None] = time.sleep,
+    now_fn: Callable[[], float] = time.monotonic,
+    fetch_fn: Callable[[], list[dict[str, Any]]] | None = None,
+    redraw_fn: Callable[[str], None] | None = None,
+) -> tuple[bool, dict[str, list[str]]]:
+    """
+    Polls `get_k8s_state()` every `poll_interval` seconds, grouping with
+    `group_services_by_health()` — the same pair `cds state --target
+    helm` calls — until every workload settles into a terminal bucket or
+    `timeout` seconds elapse. Shares its settle/failure rule with
+    Compose's `poll_state_until_settled()` via `_poll_until_settled()`.
+
+    `fetch_fn` is injectable for tests; defaults to a real
+    `get_k8s_state(namespace, release, kube_context)` call.
+    """
+    if fetch_fn is None:
+        def fetch_fn() -> list[dict[str, Any]]:
+            return get_k8s_state(namespace, release, kube_context)
+
+    def fetch_grouped() -> dict[str, list[str]]:
+        return group_services_by_health(fetch_fn())
+
+    return _poll_until_settled(
+        fetch_grouped,
+        expected_service_count=expected_service_count,
+        poll_interval=poll_interval,
+        timeout=timeout,
+        use_color=use_color,
+        sleep_fn=sleep_fn,
+        now_fn=now_fn,
+        redraw_fn=redraw_fn,
+        up_done_fn=None,
+        on_up_finished=None,
+    )
+
+
 def helm_down(
     *,
     namespace: str,

cli/main.py

@@ -33,7 +33,7 @@
 from .k8s_runner import get_k8s_state, helm_down, helm_up
 from .k8s_security import scan_k8s_security
 from .loader import save_generated_profile
-from .overlay import resolve_extends, resolve_profile
+from .overlay import _merge_profile_docs, resolve_extends, resolve_profile
 from .planner import build_plan
 from .preflight import preflight_passed, run_preflight
 from .renderer import render_compose
@@ -97,19 +97,73 @@ def print_diagnostics(diagnostics) -> None:
         print(f"{prefix} {d.format()}\n")
 
 
-def _k8s_runtime_defaults(profile_path: str) -> tuple[str, str]:
-    """Read non-secret Helm defaults without planning or loading environment values."""
+def _helm_identity(
+    metadata: dict[str, Any] | None,
+    runtime: dict[str, Any] | None,
+    profile_path: str,
+) -> tuple[str, str]:
+    """
+    Shared release/namespace fallback rule: release from metadata.name (or
+    the profile's directory name), namespace from runtime.namespace (or
+    "cds-local"). `metadata`/`runtime` can come either from a fully built
+    plan (`up`) or from a best-effort document resolution (`down`/`state`)
+    as long as both use this same rule, so a stack brought up with a given
+    --environment is torn down/inspected against the same release and
+    namespace.
+    """
     fallback_release = Path(profile_path).parent.name or "cds"
+    release = str((metadata or {}).get("name") or fallback_release)
+    namespace = str((runtime or {}).get("namespace") or "cds-local")
+    return release, namespace
+
+
+def _resolve_profile_document_best_effort(
+    profile_path: str, environment: str | None = None
+) -> dict[str, Any]:
+    """
+    Resolves a profile's `extends` chain and, if given, an `--environment`
+    overlay, WITHOUT running the full validate_loaded_profile() gate that
+    resolve_profile() applies. This is deliberate: down/state must be able
+    to resolve the release/namespace a stack was brought up under even if
+    the profile has since been edited into an invalid state (e.g. an
+    unrelated module config error) -- falling back to generic defaults in
+    that case would target the wrong release/namespace instead of tearing
+    down (or reporting on) the one that's actually running.
+    """
+    document, provenance, diagnostics = resolve_extends(profile_path)
+    if document is None or any(d.level == "error" for d in diagnostics):
+        return {}
+
+    if not environment:
+        return document
+
+    profile_file = Path(profile_path)
+    overlay_file = profile_file.parent / "environments" / f"{environment}.yaml"
+    if not overlay_file.is_file():
+        return document
+
     try:
-        document = yaml.safe_load(Path(profile_path).read_text(encoding="utf-8")) or {}
+        overlay = yaml.safe_load(overlay_file.read_text(encoding="utf-8")) or {}
     except (OSError, UnicodeDecodeError, yaml.YAMLError):
-        return fallback_release, "cds-local"
-    release = str((document.get("metadata") or {}).get("name") or fallback_release)
-    namespace = str(
-        (((document.get("spec") or {}).get("runtime") or {}).get("namespace"))
-        or "cds-local"
+        return document
+
+    merged, merge_diagnostics = _merge_profile_docs(
+        document, overlay, str(profile_file), str(overlay_file), provenance
     )
-    return release, namespace
+    if merge_diagnostics:
+        return document
+    return merged
+
+
+def _k8s_runtime_defaults(profile_path: str, environment: str | None = None) -> tuple[str, str]:
+    """
+    Read non-secret Helm defaults without planning, using the same
+    release/namespace rule the built plan uses (`_helm_identity`), so
+    `down`/`state` target the same release an `up --environment ...` run
+    created.
+    """
+    document = _resolve_profile_document_best_effort(profile_path, environment)
+    return _helm_identity(document.get("metadata"), (document.get("spec") or {}).get("runtime"), profile_path)
 
 
 def profile_completer(prefix, parsed_args, **kwargs):
@@ -1123,6 +1177,7 @@ def main() -> int:
 
     down_parser = subparsers.add_parser("down", help="Stop or uninstall a running profile")
     _add_profile_arg(down_parser)
+    _add_environment_arg(down_parser)
     down_parser.add_argument(
         "--target", choices=["compose", "helm"], default="compose", help="Runtime target."
     )
@@ -1174,6 +1229,7 @@ def main() -> int:
         help="Show running service status grouped by health",
     )
     _add_profile_arg(state_parser)
+    _add_environment_arg(state_parser)
     state_parser.add_argument(
         "--target", choices=["compose", "helm"], default="compose", help="Runtime target."
     )
@@ -1348,7 +1404,7 @@ def main() -> int:
 
     args = parser.parse_args()
     environment_explicit = hasattr(args, "environment")
-    if args.command in {"validate", "plan", "render", "up", "test", "preflight", "init", "security"}:
+    if args.command in {"validate", "plan", "render", "up", "test", "preflight", "init", "security", "down", "state"}:
         args.environment = getattr(args, "environment", None) or load_saved_environment()
     if args.command in {"render", "up", "test"}:
         args.image_source = getattr(args, "image_source", None) or load_saved_image_source()
@@ -1618,8 +1674,16 @@ def main() -> int:
                 print("Cannot start stack because Helm rendering failed.")
                 return code or 1
 
-            namespace = args.namespace or plan.get("runtime", {}).get("namespace") or "cds-local"
-            release = args.release or plan.get("metadata", {}).get("name") or "cds"
+            plan_release, plan_namespace = _helm_identity(
+                plan.get("metadata"), plan.get("runtime"), profile_path
+            )
+            namespace = args.namespace or plan_namespace
+            release = args.release or plan_release
+            if args.no_build:
+                print(
+                    "NOTE --no-build has no effect with --target helm: "
+                    "the Helm target does not build local images yet."
+                )
             log_path = (
                 Path(args.log_file) if args.log_file else default_log_path(Path(profile_path).parent.name)
             ).resolve()
@@ -1635,6 +1699,7 @@ def main() -> int:
                         timeout=args.timeout,
                         detach=args.detach,
                         log_file=log_file,
+                        use_color=(not args.no_color) and sys.stdout.isatty(),
                     )
             except (FileNotFoundError, OSError, RuntimeError, ValueError) as exc:
                 print(f"ERROR {exc}")
@@ -1837,7 +1902,7 @@ def _begin_log_tail(_up_exit_code: int) -> None:
             print(f"ERROR {exc}")
             return 1
         project_root = resolve_project_root(profile_path)
-        profile_release, profile_namespace = _k8s_runtime_defaults(profile_path)
+        profile_release, profile_namespace = _k8s_runtime_defaults(profile_path, args.environment)
         log_path = (
             Path(args.log_file) if args.log_file else default_log_path(f"down-{Path(profile_path).parent.name}")
         ).resolve()
@@ -2038,7 +2103,7 @@ def _begin_log_tail(_up_exit_code: int) -> None:
             return 1
 
         if args.target == "helm":
-            profile_release, profile_namespace = _k8s_runtime_defaults(profile_path)
+            profile_release, profile_namespace = _k8s_runtime_defaults(profile_path, args.environment)
             try:
                 services = get_k8s_state(
                     args.namespace or profile_namespace,