Genarrative/server-rs/crates/agent-runtime-orchestration/tests/dynamic_proposal.rs

use agent_runtime_core::{AgentCatalog, AgentDescriptor};
use agent_runtime_orchestration::{
    GraphEdge, GraphLimits, GraphProposal, OrchestrationErrorKind, PlanSelection, TaskGraph,
    TaskNode, TaskProposal, TaskStatus,
};

fn task(id: &str, agent_id: &str, status: TaskStatus, dependencies: &[&str]) -> TaskNode {
    TaskNode::try_new(id, agent_id, status, dependencies.iter().copied()).expect("valid task")
}

fn catalog() -> AgentCatalog {
    AgentCatalog::try_new([
        AgentDescriptor::try_new("researcher", "research", std::iter::empty::<&str>())
            .expect("researcher"),
        AgentDescriptor::try_new("reviewer", "review", std::iter::empty::<&str>())
            .expect("reviewer"),
        AgentDescriptor::try_new("writer", "writing", std::iter::empty::<&str>()).expect("writer"),
    ])
    .expect("catalog")
}

fn base_graph() -> TaskGraph {
    TaskGraph::try_new(
        "Review a document collection",
        [
            task("collect", "researcher", TaskStatus::Completed, &[]),
            task("draft", "writer", TaskStatus::Pending, &["collect"]),
        ],
    )
    .expect("base graph")
}

fn proposal(nodes: &[(&str, &str)], edges: &[(&str, &str)]) -> GraphProposal {
    GraphProposal::try_new(
        nodes
            .iter()
            .map(|(id, agent)| TaskProposal::try_new(*id, *agent).expect("valid proposal node")),
        edges
            .iter()
            .map(|(from, to)| GraphEdge::try_new(*from, *to).expect("valid proposal edge")),
    )
    .expect("valid proposal")
}

#[test]
fn llm_proposal_creates_a_new_pending_subgraph_without_mutating_the_old_graph() {
    let graph = base_graph();
    let candidate = graph
        .apply_proposal(
            &proposal(
                &[("review", "reviewer"), ("publish", "writer")],
                &[
                    ("collect", "review"),
                    ("draft", "publish"),
                    ("review", "publish"),
                ],
            ),
            &catalog(),
            &GraphLimits::default(),
        )
        .expect("proposal should be accepted");

    assert_eq!(graph.task_count(), 2);
    assert_eq!(graph.edge_count(), 1);
    assert_eq!(candidate.task_count(), 4);
    assert_eq!(candidate.edge_count(), 4);
    assert_eq!(
        candidate.get("review").expect("review task").status(),
        TaskStatus::Pending
    );
    assert_eq!(
        candidate.get("review").expect("review task").dependencies(),
        &["collect".to_string()]
    );
    assert_eq!(
        candidate
            .get("publish")
            .expect("publish task")
            .dependencies(),
        &["draft".to_string(), "review".to_string()]
    );
    assert_eq!(candidate.ready_task_ids(), vec!["draft", "review"]);
    let plan = candidate
        .plan(PlanSelection::All)
        .expect("expanded graph should produce dependency waves");
    assert_eq!(
        plan.dependency_waves(),
        &[
            vec!["collect".to_string()],
            vec!["draft".to_string(), "review".to_string()],
            vec!["publish".to_string()],
        ]
    );
}

#[test]
fn expansion_report_contains_only_the_validated_delta() {
    let graph = base_graph();
    let change = proposal(&[("review", "reviewer")], &[("collect", "review")]);
    let expansion = graph
        .expand_with_proposal(&change, &catalog(), &GraphLimits::default())
        .expect("proposal should be accepted");

    assert_eq!(expansion.added_task_ids(), ["review"]);
    assert_eq!(expansion.added_edges(), change.edges());
    assert_eq!(
        expansion.graph().get("review").map(TaskNode::id),
        Some("review")
    );
}

#[test]
fn proposal_rejects_unknown_agents_and_keeps_the_current_graph_intact() {
    let graph = base_graph();
    let error = graph
        .apply_proposal(
            &proposal(&[("review", "unknown-agent")], &[]),
            &catalog(),
            &GraphLimits::default(),
        )
        .expect_err("unknown agent");
    assert_eq!(error.kind(), OrchestrationErrorKind::UnknownAgent);
    assert_eq!(graph.task_count(), 2);
    assert!(graph.get("review").is_none());
}

#[test]
fn proposal_cycle_is_rejected_atomically() {
    let graph = base_graph();
    let error = graph
        .apply_proposal(
            &proposal(
                &[("left", "researcher"), ("right", "reviewer")],
                &[("left", "right"), ("right", "left")],
            ),
            &catalog(),
            &GraphLimits::default(),
        )
        .expect_err("cycle");
    assert_eq!(error.kind(), OrchestrationErrorKind::Cycle);
    assert_eq!(graph.task_count(), 2);
    assert!(graph.get("left").is_none());
}

#[test]
fn proposal_cannot_add_a_prerequisite_to_an_existing_task() {
    let graph = base_graph();
    let error = graph
        .apply_proposal(
            &proposal(&[("review", "reviewer")], &[("review", "draft")]),
            &catalog(),
            &GraphLimits::default(),
        )
        .expect_err("existing task mutation");
    assert_eq!(error.kind(), OrchestrationErrorKind::ExistingTaskMutation);
    assert_eq!(
        graph.get("draft").expect("draft").dependencies(),
        &["collect".to_string()]
    );
}

