klieo-workflow 3.4.0

//! Lower a validated `WorkflowDef` into a `GraphFlow`. Edges are either
//! unconditional (`to`) or conditional (`when` → `then`/`else`).

use crate::def::{EdgeDef, NodeDef, NodeKind, WorkflowDef};
use crate::error::CompileError;
use crate::node::{build_agent_node, EnvelopeGuard, ToolNode};
use crate::registry::Registry;
use klieo_flows::{Flow, GraphFlow};
use std::collections::HashSet;
use std::sync::Arc;

/// Compile a workflow into a runnable flow. Pure — no I/O.
///
/// Validates the entire definition — id uniqueness, entry presence, every
/// node's registry references (model / tool allow-list / subflow) and
/// required `input_from`/`output_to` fields, and every edge endpoint —
/// before any node is constructed.
///
/// The returned flow expects its run input to be a JSON object (the shared
/// envelope); a non-object input fails the run with
/// [`klieo_flows::FlowError::Agent`] rather than being silently coerced.
pub fn compile(def: &WorkflowDef, registry: &Registry) -> Result<Arc<dyn Flow>, CompileError> {
    validate(def, registry)?;
    let mut graph = GraphFlow::new(def.id.clone(), def.entry.clone());
    for node in &def.nodes {
        graph = graph.node(node.id.clone(), build_node(node, registry)?);
    }
    for edge in &def.edges {
        graph = add_edge(graph, edge)?;
    }
    Ok(Arc::new(EnvelopeGuard {
        inner: Arc::new(graph),
    }))
}

fn add_edge(graph: GraphFlow, edge: &EdgeDef) -> Result<GraphFlow, CompileError> {
    match (&edge.to, &edge.when) {
        (Some(_), Some(_)) => Err(CompileError::EdgeHasBothToAndWhen(edge.from.clone())),
        (Some(to), None) => Ok(graph.edge_to(edge.from.clone(), to.clone())),
        (None, Some(cond)) => {
            let (then, otherwise) = branch_targets(edge)?;
            let predicate = crate::condition::compile_condition(cond);
            Ok(graph.edge_branch(edge.from.clone(), predicate, then, otherwise))
        }
        (None, None) => Err(CompileError::EdgeMissingBranchTarget(edge.from.clone())),
    }
}

fn branch_targets(edge: &EdgeDef) -> Result<(String, String), CompileError> {
    match (edge.then.as_ref(), edge.otherwise.as_ref()) {
        (Some(then), Some(otherwise)) => Ok((then.clone(), otherwise.clone())),
        _ => Err(CompileError::EdgeMissingBranchTarget(edge.from.clone())),
    }
}

fn build_node(node: &NodeDef, registry: &Registry) -> Result<Arc<dyn Flow>, CompileError> {
    match &node.kind {
        NodeKind::Agent { agent } => {
            let input_from = require_field(node, "input_from", node.input_from.as_deref())?;
            let output_to = require_field(node, "output_to", node.output_to.as_deref())?;
            Ok(Arc::new(build_agent_node(
                &node.id, agent, input_from, output_to, registry,
            )?))
        }
        NodeKind::Tool { tool } => {
            let input_from = require_field(node, "input_from", node.input_from.as_deref())?;
            let output_to = require_field(node, "output_to", node.output_to.as_deref())?;
            Ok(Arc::new(ToolNode {
                name: node.id.clone(),
                tool: tool.clone(),
                input_from: input_from.to_string(),
                output_to: output_to.to_string(),
            }))
        }
        NodeKind::Subflow { target } => registry
            .subflow(target)
            .ok_or_else(|| unknown_ref("subflow", node, target)),
    }
}

fn require_field<'a>(
    node: &NodeDef,
    field: &'static str,
    value: Option<&'a str>,
) -> Result<&'a str, CompileError> {
    value.ok_or_else(|| CompileError::MissingField {
        node: node.id.clone(),
        field,
    })
}

fn unknown_ref(kind: &'static str, node: &NodeDef, id: &str) -> CompileError {
    CompileError::UnknownRef {
        kind,
        node: node.id.clone(),
        id: id.to_string(),
    }
}

