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

use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque};

use agent_runtime_core::AgentCatalog;
use serde::{Deserialize, Deserializer, Serialize};

use crate::{OrchestrationError, OrchestrationErrorKind};

const IDENTIFIER_MAX_CHARS: usize = 128;
const GOAL_MAX_CHARS: usize = 4_000;

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum TaskStatus {
    Pending,
    Running,
    Waiting,
    Completed,
    Failed,
    Cancelled,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct TaskNode {
    id: String,
    agent_id: String,
    status: TaskStatus,
    dependencies: Vec<String>,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct TaskNodeInput {
    id: String,
    agent_id: String,
    status: TaskStatus,
    dependencies: Vec<String>,
}

impl<'de> Deserialize<'de> for TaskNode {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let input = TaskNodeInput::deserialize(deserializer)?;
        Self::try_new(input.id, input.agent_id, input.status, input.dependencies)
            .map_err(serde::de::Error::custom)
    }
}

impl TaskNode {
    pub fn try_new(
        id: impl Into<String>,
        agent_id: impl Into<String>,
        status: TaskStatus,
        dependencies: impl IntoIterator<Item = impl Into<String>>,
    ) -> Result<Self, OrchestrationError> {
        let id = id.into();
        let agent_id = agent_id.into();
        validate_identifier(&id, "task id")?;
        validate_identifier(&agent_id, "task agent id")?;

        let mut seen = BTreeSet::new();
        let mut dependencies_output = Vec::new();
        for dependency in dependencies {
            let dependency = dependency.into();
            validate_identifier(&dependency, "task dependency")?;
            if dependency == id {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::SelfDependency,
                    format!("task {id} 不能依赖自身"),
                ));
            }
            if !seen.insert(dependency.clone()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateDependency,
                    format!("task {id} 重复依赖：{dependency}"),
                ));
            }
            dependencies_output.push(dependency);
        }

        Ok(Self {
            id,
            agent_id,
            status,
            dependencies: dependencies_output,
        })
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn agent_id(&self) -> &str {
        &self.agent_id
    }

    pub fn status(&self) -> TaskStatus {
        self.status
    }

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

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct TaskGraph {
    goal: String,
    tasks: Vec<TaskNode>,
    #[serde(skip)]
    by_id: BTreeMap<String, usize>,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct TaskGraphInput {
    goal: String,
    tasks: Vec<TaskNode>,
}

impl<'de> Deserialize<'de> for TaskGraph {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let input = TaskGraphInput::deserialize(deserializer)?;
        Self::try_new(input.goal, input.tasks).map_err(serde::de::Error::custom)
    }
}

