Resolve Pimox worker count from cluster state

jeannie

@@ -3893,13 +3893,68 @@ start_remote_kubernetes_services() {
 
 pimox_worker_count_default() {
     local spec_file="${REPO_ROOT}/.lab/pimox-workers.tsv"
+    local var_file="${REPO_ROOT}/.lab/cluster-workers.auto.tfvars.json"
+    local worker_key_prefix="${LAB_PIMOX_WORKER_KEY_PREFIX:-pimox}"
+    local worker_node_prefix="${LAB_PIMOX_WORKER_NODE_PREFIX:-pimox-worker}"
+    local max_count=0
+    local count
+    local state_count
 
     if [[ -s "${spec_file}" ]]; then
-        awk 'NF { count++ } END { print count + 0 }' "${spec_file}"
+        count="$(awk -F'\t' -v key_prefix="${worker_key_prefix}" -v node_prefix="${worker_node_prefix}-" '
-        return 0
+            $1 ~ "^" key_prefix "[0-9]+$" {
+                index = substr($1, length(key_prefix) + 1) + 0
+                if (index > max) max = index
+            }
+            $4 ~ "^" node_prefix "[0-9]+$" {
+                index = substr($4, length(node_prefix) + 1) + 0
+                if (index > max) max = index
+            }
+            END { print max + 0 }
+        ' "${spec_file}")"
+        if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
+            max_count="${count}"
+        fi
     fi
 
-    printf '1\n'
+    if [[ -s "${var_file}" ]] && command -v python3 >/dev/null 2>&1; then
+        count="$(python3 - "${var_file}" "${worker_key_prefix}" "${worker_node_prefix}" <<'PY'
+import json
+import re
+import sys
+
+var_file, worker_key_prefix, worker_node_prefix = sys.argv[1:4]
+with open(var_file, encoding="utf-8") as handle:
+    document = json.load(handle)
+nodes = document.get("worker_nodes") or {}
+highest = 0
+key_pattern = re.compile(rf"^{re.escape(worker_key_prefix)}(\d+)$")
+node_pattern = re.compile(rf"^{re.escape(worker_node_prefix)}-(\d+)$")
+for key, node in nodes.items():
+    for candidate in (key, str(node.get("node_name", ""))):
+        match = key_pattern.match(candidate) or node_pattern.match(candidate)
+        if match:
+            highest = max(highest, int(match.group(1)))
+print(highest)
+PY
+)"
+        if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
+            max_count="${count}"
+        fi
+    fi
+
+    state_count="$(tofu -chdir="${REPO_ROOT}/bootstrap/cluster" state show null_resource.worker_nodes_required 2>/dev/null |
+        awk -F'"' '/"worker_count"[[:space:]]*=/ { print $4; found = 1 } END { exit found ? 0 : 1 }' || true)"
+    if [[ "${state_count}" =~ ^[0-9]+$ && "${state_count}" -gt "${max_count}" ]]; then
+        max_count="${state_count}"
+    fi
+
+    count="${LAB_PIMOX_DEFAULT_WORKER_COUNT:-1}"
+    if [[ "${count}" =~ ^[0-9]+$ && "${count}" -gt "${max_count}" ]]; then
+        max_count="${count}"
+    fi
+
+    printf '%s\n' "${max_count}"
 }
 
 stop_pimox_worker_vms() {