poimen-memory/crates/mem-cli/src/handlers/workflow_poller.rs

// Workflow Status Poller
// Encapsulates polling logic with retry + timeout (reduces CRAP)

use crate::agent::client_sdk::SynthesisClient;
use crate::handlers::workflow_builder::WorkflowQueryBuilder;
use serde_json::Value;
use tracing::{debug, warn};

/// Poll configuration
pub struct PollConfig {
    pub max_polls: u32,
    pub poll_delay_ms: u64,
    pub timeout_ms: u64,
}

impl Default for PollConfig {
    fn default() -> Self {
        Self {
            max_polls: 30,
            poll_delay_ms: 100,
            timeout_ms: 3000,
        }
    }
}

/// Parse workflow status from DESCRIBE_WORKFLOW response
pub fn parse_workflow_status(response: &Value) -> Option<String> {
    response
        .get("data")
        .and_then(|data| data.get("status"))
        .and_then(|status| status.as_str())
        .map(|s| s.to_string())
}

/// Parse error message from failed workflow
pub fn parse_workflow_error(response: &Value) -> Option<String> {
    response
        .get("data")
        .and_then(|data| data.get("error"))
        .and_then(|error| error.as_str())
        .map(|s| s.to_string())
}

/// Extract result from completed workflow
pub fn extract_workflow_result(response: &Value) -> Option<Value> {
    response
        .get("data")
        .and_then(|data| data.get("result"))
        .cloned()
}

/// Poll workflow until completion or timeout
/// Returns: (status, result_or_error)
pub async fn poll_workflow_until_complete(
    client: &SynthesisClient,
    workflow_id: &str,
    config: PollConfig,
) -> Result<(String, Option<Value>), String> {
    let start = std::time::Instant::now();
    let mut poll_count = 0;

    loop {
        // Check timeout
        if start.elapsed().as_millis() as u64 > config.timeout_ms {
            return Err("Workflow polling timeout".to_string());
        }

        // Check max polls
        if poll_count >= config.max_polls {
            return Err("Max workflow polls exceeded".to_string());
        }

        // Query workflow status
        let query_req = WorkflowQueryBuilder::new(workflow_id).describe();

        match client.execute_workflow(query_req).await {
            Ok(response) => {
                if let Some(status) = parse_workflow_status(&response) {
                    match status.as_str() {
                        "COMPLETED" => {
                            let result = extract_workflow_result(&response);
                            return Ok(("COMPLETED".to_string(), result));
                        }
                        "FAILED" => {
                            let error = parse_workflow_error(&response)
                                .unwrap_or_else(|| "Unknown error".to_string());
                            return Err(format!("Workflow failed: {}", error));
                        }
                        _ => {
                            debug!(
                                "Workflow status: {}, poll_count: {}/{}",
                                status, poll_count, config.max_polls
                            );
                        }
                    }
                } else {
                    warn!("Could not parse workflow status from response");
                }
            }
            Err(e) => {
                return Err(format!("Workflow query failed: {}", e));
            }
        }

        poll_count += 1;
        tokio::time::sleep(tokio::time::Duration::from_millis(config.poll_delay_ms)).await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_parse_workflow_status_success() {
        let response = json!({
            "data": {
                "status": "COMPLETED",
                "result": {"answer": "test"}
            }
        });

        let status = parse_workflow_status(&response);
        assert_eq!(status, Some("COMPLETED".to_string()));
    }

    #[test]
    fn test_parse_workflow_status_running() {
        let response = json!({
            "data": {
                "status": "RUNNING"
            }
        });

        let status = parse_workflow_status(&response);
        assert_eq!(status, Some("RUNNING".to_string()));
    }

    #[test]
    fn test_parse_workflow_status_missing() {
        let response = json!({"data": {}});
        let status = parse_workflow_status(&response);
        assert_eq!(status, None);
    }

    #[test]
    fn test_parse_workflow_error() {
        let response = json!({
            "data": {
                "error": "Activity failed: LLM timeout"
            }
        });

        let error = parse_workflow_error(&response);
        assert_eq!(error, Some("Activity failed: LLM timeout".to_string()));
    }

    #[test]
    fn test_extract_workflow_result() {
        let response = json!({
            "data": {
                "result": {
                    "question": "test",
                    "answers": ["answer1"]
                }
            }
        });

        let result = extract_workflow_result(&response);
        assert!(result.is_some());
        assert_eq!(result.unwrap()["question"], "test");
    }

    #[test]
    fn test_extract_workflow_result_none() {
        let response = json!({"data": {}});
        let result = extract_workflow_result(&response);
        assert_eq!(result, None);
    }

    #[test]
    fn test_poll_config_default() {
        let config = PollConfig::default();
        assert_eq!(config.max_polls, 30);
        assert_eq!(config.poll_delay_ms, 100);
        assert_eq!(config.timeout_ms, 3000);
    }
}