impl TaskGraph {
    pub fn try_new(
        goal: impl Into<String>,
        tasks: impl IntoIterator<Item = TaskNode>,
    ) -> Result<Self, OrchestrationError> {
        let goal = goal.into();
        validate_goal(&goal)?;
        let tasks = tasks.into_iter().collect::<Vec<_>>();
        if tasks.is_empty() {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::InvalidInput,
                "task graph 至少需要一个任务",
            ));
        }

        let mut by_id = BTreeMap::new();
        for (index, task) in tasks.iter().enumerate() {
            if by_id.insert(task.id.clone(), index).is_some() {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateTask,
                    format!("task id 重复：{}", task.id),
                ));
            }
        }
        for task in &tasks {
            for dependency in &task.dependencies {
                if !by_id.contains_key(dependency) {
                    return Err(OrchestrationError::new(
                        OrchestrationErrorKind::UnknownDependency,
                        format!("task {} 引用了未知依赖：{dependency}", task.id),
                    ));
                }
            }
        }

        validate_acyclic(&tasks, &by_id)?;
        Ok(Self { goal, tasks, by_id })
    }

    pub fn goal(&self) -> &str {
        &self.goal
    }

    pub fn tasks(&self) -> &[TaskNode] {
        &self.tasks
    }

    pub fn get(&self, task_id: &str) -> Option<&TaskNode> {
        self.by_id
            .get(task_id)
            .and_then(|index| self.tasks.get(*index))
    }

    /// Number of tasks in this graph.
    pub fn task_count(&self) -> usize {
        self.tasks.len()
    }

    /// Alias for [`TaskGraph::task_count`] using graph terminology.
    pub fn node_count(&self) -> usize {
        self.task_count()
    }

    /// Number of prerequisite edges in this graph.
    pub fn edge_count(&self) -> usize {
        self.tasks.iter().map(|task| task.dependencies.len()).sum()
    }

    /// Longest dependency path measured in task layers. A root task has
    /// depth 1. Graph construction rejects cycles, so this calculation is
    /// total for every `TaskGraph` value.
    pub fn depth(&self) -> usize {
        graph_depth(&self.tasks, &self.by_id)
    }

    /// Number of direct dependents of a prerequisite task.
    pub fn fan_out(&self, task_id: &str) -> Option<usize> {
        self.get(task_id)?;
        Some(
            self.tasks
                .iter()
                .filter(|task| task.dependencies.iter().any(|id| id == task_id))
                .count(),
        )
    }

    /// Alias for [`TaskGraph::fan_out`].
    pub fn out_degree(&self, task_id: &str) -> Option<usize> {
        self.fan_out(task_id)
    }

    pub fn validate_agents(&self, catalog: &AgentCatalog) -> Result<(), OrchestrationError> {
        for task in &self.tasks {
            if catalog.get(&task.agent_id).is_none() {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::UnknownAgent,
                    format!("task {} 引用了未注册 Agent：{}", task.id, task.agent_id),
                ));
            }
        }
        Ok(())
    }

    pub fn ready_task_ids(&self) -> Vec<&str> {
        let completed = self
            .tasks
            .iter()
            .filter(|task| task.status == TaskStatus::Completed)
            .map(|task| task.id.as_str())
            .collect::<HashSet<_>>();

        self.tasks
            .iter()
            .filter(|task| {
                task.status == TaskStatus::Pending
                    && task
                        .dependencies
                        .iter()
                        .all(|dependency| completed.contains(dependency.as_str()))
            })
            .map(|task| task.id.as_str())
            .collect()
    }

    pub fn expand_downstream<T: AsRef<str>>(
        &self,
        task_ids: &[T],
    ) -> Result<Vec<String>, OrchestrationError> {
        let seeds = self.collect_known_task_ids(task_ids, "downstream seeds")?;
        if seeds.is_empty() {
            return Ok(Vec::new());
        }
        let mut impacted = seeds;
        let mut changed = true;
        while changed {
            changed = false;
            for task in &self.tasks {
                if impacted.contains(&task.id) {
                    continue;
                }
                if task
                    .dependencies
                    .iter()
                    .any(|dependency| impacted.contains(dependency))
                {
                    impacted.insert(task.id.clone());
                    changed = true;
                }
            }
        }

        Ok(self
            .tasks
            .iter()
            .filter(|task| impacted.contains(&task.id))
            .map(|task| task.id.clone())
            .collect())
    }

    pub fn dependency_waves<A: AsRef<str>, S: AsRef<str>>(
        &self,
        active_task_ids: &[A],
        satisfied_task_ids: &[S],
    ) -> Result<Vec<Vec<String>>, OrchestrationError> {
        let active = self.collect_known_task_ids(active_task_ids, "active tasks")?;
        let satisfied = self.collect_known_task_ids(satisfied_task_ids, "satisfied tasks")?;
        if let Some(task_id) = active.iter().find(|task_id| satisfied.contains(*task_id)) {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::ConflictingTaskSet,
                format!("task 同时位于 active 与 satisfied：{task_id}"),
            ));
        }

        for task in self.tasks.iter().filter(|task| active.contains(&task.id)) {
            for dependency in &task.dependencies {
                if !active.contains(dependency) && !satisfied.contains(dependency) {
                    return Err(OrchestrationError::new(
                        OrchestrationErrorKind::UnsatisfiedDependency,
                        format!(
                            "active task {} 的依赖既未 active 也未 satisfied：{dependency}",
                            task.id
                        ),
                    ));
                }
            }
        }

        let mut remaining = self
            .tasks
            .iter()
            .filter(|task| active.contains(&task.id))
            .map(|task| task.id.clone())
            .collect::<Vec<_>>();
        let mut completed = satisfied;
        let mut waves = Vec::new();
        while !remaining.is_empty() {
            let wave = remaining
                .iter()
                .filter(|task_id| {
                    self.get(task_id).is_some_and(|task| {
                        task.dependencies
                            .iter()
                            .all(|dependency| completed.contains(dependency))
                    })
                })
                .cloned()
                .collect::<Vec<_>>();
            if wave.is_empty() {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::Cycle,
                    format!(
                        "active task graph 无法生成下一依赖波次：{}",
                        remaining.join(", ")
                    ),
                ));
            }
            for task_id in &wave {
                completed.insert(task_id.clone());
            }
            remaining.retain(|task_id| !completed.contains(task_id));
            waves.push(wave);
        }
        Ok(waves)
    }

    pub(crate) fn all_task_ids(&self) -> Vec<String> {
        self.tasks.iter().map(|task| task.id.clone()).collect()
    }

    fn collect_known_task_ids<T: AsRef<str>>(
        &self,
        task_ids: &[T],
        label: &str,
    ) -> Result<BTreeSet<String>, OrchestrationError> {
        let mut output = BTreeSet::new();
        for task_id in task_ids {
            let task_id = task_id.as_ref();
            if self.get(task_id).is_none() {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::UnknownTask,
                    format!("{label} 包含未知 task：{task_id}"),
                ));
            }
            if !output.insert(task_id.to_string()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateTask,
                    format!("{label} 包含重复 task：{task_id}"),
                ));
            }
        }
        Ok(output)
    }
}

