1 files changed, 55 insertions, 4 deletions
diff --git a/makima/src/bin/makima.rs b/makima/src/bin/makima.rs
index f0db69c..1a307d1 100644
--- a/makima/src/bin/makima.rs
+++ b/makima/src/bin/makima.rs
@@ -335,11 +335,62 @@ async fn run_supervisor(
         SupervisorCommand::Wait(args) => {
             let client = ApiClient::new(args.common.api_url, args.common.api_key)?;
             eprintln!(
-                "Waiting for task {} (timeout: {}s)...",
-                args.task_id, args.timeout
+                "Waiting for task {} (timeout: {}s, poll interval: {}s)...",
+                args.task_id, args.timeout, args.poll_interval
             );
-            let result = client.supervisor_wait(args.task_id, args.timeout).await?;
-            println!("{}", serde_json::to_string(&result.0)?);
+
+            let start_time = std::time::Instant::now();
+            let timeout_duration = std::time::Duration::from_secs(args.timeout as u64);
+            let poll_interval = std::time::Duration::from_secs(args.poll_interval);
+            let server_wait_timeout = 30i32; // Short timeout for server-side wait
+
+            loop {
+                // Check if we've exceeded the total timeout
+                let remaining = timeout_duration.saturating_sub(start_time.elapsed());
+                if remaining.is_zero() {
+                    eprintln!("Timeout reached after {}s", args.timeout);
+                    let result = client.supervisor_get_task(args.task_id).await?;
+                    println!("{}", serde_json::to_string(&result.0)?);
+                    break;
+                }
+
+                // Try server-side wait with short timeout
+                let wait_timeout = std::cmp::min(server_wait_timeout, remaining.as_secs() as i32);
+
+                match client.supervisor_wait(args.task_id, wait_timeout).await {
+                    Ok(result) => {
+                        if let Some(completed) = result.0.get("completed").and_then(|c| c.as_bool()) {
+                            if completed {
+                                println!("{}", serde_json::to_string(&result.0)?);
+                                break;
+                            }
+                        }
+                        // Not completed yet, continue loop
+                        eprintln!("Task still running (elapsed: {:?})", start_time.elapsed());
+                    }
+                    Err(e) => {
+                        eprintln!("Warning: Server wait failed: {}. Falling back to polling...", e);
+                        // Fall back to simple status poll
+                        if let Ok(result) = client.supervisor_get_task(args.task_id).await {
+                            if let Some(status) = result.0.get("status").and_then(|s| s.as_str()) {
+                                if status == "done" || status == "failed" || status == "merged" {
+                                    let wait_response = serde_json::json!({
+                                        "taskId": args.task_id,
+                                        "status": status,
+                                        "completed": true,
+                                        "outputSummary": result.0.get("progressSummary")
+                                    });
+                                    println!("{}", serde_json::to_string(&wait_response)?);
+                                    break;
+                                }
+                            }
+                        }
+                    }
+                }
+
+                // Small delay before retrying
+                tokio::time::sleep(poll_interval).await;
+            }
         }
         SupervisorCommand::ReadFile(args) => {
             let client = ApiClient::new(args.common.api_url, args.common.api_key)?;