Genarrative/server-rs/crates/agent-runtime-orchestration/src/plan.rs

use serde::{Deserialize, Serialize};

use crate::{OrchestrationError, OrchestrationErrorKind, TaskGraph};

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PlanSelection {
    All,
    Repair { task_ids: Vec<String> },
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct OrchestrationPlan {
    active_task_ids: Vec<String>,
    carried_task_ids: Vec<String>,
    dependency_waves: Vec<Vec<String>>,
}

impl OrchestrationPlan {
    pub fn active_task_ids(&self) -> &[String] {
        &self.active_task_ids
    }

    pub fn carried_task_ids(&self) -> &[String] {
        &self.carried_task_ids
    }

    pub fn dependency_waves(&self) -> &[Vec<String>] {
        &self.dependency_waves
    }
}

impl TaskGraph {
    pub fn plan(&self, selection: PlanSelection) -> Result<OrchestrationPlan, OrchestrationError> {
        let all_task_ids = self.all_task_ids();
        let active_task_ids = match selection {
            PlanSelection::All => all_task_ids.clone(),
            PlanSelection::Repair { task_ids } => {
                if task_ids.is_empty() {
                    return Err(OrchestrationError::new(
                        OrchestrationErrorKind::InvalidInput,
                        "repair selection 至少需要一个 task",
                    ));
                }
                self.expand_downstream(&task_ids)?
            }
        };
        let active = active_task_ids
            .iter()
            .map(String::as_str)
            .collect::<std::collections::HashSet<_>>();
        let carried_task_ids = all_task_ids
            .into_iter()
            .filter(|task_id| !active.contains(task_id.as_str()))
            .collect::<Vec<_>>();
        let dependency_waves = self.dependency_waves(&active_task_ids, &carried_task_ids)?;

        Ok(OrchestrationPlan {
            active_task_ids,
            carried_task_ids,
            dependency_waves,
        })
    }
}