Fix wait command to always check replica count (#106)

Author: firecow

src/commands/wait-cmd.ts

@@ -41,24 +41,21 @@ export async function handler (args: ArgumentsCamelCase) {
 
         serviceStateMap = new Map<string, string>();
         services = await dockerode.listServices({filters: {label: [`com.docker.stack.namespace=${appName}`]}});
-        tasks = await dockerode.listTasks({filters: {"label": [`com.docker.stack.namespace=${appName}`], "desired-state": ["running"]}}) as Task[];
+        tasks = await dockerode.listTasks({filters: {"label": [`com.docker.stack.namespace=${appName}`]}}) as Task[];
 
         for (const s of services) {
-            if (s.UpdateStatus?.State) {
+            const runningTasks = tasks.filter((t) => t.Status.State === "running" && t.ServiceID === s.ID);
+            const desiredReplicas = s.Spec?.Mode?.Replicated?.Replicas ?? 0;
+
+            // Always check replica count first - compare running tasks against desired replicas
+            if (runningTasks.length < desiredReplicas) {
+                serviceStateMap.set(s.ID, "replicating");
+            } else if (s.UpdateStatus?.State && ["updating", "paused", "rollback_started", "rollback_paused"].includes(s.UpdateStatus.State)) {
                 serviceStateMap.set(s.ID, s.UpdateStatus.State);
-            } else {
-                const runningTasks = tasks.filter((t) => t.Status.State === "running" && t.ServiceID === s.ID);
-                const totalTasks = tasks.filter((t) => t.ServiceID === s.ID);
-                if (totalTasks.length > runningTasks.length) {
-                    serviceStateMap.set(s.ID, "replicating");
-                }
             }
         }
 
-        const servicesUpdating = [];
-        for (const [id, state] of serviceStateMap) {
-            if (!["completed", "rollback_completed"].includes(state)) servicesUpdating.push([id, state]);
-        }
+        const servicesUpdating = [...serviceStateMap.entries()];
 
         bail = servicesUpdating.length === 0;
         if (!bail) {