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

//! Structured, host-agnostic graph expansion proposed by an LLM or another
//! planner.
//!
//! This module deliberately stops at validation and candidate construction.
//! It does not call a provider, execute an agent, or persist an epoch. A host
//! can deserialize a provider function-call argument into [`GraphProposal`],
//! pass it to [`TaskGraph::expand_with_proposal`], and persist the returned
//! graph as the next epoch if the result is accepted.

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

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

use crate::{
    OrchestrationError, OrchestrationErrorKind, TaskGraph, TaskNode, TaskStatus,
    graph::validate_identifier,
};

/// Default maximum number of tasks in a candidate graph.
pub const DEFAULT_GRAPH_MAX_TASKS: usize = 128;
/// Default maximum number of dependency edges in a candidate graph.
pub const DEFAULT_GRAPH_MAX_EDGES: usize = 512;
/// Default maximum number of task layers in a candidate graph.
pub const DEFAULT_GRAPH_MAX_DEPTH: usize = 32;
/// Default maximum number of direct dependents of one task.
pub const DEFAULT_GRAPH_MAX_OUT_DEGREE: usize = 32;

/// A task that a planner proposes to add to a graph.
///
/// New tasks are always inserted with [`TaskStatus::Pending`]. Product
/// metadata such as a title, artifact path, or acceptance text belongs in the
/// host adapter and is intentionally not part of this generic DTO.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct TaskProposal {
    id: String,
    agent_id: String,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct TaskProposalInput {
    id: String,
    agent_id: String,
}

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

impl TaskProposal {
    /// Creates a validated task proposal.
    pub fn try_new(
        id: impl Into<String>,
        agent_id: impl Into<String>,
    ) -> Result<Self, OrchestrationError> {
        let id = id.into();
        let agent_id = agent_id.into();
        validate_identifier(&id, "proposal task id")?;
        validate_identifier(&agent_id, "proposal task agent id")?;
        Ok(Self { id, agent_id })
    }

    /// Alias for [`TaskProposal::try_new`] for hosts that use `new` for DTO
    /// construction while still handling validation errors.
    pub fn new(
        id: impl Into<String>,
        agent_id: impl Into<String>,
    ) -> Result<Self, OrchestrationError> {
        Self::try_new(id, agent_id)
    }

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

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

/// A directed dependency edge in a proposal.
///
/// `from` is the prerequisite and `to` is the task that depends on it. The
/// edge therefore corresponds to adding `from` to `to.dependencies` in the
/// resulting graph.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct GraphEdge {
    from: String,
    to: String,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct GraphEdgeInput {
    #[serde(alias = "source")]
    from: String,
    #[serde(alias = "target")]
    to: String,
}

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

impl GraphEdge {
    /// Creates a validated prerequisite-to-dependent edge.
    pub fn try_new(
        from: impl Into<String>,
        to: impl Into<String>,
    ) -> Result<Self, OrchestrationError> {
        let from = from.into();
        let to = to.into();
        validate_identifier(&from, "proposal edge from")?;
        validate_identifier(&to, "proposal edge to")?;
        if from == to {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::SelfDependency,
                format!("proposal edge 不能连接任务自身：{from}"),
            ));
        }
        Ok(Self { from, to })
    }

    /// Alias for [`GraphEdge::try_new`].
    pub fn new(from: impl Into<String>, to: impl Into<String>) -> Result<Self, OrchestrationError> {
        Self::try_new(from, to)
    }

    /// Convenience constructor whose names make the dependency direction
    /// explicit at call sites.
    pub fn dependency(
        prerequisite: impl Into<String>,
        dependent: impl Into<String>,
    ) -> Result<Self, OrchestrationError> {
        Self::try_new(prerequisite, dependent)
    }

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

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

    /// Alias for [`GraphEdge::from`], useful when a host calls the fields
    /// source/target in its own graph model.
    pub fn source(&self) -> &str {
        &self.from
    }

    /// Alias for [`GraphEdge::to`].
    pub fn target(&self) -> &str {
        &self.to
    }
}