fn validate(def: &WorkflowDef, registry: &Registry) -> Result<(), CompileError> {
    let ids = collect_node_ids(def)?;
    if !ids.contains(def.entry.as_str()) {
        return Err(CompileError::MissingEntry(def.entry.clone()));
    }
    for node in &def.nodes {
        validate_node(node, registry)?;
    }
    validate_edges(&def.edges, &ids)
}

fn collect_node_ids(def: &WorkflowDef) -> Result<HashSet<&str>, CompileError> {
    let mut ids = HashSet::new();
    for node in &def.nodes {
        if !ids.insert(node.id.as_str()) {
            return Err(CompileError::DuplicateNodeId(node.id.clone()));
        }
    }
    Ok(ids)
}

fn validate_node(node: &NodeDef, registry: &Registry) -> Result<(), CompileError> {
    match &node.kind {
        NodeKind::Agent { agent } => validate_agent_node(node, agent, registry),
        NodeKind::Tool { tool } => {
            require_io_fields(node)?;
            if registry.allows_tool(tool) {
                Ok(())
            } else {
                Err(unknown_ref("tool", node, tool))
            }
        }
        NodeKind::Subflow { target } => match registry.subflow(target) {
            Some(_) => Ok(()),
            None => Err(unknown_ref("subflow", node, target)),
        },
    }
}

fn validate_agent_node(
    node: &NodeDef,
    cfg: &crate::def::AgentConfig,
    registry: &Registry,
) -> Result<(), CompileError> {
    require_io_fields(node)?;
    if registry.model(&cfg.model).is_none() {
        return Err(unknown_ref("model", node, &cfg.model));
    }
    for tool_id in &cfg.tools {
        if !registry.allows_tool(tool_id) {
            return Err(unknown_ref("tool", node, tool_id));
        }
    }
    Ok(())
}

fn require_io_fields(node: &NodeDef) -> Result<(), CompileError> {
    require_field(node, "input_from", node.input_from.as_deref())?;
    require_field(node, "output_to", node.output_to.as_deref())?;
    Ok(())
}

fn validate_edges(edges: &[EdgeDef], ids: &HashSet<&str>) -> Result<(), CompileError> {
    let mut seen_sources = HashSet::new();
    for edge in edges {
        if !ids.contains(edge.from.as_str()) {
            return Err(CompileError::MissingEndpoint(edge.from.clone()));
        }
        validate_edge_targets(edge, ids)?;
        if !seen_sources.insert(edge.from.as_str()) {
            return Err(CompileError::DuplicateEdgeSource(edge.from.clone()));
        }
    }
    Ok(())
}

fn validate_edge_targets(edge: &EdgeDef, ids: &HashSet<&str>) -> Result<(), CompileError> {
    match (&edge.to, &edge.when) {
        (Some(_), Some(_)) => Err(CompileError::EdgeHasBothToAndWhen(edge.from.clone())),
        (Some(to), None) => require_endpoint(to, ids),
        (None, Some(_)) => {
            let (then, otherwise) = branch_targets(edge)?;
            require_endpoint(&then, ids)?;
            require_endpoint(&otherwise, ids)
        }
        (None, None) => Err(CompileError::EdgeMissingBranchTarget(edge.from.clone())),
    }
}