#[test]
fn proposal_limits_cover_nodes_edges_depth_and_fan_out() {
    let graph = base_graph();
    let limits = GraphLimits::new(8, 8, 8, 1);
    let fan_out = graph
        .apply_proposal(
            &proposal(
                &[("review", "reviewer"), ("verify", "reviewer")],
                &[("collect", "review"), ("collect", "verify")],
            ),
            &catalog(),
            &limits,
        )
        .expect_err("fan-out budget");
    assert_eq!(fan_out.kind(), OrchestrationErrorKind::FanOutBudgetExceeded);

    let node_budget = graph
        .apply_proposal(
            &proposal(&[("review", "reviewer"), ("verify", "reviewer")], &[]),
            &catalog(),
            &GraphLimits::new(3, 8, 8, 8),
        )
        .expect_err("node budget");
    assert_eq!(
        node_budget.kind(),
        OrchestrationErrorKind::NodeBudgetExceeded
    );

    let edge_budget = graph
        .apply_proposal(
            &proposal(
                &[("review", "reviewer"), ("verify", "reviewer")],
                &[("collect", "review"), ("collect", "verify")],
            ),
            &catalog(),
            &GraphLimits::new(8, 2, 8, 8),
        )
        .expect_err("edge budget");
    assert_eq!(
        edge_budget.kind(),
        OrchestrationErrorKind::EdgeBudgetExceeded
    );
}

#[test]
fn strict_json_round_trips_proposals_and_rejects_unknown_fields() {
    let change = proposal(&[("review", "reviewer")], &[("collect", "review")]);
    let json = serde_json::to_value(&change).expect("serialize proposal");
    assert_eq!(
        json,
        serde_json::json!({
            "nodes": [{"id": "review", "agentId": "reviewer"}],
            "edges": [{"from": "collect", "to": "review"}]
        })
    );
    let decoded: GraphProposal = serde_json::from_value(json).expect("decode proposal");
    assert_eq!(decoded, change);

    let unknown = serde_json::from_str::<GraphProposal>(
        r#"{"nodes":[{"id":"review","agentId":"reviewer","title":"not allowed"}],"edges":[]}"#,
    )
    .expect_err("unknown proposal field");
    assert!(unknown.to_string().contains("unknown field"));

    let graph = base_graph();
    let graph_json = serde_json::to_value(&graph).expect("serialize graph");
    let restored: TaskGraph = serde_json::from_value(graph_json).expect("decode graph");
    assert_eq!(restored, graph);

    let limits = GraphLimits::default();
    let limits_json = serde_json::to_value(limits).expect("serialize limits");
    assert_eq!(
        limits_json,
        serde_json::json!({
            "maxTasks": 128,
            "maxEdges": 512,
            "maxDepth": 32,
            "maxOutDegree": 32
        })
    );
    let restored_limits: GraphLimits = serde_json::from_value(limits_json).expect("decode limits");
    assert_eq!(restored_limits, limits);

    let unknown_limits = serde_json::from_str::<GraphLimits>(
        r#"{"maxTasks":1,"maxEdges":1,"maxDepth":1,"maxOutDegree":1,"extra":true}"#,
    )
    .expect_err("unknown limits field");
    assert!(unknown_limits.to_string().contains("unknown field"));
}

#[test]
fn proposal_rejects_duplicate_edges_unknown_endpoints_and_invalid_limits() {
    let duplicate = GraphProposal::try_new(
        [TaskProposal::try_new("review", "reviewer").expect("node")],
        [
            GraphEdge::try_new("collect", "review").expect("edge"),
            GraphEdge::try_new("collect", "review").expect("edge"),
        ],
    )
    .expect_err("duplicate edge");
    assert_eq!(duplicate.kind(), OrchestrationErrorKind::DuplicateEdge);

    let unknown = base_graph()
        .apply_proposal(
            &proposal(&[("review", "reviewer")], &[("missing", "review")]),
            &catalog(),
            &GraphLimits::default(),
        )
        .expect_err("unknown edge source");
    assert_eq!(unknown.kind(), OrchestrationErrorKind::UnknownDependency);

    let invalid_limits = GraphLimits::try_new(0, 1, 1, 1).expect_err("zero limit");
    assert_eq!(invalid_limits.kind(), OrchestrationErrorKind::InvalidLimits);
}

#[test]
fn graph_reports_layer_depth_and_fan_out() {
    let graph = TaskGraph::try_new(
        "depth",
        [
            task("root", "researcher", TaskStatus::Pending, &[]),
            task("middle", "reviewer", TaskStatus::Pending, &["root"]),
            task("leaf", "writer", TaskStatus::Pending, &["middle"]),
        ],
    )
    .expect("graph");
    assert_eq!(graph.depth(), 3);
    assert_eq!(graph.fan_out("root"), Some(1));
    assert_eq!(graph.node_count(), 3);
    assert_eq!(graph.fan_out("missing"), None);
}

#[test]
fn proposal_rejects_empty_payload_and_depth_overflow() {
    let empty = GraphProposal::try_new(
        std::iter::empty::<TaskProposal>(),
        std::iter::empty::<GraphEdge>(),
    )
    .expect_err("empty proposal");
    assert_eq!(empty.kind(), OrchestrationErrorKind::EmptyProposal);

    let graph = base_graph();
    let error = graph
        .apply_proposal(
            &proposal(
                &[("review", "reviewer"), ("publish", "writer")],
                &[("collect", "review"), ("review", "publish")],
            ),
            &catalog(),
            &GraphLimits::new(8, 8, 2, 8),
        )
        .expect_err("candidate depth should exceed the limit");
    assert_eq!(error.kind(), OrchestrationErrorKind::DepthBudgetExceeded);
}