/// A structured graph change returned by a planner.
///
/// The proposal contains only additions. Edges whose target is an existing
/// task are rejected so a running task never acquires a new prerequisite in
/// place. To replace existing dependencies, a host must build a complete
/// candidate graph and install it as a new epoch with its own persistence/CAS
/// contract.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct GraphProposal {
    nodes: Vec<TaskProposal>,
    edges: Vec<GraphEdge>,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct GraphProposalInput {
    nodes: Vec<TaskProposal>,
    edges: Vec<GraphEdge>,
}

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

impl GraphProposal {
    pub fn try_new(
        nodes: impl IntoIterator<Item = TaskProposal>,
        edges: impl IntoIterator<Item = GraphEdge>,
    ) -> Result<Self, OrchestrationError> {
        let nodes = nodes.into_iter().collect::<Vec<_>>();
        let edges = edges.into_iter().collect::<Vec<_>>();
        if nodes.is_empty() && edges.is_empty() {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::EmptyProposal,
                "graph proposal 至少需要一个新节点或一条新边",
            ));
        }

        let mut node_ids = BTreeSet::new();
        for node in &nodes {
            if !node_ids.insert(node.id.clone()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateTask,
                    format!("proposal task id 重复：{}", node.id),
                ));
            }
        }

        let mut edge_ids = BTreeSet::new();
        for edge in &edges {
            if !edge_ids.insert((edge.from.clone(), edge.to.clone())) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateEdge,
                    format!("proposal edge 重复：{} -> {}", edge.from, edge.to),
                ));
            }
        }

        Ok(Self { nodes, edges })
    }

    pub fn new(
        nodes: impl IntoIterator<Item = TaskProposal>,
        edges: impl IntoIterator<Item = GraphEdge>,
    ) -> Result<Self, OrchestrationError> {
        Self::try_new(nodes, edges)
    }

    pub fn nodes(&self) -> &[TaskProposal] {
        &self.nodes
    }

    /// Alias for [`GraphProposal::nodes`] for callers that use task-oriented
    /// terminology.
    pub fn tasks(&self) -> &[TaskProposal] {
        &self.nodes
    }

    pub fn edges(&self) -> &[GraphEdge] {
        &self.edges
    }

    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty() && self.edges.is_empty()
    }

    pub fn validate(&self) -> Result<(), OrchestrationError> {
        // The fields are private and constructors/deserialization already
        // enforce these invariants. Re-running the cheap checks keeps this
        // method useful as an explicit boundary for host adapters.
        if self.is_empty() {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::EmptyProposal,
                "graph proposal 至少需要一个新节点或一条新边",
            ));
        }
        let mut node_ids = BTreeSet::new();
        for node in &self.nodes {
            validate_identifier(&node.id, "proposal task id")?;
            validate_identifier(&node.agent_id, "proposal task agent id")?;
            if !node_ids.insert(node.id.as_str()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateTask,
                    format!("proposal task id 重复：{}", node.id),
                ));
            }
        }
        let mut edge_ids = BTreeSet::new();
        for edge in &self.edges {
            validate_identifier(&edge.from, "proposal edge from")?;
            validate_identifier(&edge.to, "proposal edge to")?;
            if edge.from == edge.to {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::SelfDependency,
                    format!("proposal edge 不能连接任务自身：{}", edge.from),
                ));
            }
            if !edge_ids.insert((edge.from.as_str(), edge.to.as_str())) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateEdge,
                    format!("proposal edge 重复：{} -> {}", edge.from, edge.to),
                ));
            }
        }
        Ok(())
    }
}