fn require_endpoint(node_id: &str, ids: &HashSet<&str>) -> Result<(), CompileError> {
    if ids.contains(node_id) {
        Ok(())
    } else {
        Err(CompileError::MissingEndpoint(node_id.to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::node::NON_OBJECT_ENVELOPE;
    use crate::test_support::{const_subflow, dummy_llm, test_ctx};
    use klieo_flows::FlowError;

    fn agent_node(id: &str) -> NodeDef {
        serde_json::from_value(serde_json::json!({
            "id": id, "kind": "agent", "input_from": "q", "output_to": "a",
            "agent": { "model": "m", "system_prompt": "p" }
        }))
        .unwrap()
    }

    fn subflow_node(id: &str, target: &str) -> NodeDef {
        serde_json::from_value(serde_json::json!({
            "id": id, "kind": "subflow", "ref": target
        }))
        .unwrap()
    }

    fn def_with(nodes: Vec<NodeDef>, entry: &str, edges: serde_json::Value) -> WorkflowDef {
        serde_json::from_value(serde_json::json!({
            "id": "w", "entry": entry, "nodes": nodes, "edges": edges
        }))
        .unwrap()
    }

    fn reg() -> Registry {
        Registry::new().with_model("m", dummy_llm())
    }

    /// `compile()` returns `Arc<dyn Flow>` on success, which isn't `Debug`,
    /// so `Result::unwrap_err` (bound on the `Ok` side being `Debug`) isn't
    /// available. Centralises the "must fail" assertion for every negative
    /// test below.
    fn expect_compile_err(def: &WorkflowDef, registry: &Registry) -> CompileError {
        match compile(def, registry) {
            Ok(_) => panic!("expected compile() to fail, but it succeeded"),
            Err(e) => e,
        }
    }

    #[test]
    fn duplicate_node_id_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("a")],
            "a",
            serde_json::json!([]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::DuplicateNodeId(_)
        ));
    }

    #[test]
    fn missing_entry_rejected() {
        let def = def_with(vec![agent_node("a")], "ghost", serde_json::json!([]));
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingEntry(_)
        ));
    }

    #[test]
    fn edge_to_missing_node_rejected() {
        let def = def_with(
            vec![agent_node("a")],
            "a",
            serde_json::json!([{"from": "a", "to": "ghost"}]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingEndpoint(_)
        ));
    }

    #[test]
    fn two_edges_from_same_node_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b"), agent_node("c")],
            "a",
            serde_json::json!([{"from": "a", "to": "b"}, {"from": "a", "to": "c"}]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::DuplicateEdgeSource(_)
        ));
    }

    #[test]
    fn unregistered_tool_node_rejected() {
        let tool_node: NodeDef = serde_json::from_value(serde_json::json!({
            "id": "t", "kind": "tool", "tool": "danger", "input_from": "x", "output_to": "y"
        }))
        .unwrap();
        let def = def_with(vec![tool_node], "t", serde_json::json!([]));
        assert!(matches!(
            expect_compile_err(&def, &Registry::new()),
            CompileError::UnknownRef { kind: "tool", .. }
        ));
    }

    #[test]
    fn agent_node_missing_output_to_rejected() {
        let n: NodeDef = serde_json::from_value(serde_json::json!({
            "id": "a", "kind": "agent", "input_from": "q",
            "agent": { "model": "m", "system_prompt": "p" }
        }))
        .unwrap();
        let def = def_with(vec![n], "a", serde_json::json!([]));
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingField {
                field: "output_to",
                ..
            }
        ));
    }

    #[test]
    fn agent_node_missing_input_from_rejected() {
        let n: NodeDef = serde_json::from_value(serde_json::json!({
            "id": "a", "kind": "agent", "output_to": "a",
            "agent": { "model": "m", "system_prompt": "p" }
        }))
        .unwrap();
        let def = def_with(vec![n], "a", serde_json::json!([]));
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingField {
                field: "input_from",
                ..
            }
        ));
    }

    #[test]
    fn tool_node_missing_input_from_rejected() {
        let n: NodeDef = serde_json::from_value(serde_json::json!({
            "id": "t", "kind": "tool", "tool": "lookup", "output_to": "y"
        }))
        .unwrap();
        let def = def_with(vec![n], "t", serde_json::json!([]));
        let registry = Registry::new().with_tool("lookup");
        assert!(matches!(
            expect_compile_err(&def, &registry),
            CompileError::MissingField {
                field: "input_from",
                ..
            }
        ));
    }

    #[test]
    fn edge_from_missing_node_rejected() {
        let def = def_with(
            vec![agent_node("a")],
            "a",
            serde_json::json!([{"from": "ghost", "to": "a"}]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingEndpoint(_)
        ));
    }

    #[test]
    fn edge_without_target_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b")],
            "a",
            serde_json::json!([{"from": "a"}]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::EdgeMissingBranchTarget(_)
        ));
    }

    #[test]
    fn unregistered_subflow_node_rejected() {
        let def = def_with(vec![subflow_node("s", "ghost")], "s", serde_json::json!([]));
        assert!(matches!(
            expect_compile_err(&def, &Registry::new()),
            CompileError::UnknownRef {
                kind: "subflow",
                ..
            }
        ));
    }

    #[test]
    fn valid_linear_workflow_compiles() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b")],
            "a",
            serde_json::json!([{"from": "a", "to": "b"}]),
        );
        assert_eq!(compile(&def, &reg()).unwrap().name(), "w");
    }

    #[tokio::test]
    async fn subflow_node_runs_registered_flow() {
        let registry = Registry::new().with_subflow("seed", const_subflow("SEED"));
        let def = def_with(vec![subflow_node("s", "seed")], "s", serde_json::json!([]));
        let flow = compile(&def, &registry).unwrap();
        let out = flow.run(test_ctx(), serde_json::json!({})).await.unwrap();
        assert_eq!(out["tag"], serde_json::json!("SEED"));
    }

    #[tokio::test]
    async fn non_object_envelope_rejected() {
        for bad in [
            serde_json::json!([1, 2, 3]),
            serde_json::json!("string"),
            serde_json::json!(5),
        ] {
            let registry = Registry::new().with_subflow("seed", const_subflow("SEED"));
            let def = def_with(vec![subflow_node("s", "seed")], "s", serde_json::json!([]));
            let flow = compile(&def, &registry).unwrap();
            match flow.run(test_ctx(), bad.clone()).await {
                Ok(_) => panic!("expected non-object input {bad} to be rejected"),
                Err(FlowError::Agent(msg)) => assert!(
                    msg.contains(NON_OBJECT_ENVELOPE),
                    "unexpected message: {msg}"
                ),
                Err(other) => panic!("expected FlowError::Agent, got {other:?}"),
            }
        }
    }

    #[test]
    fn edge_with_both_to_and_when_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b")],
            "a",
            serde_json::json!([{
                "from": "a", "to": "b",
                "when": {"field": "x", "op": "exists"}, "then": "b", "else": "b"
            }]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::EdgeHasBothToAndWhen(_)
        ));
    }

    #[test]
    fn conditional_edge_missing_else_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b")],
            "a",
            serde_json::json!([{
                "from": "a", "when": {"field": "x", "op": "exists"}, "then": "b"
            }]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::EdgeMissingBranchTarget(_)
        ));
    }

    #[test]
    fn conditional_edge_missing_endpoint_rejected() {
        let def = def_with(
            vec![agent_node("a"), agent_node("b")],
            "a",
            serde_json::json!([{
                "from": "a", "when": {"field": "x", "op": "exists"},
                "then": "ghost", "else": "b"
            }]),
        );
        assert!(matches!(
            expect_compile_err(&def, &reg()),
            CompileError::MissingEndpoint(_)
        ));
    }

    #[tokio::test]
    async fn conditional_edge_routes_on_field() {
        let registry = Registry::new()
            .with_subflow("seed_high", const_subflow("SEED"))
            .with_subflow("high", const_subflow("HIGH"))
            .with_subflow("low", const_subflow("LOW"));
        let def: WorkflowDef = serde_json::from_value(serde_json::json!({
            "id": "route", "entry": "seed",
            "nodes": [
                { "id": "seed", "kind": "subflow", "ref": "seed_high" },
                { "id": "high", "kind": "subflow", "ref": "high" },
                { "id": "low",  "kind": "subflow", "ref": "low" }
            ],
            "edges": [
                { "from": "seed", "when": {"field": "tag", "op": "eq", "value": "SEED"},
                  "then": "high", "else": "low" }
            ]
        }))
        .unwrap();
        let flow = compile(&def, &registry).unwrap();
        let out = flow.run(test_ctx(), serde_json::json!({})).await.unwrap();
        assert_eq!(out["tag"], serde_json::json!("HIGH"));
    }
}