pub(crate) fn validate_identifier(value: &str, field: &str) -> Result<(), OrchestrationError> {
    if value != value.trim() {
        return Err(OrchestrationError::new(
            OrchestrationErrorKind::InvalidInput,
            format!("{field} 不得包含首尾空白"),
        ));
    }
    let mut chars = value.chars();
    let first = chars.next().ok_or_else(|| {
        OrchestrationError::new(
            OrchestrationErrorKind::InvalidInput,
            format!("{field} 不能为空"),
        )
    })?;
    if value.chars().count() > IDENTIFIER_MAX_CHARS
        || !first.is_ascii_alphanumeric()
        || !chars.all(|character| {
            character.is_ascii_alphanumeric() || matches!(character, '.' | '_' | '-' | ':')
        })
    {
        return Err(OrchestrationError::new(
            OrchestrationErrorKind::InvalidInput,
            format!("{field} 不是合法稳定标识：{value}"),
        ));
    }
    Ok(())
}

fn validate_goal(goal: &str) -> Result<(), OrchestrationError> {
    if goal != goal.trim() || goal.is_empty() || goal.chars().count() > GOAL_MAX_CHARS {
        return Err(OrchestrationError::new(
            OrchestrationErrorKind::InvalidInput,
            format!("task graph goal 必须为 1..={GOAL_MAX_CHARS} 个无首尾空白字符"),
        ));
    }
    if goal.chars().any(char::is_control) {
        return Err(OrchestrationError::new(
            OrchestrationErrorKind::InvalidInput,
            "task graph goal 不能包含控制字符",
        ));
    }
    Ok(())
}

fn validate_acyclic(
    tasks: &[TaskNode],
    by_id: &BTreeMap<String, usize>,
) -> Result<(), OrchestrationError> {
    let mut indegrees = tasks
        .iter()
        .map(|task| task.dependencies.len())
        .collect::<Vec<_>>();
    let mut dependents = vec![Vec::<usize>::new(); tasks.len()];
    for (task_index, task) in tasks.iter().enumerate() {
        for dependency in &task.dependencies {
            let dependency_index = by_id[dependency];
            dependents[dependency_index].push(task_index);
        }
    }

    let mut ready = indegrees
        .iter()
        .enumerate()
        .filter_map(|(index, indegree)| (*indegree == 0).then_some(index))
        .collect::<VecDeque<_>>();
    let mut visited = 0;
    while let Some(index) = ready.pop_front() {
        visited += 1;
        for dependent in &dependents[index] {
            indegrees[*dependent] -= 1;
            if indegrees[*dependent] == 0 {
                ready.push_back(*dependent);
            }
        }
    }
    if visited == tasks.len() {
        return Ok(());
    }

    let cyclic = tasks
        .iter()
        .zip(indegrees)
        .filter_map(|(task, indegree)| (indegree > 0).then_some(task.id.as_str()))
        .collect::<Vec<_>>();
    Err(OrchestrationError::new(
        OrchestrationErrorKind::Cycle,
        format!("task graph 包含依赖环：{}", cyclic.join(", ")),
    ))
}

fn graph_depth(tasks: &[TaskNode], by_id: &BTreeMap<String, usize>) -> usize {
    if tasks.is_empty() {
        return 0;
    }

    let mut indegrees = tasks
        .iter()
        .map(|task| task.dependencies.len())
        .collect::<Vec<_>>();
    let mut dependents = vec![Vec::<usize>::new(); tasks.len()];
    for (task_index, task) in tasks.iter().enumerate() {
        for dependency in &task.dependencies {
            // `TaskGraph::try_new` proves this lookup exists. Keeping the
            // defensive branch makes this helper total if it is ever reused
            // during a future internal refactor.
            let Some(&dependency_index) = by_id.get(dependency) else {
                return 0;
            };
            dependents[dependency_index].push(task_index);
        }
    }

    let mut depths = vec![1usize; tasks.len()];
    let mut ready = indegrees
        .iter()
        .enumerate()
        .filter_map(|(index, indegree)| (*indegree == 0).then_some(index))
        .collect::<VecDeque<_>>();
    let mut visited = 0;
    let mut maximum = 1;
    while let Some(index) = ready.pop_front() {
        visited += 1;
        maximum = maximum.max(depths[index]);
        for dependent in &dependents[index] {
            depths[*dependent] = depths[*dependent].max(depths[index].saturating_add(1));
            indegrees[*dependent] -= 1;
            if indegrees[*dependent] == 0 {
                ready.push_back(*dependent);
            }
        }
    }
    if visited == tasks.len() { maximum } else { 0 }
}