/// Resource limits applied to the candidate graph produced by a proposal.
///
/// Limits are checked against the complete resulting graph, not just the
/// proposed delta. `max_depth` counts graph layers: a root task has depth 1.
/// `max_out_degree` counts dependents for one prerequisite (`from -> to`).
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct GraphLimits {
    pub max_tasks: usize,
    pub max_edges: usize,
    pub max_depth: usize,
    pub max_out_degree: usize,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
struct GraphLimitsInput {
    #[serde(alias = "maxNodes")]
    max_tasks: usize,
    max_edges: usize,
    max_depth: usize,
    #[serde(alias = "maxFanOut")]
    max_out_degree: usize,
}

impl<'de> Deserialize<'de> for GraphLimits {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let input = GraphLimitsInput::deserialize(deserializer)?;
        Self::try_new(
            input.max_tasks,
            input.max_edges,
            input.max_depth,
            input.max_out_degree,
        )
        .map_err(serde::de::Error::custom)
    }
}

impl Default for GraphLimits {
    fn default() -> Self {
        Self {
            max_tasks: DEFAULT_GRAPH_MAX_TASKS,
            max_edges: DEFAULT_GRAPH_MAX_EDGES,
            max_depth: DEFAULT_GRAPH_MAX_DEPTH,
            max_out_degree: DEFAULT_GRAPH_MAX_OUT_DEGREE,
        }
    }
}

impl GraphLimits {
    pub const fn new(
        max_tasks: usize,
        max_edges: usize,
        max_depth: usize,
        max_out_degree: usize,
    ) -> Self {
        Self {
            max_tasks,
            max_edges,
            max_depth,
            max_out_degree,
        }
    }

    pub fn try_new(
        max_tasks: usize,
        max_edges: usize,
        max_depth: usize,
        max_out_degree: usize,
    ) -> Result<Self, OrchestrationError> {
        let limits = Self::new(max_tasks, max_edges, max_depth, max_out_degree);
        limits.validate().map(|()| limits)
    }

    pub fn with_max_tasks(mut self, value: usize) -> Self {
        self.max_tasks = value;
        self
    }

    pub fn with_max_nodes(self, value: usize) -> Self {
        self.with_max_tasks(value)
    }

    pub fn with_max_edges(mut self, value: usize) -> Self {
        self.max_edges = value;
        self
    }

    pub fn with_max_depth(mut self, value: usize) -> Self {
        self.max_depth = value;
        self
    }

    pub fn with_max_out_degree(mut self, value: usize) -> Self {
        self.max_out_degree = value;
        self
    }

    pub fn with_max_fan_out(self, value: usize) -> Self {
        self.with_max_out_degree(value)
    }

    pub fn max_nodes(&self) -> usize {
        self.max_tasks
    }

    pub fn max_fan_out(&self) -> usize {
        self.max_out_degree
    }

    pub fn validate(&self) -> Result<(), OrchestrationError> {
        let invalid = [
            (self.max_tasks, "maxTasks"),
            (self.max_edges, "maxEdges"),
            (self.max_depth, "maxDepth"),
            (self.max_out_degree, "maxOutDegree"),
        ]
        .into_iter()
        .find(|(value, _)| *value == 0);
        if let Some((_, field)) = invalid {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::InvalidLimits,
                format!("graph limits 的 {field} 必须大于 0"),
            ));
        }
        Ok(())
    }
}

/// A validated candidate graph plus the delta that produced it.
///
/// The host may use this report to persist an epoch/change journal without
/// re-parsing the untrusted provider payload. The graph itself remains the
/// authoritative candidate.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GraphExpansion {
    graph: TaskGraph,
    added_task_ids: Vec<String>,
    added_edges: Vec<GraphEdge>,
}

/// Descriptive alias for [`GraphExpansion`].
pub type AppliedGraphProposal = GraphExpansion;

impl GraphExpansion {
    pub fn graph(&self) -> &TaskGraph {
        &self.graph
    }

    pub fn into_graph(self) -> TaskGraph {
        self.graph
    }

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

    pub fn added_edges(&self) -> &[GraphEdge] {
        &self.added_edges
    }
}

impl TaskGraph {
    /// Applies an additive proposal and returns a new validated graph.
    ///
    /// This method is intentionally immutable: a successful return is a
    /// candidate for a new host-managed epoch, while every error leaves the
    /// current graph untouched. New edges must target a newly proposed task;
    /// this prevents changing the prerequisites of a task that may already be
    /// running or completed.
    pub fn apply_proposal(
        &self,
        proposal: &GraphProposal,
        catalog: &AgentCatalog,
        limits: &GraphLimits,
    ) -> Result<Self, OrchestrationError> {
        self.expand_with_proposal(proposal, catalog, limits)
            .map(GraphExpansion::into_graph)
    }

