huskies: merge 476_refactor_split_agents_pool_lifecycle_rs_into_submodules

server/src/agents/pool/wait.rs

@@ -0,0 +1,190 @@
+use super::super::{AgentEvent, AgentInfo, AgentStatus};
+use super::types::{agent_info_from_entry, composite_key};
+use super::AgentPool;
+
+use tokio::sync::broadcast;
+
+impl AgentPool {
+    /// Block until the agent reaches a terminal state (completed, failed, stopped).
+    /// Returns the agent's final `AgentInfo`.
+    /// `timeout_ms` caps how long to wait; returns an error if the deadline passes.
+    pub async fn wait_for_agent(
+        &self,
+        story_id: &str,
+        agent_name: &str,
+        timeout_ms: u64,
+    ) -> Result<AgentInfo, String> {
+        // Subscribe before checking status so we don't miss the terminal event
+        // if the agent completes in the window between the two operations.
+        let mut rx = self.subscribe(story_id, agent_name)?;
+
+        // Return immediately if already in a terminal state.
+        {
+            let agents = self.agents.lock().map_err(|e| e.to_string())?;
+            let key = composite_key(story_id, agent_name);
+            if let Some(agent) = agents.get(&key)
+                && matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
+            {
+                return Ok(agent_info_from_entry(story_id, agent));
+            }
+        }
+
+        let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(timeout_ms);
+
+        loop {
+            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
+            if remaining.is_zero() {
+                return Err(format!(
+                    "Timed out after {timeout_ms}ms waiting for agent '{agent_name}' on story '{story_id}'"
+                ));
+            }
+
+            match tokio::time::timeout(remaining, rx.recv()).await {
+                Ok(Ok(event)) => {
+                    let is_terminal = match &event {
+                        AgentEvent::Done { .. } | AgentEvent::Error { .. } => true,
+                        AgentEvent::Status { status, .. } if status == "stopped" => true,
+                        _ => false,
+                    };
+                    if is_terminal {
+                        let agents = self.agents.lock().map_err(|e| e.to_string())?;
+                        let key = composite_key(story_id, agent_name);
+                        return Ok(if let Some(agent) = agents.get(&key) {
+                            agent_info_from_entry(story_id, agent)
+                        } else {
+                            // Agent was removed from map (e.g. stop_agent removes it after
+                            // the "stopped" status event is sent).
+                            let (status, session_id) = match event {
+                                AgentEvent::Done { session_id, .. } => {
+                                    (AgentStatus::Completed, session_id)
+                                }
+                                _ => (AgentStatus::Failed, None),
+                            };
+                            AgentInfo {
+                                story_id: story_id.to_string(),
+                                agent_name: agent_name.to_string(),
+                                status,
+                                session_id,
+                                worktree_path: None,
+                                base_branch: None,
+                                completion: None,
+                                log_session_id: None,
+                                throttled: false,
+                            }
+                        });
+                    }
+                }
+                Ok(Err(broadcast::error::RecvError::Lagged(_))) => {
+                    // Missed some buffered events — check current status before resuming.
+                    let agents = self.agents.lock().map_err(|e| e.to_string())?;
+                    let key = composite_key(story_id, agent_name);
+                    if let Some(agent) = agents.get(&key)
+                        && matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
+                    {
+                        return Ok(agent_info_from_entry(story_id, agent));
+                    }
+                    // Still running — continue the loop.
+                }
+                Ok(Err(broadcast::error::RecvError::Closed)) => {
+                    // Channel closed: no more events will arrive. Return current state.
+                    let agents = self.agents.lock().map_err(|e| e.to_string())?;
+                    let key = composite_key(story_id, agent_name);
+                    if let Some(agent) = agents.get(&key) {
+                        return Ok(agent_info_from_entry(story_id, agent));
+                    }
+                    return Err(format!(
+                        "Agent '{agent_name}' for story '{story_id}' channel closed unexpectedly"
+                    ));
+                }
+                Err(_) => {
+                    return Err(format!(
+                        "Timed out after {timeout_ms}ms waiting for agent '{agent_name}' on story '{story_id}'"
+                    ));
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::AgentPool;
+    use crate::agents::{AgentEvent, AgentStatus};
+
+    #[tokio::test]
+    async fn wait_for_agent_returns_immediately_if_completed() {
+        let pool = AgentPool::new_test(3001);
+        pool.inject_test_agent("s1", "bot", AgentStatus::Completed);
+
+        let info = pool.wait_for_agent("s1", "bot", 1000).await.unwrap();
+        assert_eq!(info.status, AgentStatus::Completed);
+        assert_eq!(info.story_id, "s1");
+        assert_eq!(info.agent_name, "bot");
+    }
+
+    #[tokio::test]
+    async fn wait_for_agent_returns_immediately_if_failed() {
+        let pool = AgentPool::new_test(3001);
+        pool.inject_test_agent("s2", "bot", AgentStatus::Failed);
+
+        let info = pool.wait_for_agent("s2", "bot", 1000).await.unwrap();
+        assert_eq!(info.status, AgentStatus::Failed);
+    }
+
+    #[tokio::test]
+    async fn wait_for_agent_completes_on_done_event() {
+        let pool = AgentPool::new_test(3001);
+        let tx = pool.inject_test_agent("s3", "bot", AgentStatus::Running);
+
+        // Send Done event after a short delay
+        let tx_clone = tx.clone();
+        tokio::spawn(async move {
+            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+            let _ = tx_clone.send(AgentEvent::Done {
+                story_id: "s3".to_string(),
+                agent_name: "bot".to_string(),
+                session_id: Some("sess-abc".to_string()),
+            });
+        });
+
+        let info = pool.wait_for_agent("s3", "bot", 2000).await.unwrap();
+        assert_eq!(info.story_id, "s3");
+    }
+
+    #[tokio::test]
+    async fn wait_for_agent_times_out() {
+        let pool = AgentPool::new_test(3001);
+        pool.inject_test_agent("s4", "bot", AgentStatus::Running);
+
+        let result = pool.wait_for_agent("s4", "bot", 50).await;
+        assert!(result.is_err());
+        let msg = result.unwrap_err();
+        assert!(msg.contains("Timed out"), "unexpected message: {msg}");
+    }
+
+    #[tokio::test]
+    async fn wait_for_agent_errors_for_nonexistent() {
+        let pool = AgentPool::new_test(3001);
+        let result = pool.wait_for_agent("no_story", "no_bot", 100).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn wait_for_agent_completes_on_stopped_status_event() {
+        let pool = AgentPool::new_test(3001);
+        let tx = pool.inject_test_agent("s5", "bot", AgentStatus::Running);
+
+        let tx_clone = tx.clone();
+        tokio::spawn(async move {
+            tokio::time::sleep(std::time::Duration::from_millis(30)).await;
+            let _ = tx_clone.send(AgentEvent::Status {
+                story_id: "s5".to_string(),
+                agent_name: "bot".to_string(),
+                status: "stopped".to_string(),
+            });
+        });
+
+        let info = pool.wait_for_agent("s5", "bot", 2000).await.unwrap();
+        assert_eq!(info.story_id, "s5");
+    }
+}