Audapure/wo cli improvements

src/workload-orchestration/azext_workload_orchestration/common/context.py

@@ -288,8 +288,6 @@ def capability_add(cli_ctx, resource_group, context_name, name=None,
         return ctx
 
     merged = existing + added
-    names_str = ", ".join(c["name"] for c in added)
-    _log(f"Adding {len(added)}: {names_str}")
 
     updated = _patch_context_capabilities(
         cli_ctx, sub_id, resource_group, context_name, merged

src/workload-orchestration/azext_workload_orchestration/common/target.py

@@ -40,6 +40,7 @@
 )
 from azext_workload_orchestration.common.utils import (
     _eprint,
+    ensure_required_cli_extensions,
     invoke_cli_command,
 )
 
@@ -84,6 +85,17 @@ def target_prepare(
 
     step_results = {}
 
+    # Step 0: Ensure required az CLI extensions are installed
+    # (connectedk8s, k8s-extension, customlocation are called via invoke_cli_command
+    # and would fail with opaque errors if missing)
+    try:
+        ensure_required_cli_extensions()
+        step_results["cli-extensions"] = "Ready"
+    except Exception as exc:
+        step_results["cli-extensions"] = f"FAILED: {exc}"
+        _print_failure_hint(step_results)
+        raise
+
     try:
         connected_cluster_id = _preflight_checks(cmd, cluster_name, resource_group)
         step_results["preflight"] = "Passed"

src/workload-orchestration/azext_workload_orchestration/common/utils.py

@@ -156,3 +156,74 @@ def print_step(step_num, total, message, status=""):
         _eprint(f"{connector} {message} {status}")
     else:
         _eprint(f"{connector} {message}...")
+
+
+# ---------------------------------------------------------------------------
+# CLI extension dependency check
+# ---------------------------------------------------------------------------
+
+# az CLI extensions that workload-orchestration calls at runtime via
+# invoke_cli_command. These MUST be installed before `cluster init` runs,
+# otherwise sub-command invocations fail with opaque "command not recognized"
+# errors. Mirrors the azext_vme.utils.check_and_add_cli_extension pattern.
+REQUIRED_CLI_EXTENSIONS = [
+    "connectedk8s",     # `connectedk8s show` in _preflight_checks
+    "k8s-extension",    # `k8s-extension list/create` for aio-certmgr + wo-extension
+    "customlocation",   # `customlocation create` for Step 4
+]
+
+
+def check_and_add_cli_extension(extension_name):
+    """Check if an az CLI extension is installed; install it if missing.
+
+    Uses subprocess (not invoke_cli_command) so that `az extension add`
+    runs in its own CLI process — needed because in-process invoke does
+    not pick up newly installed extensions in the same Python process.
+
+    Raises CLIInternalError if the install fails.
+    """
+    import shutil
+    import subprocess
+
+    az = shutil.which("az") or "az"
+
+    # Check if installed
+    try:
+        result = subprocess.run(
+            [az, "extension", "list",
+             "--query", f"[?name=='{extension_name}'].name",
+             "-o", "tsv"],
+            capture_output=True, text=True, check=True, encoding="utf-8"
+        )
+        if extension_name in (result.stdout or "").strip():
+            logger.debug("az cli extension '%s' already installed", extension_name)
+            return
+    except subprocess.CalledProcessError as exc:
+        raise CLIInternalError(
+            f"Failed to check for az cli extension '{extension_name}': {exc.stderr or exc}"
+        ) from None
+
+    _eprint(f"  ├── Installing required az cli extension: {extension_name}...")
+    try:
+        subprocess.run(
+            [az, "extension", "add", "--name", extension_name, "--yes"],
+            capture_output=True, text=True, check=True, encoding="utf-8"
+        )
+        _eprint(f"  ├── Installed az cli extension: {extension_name} ✓")
+    except subprocess.CalledProcessError as exc:
+        raise CLIInternalError(
+            f"Failed to install required az cli extension '{extension_name}': "
+            f"{exc.stderr or exc}\n"
+            f"Please install manually: az extension add --name {extension_name}"
+        ) from None
+
+
+def ensure_required_cli_extensions(extension_names=None):
+    """Ensure all required az CLI extensions are installed.
+
+    Defaults to REQUIRED_CLI_EXTENSIONS. Idempotent — skips already-installed
+    extensions. Fails fast with a clear message if any install fails.
+    """
+    extensions = extension_names if extension_names is not None else REQUIRED_CLI_EXTENSIONS
+    for ext in extensions:
+        check_and_add_cli_extension(ext)