    /// Returns the candidate graph together with its validated additive delta.
    pub fn expand_with_proposal(
        &self,
        proposal: &GraphProposal,
        catalog: &AgentCatalog,
        limits: &GraphLimits,
    ) -> Result<GraphExpansion, OrchestrationError> {
        limits.validate()?;
        proposal.validate()?;
        self.validate_agents(catalog)?;

        let existing_task_count = self.tasks().len();
        if existing_task_count > limits.max_tasks {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::NodeBudgetExceeded,
                format!(
                    "现有 task 数量 {} 已超过 maxTasks {}",
                    existing_task_count, limits.max_tasks
                ),
            ));
        }
        let resulting_task_count = existing_task_count
            .checked_add(proposal.nodes.len())
            .ok_or_else(|| {
                OrchestrationError::new(
                    OrchestrationErrorKind::NodeBudgetExceeded,
                    "proposal task 数量计算溢出",
                )
            })?;
        if resulting_task_count > limits.max_tasks {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::NodeBudgetExceeded,
                format!(
                    "扩图后 task 数量 {} 超过 maxTasks {}",
                    resulting_task_count, limits.max_tasks
                ),
            ));
        }

        let existing_ids = self
            .tasks()
            .iter()
            .map(|task| task.id().to_string())
            .collect::<BTreeSet<_>>();
        let proposed_ids = proposal
            .nodes
            .iter()
            .map(|node| node.id.clone())
            .collect::<BTreeSet<_>>();
        for node in &proposal.nodes {
            if existing_ids.contains(&node.id) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateTask,
                    format!("proposal task 已存在于当前 graph：{}", node.id),
                ));
            }
            if catalog.get(&node.agent_id).is_none() {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::UnknownAgent,
                    format!(
                        "proposal task {} 引用了未注册 Agent：{}",
                        node.id, node.agent_id
                    ),
                ));
            }
        }

        let all_ids = existing_ids
            .iter()
            .chain(proposed_ids.iter())
            .cloned()
            .collect::<BTreeSet<_>>();
        let existing_edges = dependency_edges(self);
        let mut proposed_edges = BTreeSet::new();
        let mut dependencies_by_target = BTreeMap::<String, Vec<String>>::new();
        for edge in &proposal.edges {
            if !all_ids.contains(edge.from()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::UnknownDependency,
                    format!(
                        "proposal edge {} -> {} 引用了未知依赖：{}",
                        edge.from(),
                        edge.to(),
                        edge.from()
                    ),
                ));
            }
            if !all_ids.contains(edge.to()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::UnknownTask,
                    format!(
                        "proposal edge {} -> {} 引用了未知目标 task：{}",
                        edge.from(),
                        edge.to(),
                        edge.to()
                    ),
                ));
            }
            if !proposed_ids.contains(edge.to()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::ExistingTaskMutation,
                    format!(
                        "proposal edge {} -> {} 不能修改已有 task 的依赖",
                        edge.from(),
                        edge.to()
                    ),
                ));
            }
            let edge_key = (edge.from().to_string(), edge.to().to_string());
            if !proposed_edges.insert(edge_key.clone()) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateEdge,
                    format!("proposal edge 重复：{} -> {}", edge.from(), edge.to()),
                ));
            }
            if existing_edges.contains(&edge_key) {
                return Err(OrchestrationError::new(
                    OrchestrationErrorKind::DuplicateEdge,
                    format!("proposal edge 已存在：{} -> {}", edge.from(), edge.to()),
                ));
            }
            dependencies_by_target
                .entry(edge.to().to_string())
                .or_default()
                .push(edge.from().to_string());
        }

        let resulting_edge_count = self
            .edge_count()
            .checked_add(proposal.edges.len())
            .ok_or_else(|| {
                OrchestrationError::new(
                    OrchestrationErrorKind::EdgeBudgetExceeded,
                    "proposal edge 数量计算溢出",
                )
            })?;
        if resulting_edge_count > limits.max_edges {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::EdgeBudgetExceeded,
                format!(
                    "扩图后 dependency edge 数量 {} 超过 maxEdges {}",
                    resulting_edge_count, limits.max_edges
                ),
            ));
        }

        validate_out_degree(self, &proposal.edges, limits.max_out_degree)?;

        let mut tasks = self.tasks().to_vec();
        let added_task_ids = proposal
            .nodes
            .iter()
            .map(|node| node.id.clone())
            .collect::<Vec<_>>();
        for node in &proposal.nodes {
            let dependencies = dependencies_by_target.remove(&node.id).unwrap_or_default();
            tasks.push(TaskNode::try_new(
                node.id.clone(),
                node.agent_id.clone(),
                TaskStatus::Pending,
                dependencies,
            )?);
        }

        // TaskGraph::try_new performs the final unknown-dependency and cycle
        // checks over the complete candidate, so no partially built graph can
        // escape this method.
        let candidate = Self::try_new(self.goal().to_string(), tasks)?;
        if candidate.depth() > limits.max_depth {
            return Err(OrchestrationError::new(
                OrchestrationErrorKind::DepthBudgetExceeded,
                format!(
                    "扩图后 graph depth {} 超过 maxDepth {}",
                    candidate.depth(),
                    limits.max_depth
                ),
            ));
        }

        Ok(GraphExpansion {
            graph: candidate,
            added_task_ids,
            added_edges: proposal.edges.clone(),
        })
    }

    /// Parameter-order variant for hosts that keep limits before the catalog.
    pub fn apply_proposal_with_limits(
        &self,
        proposal: &GraphProposal,
        limits: &GraphLimits,
        catalog: &AgentCatalog,
    ) -> Result<Self, OrchestrationError> {
        self.apply_proposal(proposal, catalog, limits)
    }

    /// Short alias for [`TaskGraph::apply_proposal`].
    pub fn expand(
        &self,
        proposal: &GraphProposal,
        catalog: &AgentCatalog,
        limits: &GraphLimits,
    ) -> Result<Self, OrchestrationError> {
        self.apply_proposal(proposal, catalog, limits)
    }
}

fn dependency_edges(graph: &TaskGraph) -> BTreeSet<(String, String)> {
    graph
        .tasks()
        .iter()
        .flat_map(|task| {
            task.dependencies()
                .iter()
                .map(|dependency| (dependency.clone(), task.id().to_string()))
        })
        .collect()
}

fn validate_out_degree(
    graph: &TaskGraph,
    proposed_edges: &[GraphEdge],
    max_out_degree: usize,
) -> Result<(), OrchestrationError> {
    let mut out_degree = BTreeMap::<String, usize>::new();
    for (from, _) in dependency_edges(graph) {
        let count = out_degree.entry(from.clone()).or_default();
        *count = count.checked_add(1).ok_or_else(|| {
            OrchestrationError::new(
                OrchestrationErrorKind::FanOutBudgetExceeded,
                format!("task {from} 的 fan-out 数量计算溢出"),
            )
        })?;
    }
    for edge in proposed_edges {
        let count = out_degree.entry(edge.from().to_string()).or_default();
        *count = count.checked_add(1).ok_or_else(|| {
            OrchestrationError::new(
                OrchestrationErrorKind::FanOutBudgetExceeded,
                format!("task {} 的 fan-out 数量计算溢出", edge.from()),
            )
        })?;
    }
    if let Some((task_id, count)) = out_degree
        .iter()
        .find(|(_, count)| **count > max_out_degree)
    {
        return Err(OrchestrationError::new(
            OrchestrationErrorKind::FanOutBudgetExceeded,
            format!("task {task_id} 的 fan-out {count} 超过 maxOutDegree {max_out_degree}"),
        ));
    }
    Ok(())
}