tinyagents 0.1.0

//! Tests for the expressive language pipeline: lexer, parser, compiler,
//! capability binding, and graph materialisation.

use crate::graph::{Node, NodeOutput};
use crate::language::compiler::{
    CapabilityResolver, NodeFactory, bind_capabilities, build_graph, compile,
};
use crate::language::lexer::tokenize;
use crate::language::parser::{parse, parse_str};
use crate::language::types::{Literal, NodeSpec, Routing, Token};

/// The `support_agent` fixture from the module spec: an agent node with a tool
/// loop plus conditional routing to `END`.
const SUPPORT_AGENT: &str = r#"
// A support workflow with a tool loop.
graph support_agent {
  start agent

  defaults {
    recursion_limit 50
    backoff "exponential"
    checkpoint inherit
  }

  channel messages messages
  channel tool_calls append

  node agent {
    kind agent
    model "default"
    system "Resolve support requests using tools when useful."
    tools ["lookup_user", "create_ticket"]
    routes {
      tool_call -> tools
      final -> END
    }
  }

  node tools {
    kind tool_executor
    next agent
  }
}
"#;

// ---------------------------------------------------------------------------
// Lexer
// ---------------------------------------------------------------------------

#[test]
fn tokenizes_punctuation_and_arrow() {
    let tokens = tokenize("a -> b { } [ ] ,").unwrap();
    let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect();
    assert_eq!(
        kinds,
        vec![
            Token::Ident("a".into()),
            Token::Arrow,
            Token::Ident("b".into()),
            Token::LBrace,
            Token::RBrace,
            Token::LBracket,
            Token::RBracket,
            Token::Comma,
            Token::Eof,
        ]
    );
}

#[test]
fn tokenizes_strings_numbers_and_comments() {
    let tokens = tokenize("// comment\n\"hi\\n\" 50 1.5 -3").unwrap();
    let kinds: Vec<_> = tokens.into_iter().map(|t| t.token).collect();
    assert_eq!(
        kinds,
        vec![
            Token::Str("hi\n".into()),
            Token::Num(50.0),
            Token::Num(1.5),
            Token::Num(-3.0),
            Token::Eof,
        ]
    );
}

#[test]
fn tracks_line_and_column_spans() {
    let tokens = tokenize("graph\n  foo").unwrap();
    assert_eq!(tokens[0].span.line, 1);
    assert_eq!(tokens[0].span.column, 1);
    assert_eq!(tokens[1].span.line, 2);
    assert_eq!(tokens[1].span.column, 3);
}

#[test]
fn unterminated_string_is_a_parse_error() {
    let err = tokenize("\"oops").unwrap_err();
    assert!(matches!(err, crate::error::RustAgentsError::Parse { .. }));
}

#[test]
fn invalid_escape_is_a_parse_error() {
    let err = tokenize("\"bad\\x\"").unwrap_err();
    assert!(matches!(err, crate::error::RustAgentsError::Parse { .. }));
}

// ---------------------------------------------------------------------------
// Parser
// ---------------------------------------------------------------------------

#[test]
fn parses_support_agent_into_ast() {
    let program = parse_str(SUPPORT_AGENT).unwrap();
    assert_eq!(program.graphs.len(), 1);
    let graph = &program.graphs[0];

    assert_eq!(graph.name, "support_agent");
    assert_eq!(graph.start.as_deref(), Some("agent"));
    assert_eq!(graph.channels.len(), 2);
    assert_eq!(graph.channels[0].name, "messages");
    assert_eq!(graph.channels[0].reducer, "messages");
    assert_eq!(graph.channels[1].reducer, "append");

    // Defaults preserve declared order and literal kinds.
    assert_eq!(graph.defaults.len(), 3);
    assert_eq!(graph.defaults[0].0, "recursion_limit");
    assert_eq!(graph.defaults[0].1, Literal::Num(50.0));
    assert_eq!(graph.defaults[1].1, Literal::Str("exponential".into()));
    assert_eq!(graph.defaults[2].1, Literal::Ident("inherit".into()));

    assert_eq!(graph.nodes.len(), 2);
    let agent = &graph.nodes[0];
    assert_eq!(agent.kind.as_deref(), Some("agent"));
    assert_eq!(agent.model.as_deref(), Some("default"));
    assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]);
    assert_eq!(agent.routes.len(), 2);
    assert_eq!(agent.routes[0].label, "tool_call");
    assert_eq!(agent.routes[0].target, "tools");
    assert_eq!(agent.routes[1].target, "END");

    let tools = &graph.nodes[1];
    assert_eq!(tools.kind.as_deref(), Some("tool_executor"));
    assert_eq!(tools.next.as_deref(), Some("agent"));
}

#[test]
fn parses_top_level_edge() {
    let src = "graph g { start a node a { } node b { } a -> b b -> END }";
    let program = parse_str(src).unwrap();
    let graph = &program.graphs[0];
    assert_eq!(graph.edges.len(), 2);
    assert_eq!(graph.edges[0].from, "a");
    assert_eq!(graph.edges[0].to, "b");
    assert_eq!(graph.edges[1].to, "END");
}

#[test]
fn parse_reports_unexpected_token() {
    // Missing graph name.
    let tokens = tokenize("graph { }").unwrap();
    let err = parse(&tokens).unwrap_err();
    match err {
        crate::error::RustAgentsError::Parse { message, .. } => {
            assert!(message.contains("expected identifier"), "{message}");
        }
        other => panic!("expected parse error, got {other:?}"),
    }
}

#[test]
fn parse_rejects_unknown_node_item() {
    let src = "graph g { start a node a { bogus x } }";
    let err = parse_str(src).unwrap_err();
    match err {
        crate::error::RustAgentsError::Parse { message, .. } => {
            assert!(message.contains("unknown node item"), "{message}");
        }
        other => panic!("expected parse error, got {other:?}"),
    }
}

// ---------------------------------------------------------------------------
// Compiler: AST -> Blueprint
// ---------------------------------------------------------------------------

#[test]
fn compiles_support_agent_blueprint() {
    let program = parse_str(SUPPORT_AGENT).unwrap();
    let blueprints = compile(&program).unwrap();
    assert_eq!(blueprints.len(), 1);
    let bp = &blueprints[0];

    assert_eq!(bp.graph_id, "support_agent");
    assert_eq!(bp.start, "agent");
    assert_eq!(bp.channels.len(), 2);
    assert_eq!(bp.defaults.len(), 3);
    assert_eq!(bp.nodes.len(), 2);

    let agent = &bp.nodes[0];
    assert_eq!(agent.kind, "agent");
    assert_eq!(agent.tools, vec!["lookup_user", "create_ticket"]);
    match &agent.routing {
        Routing::Conditional(routes) => {
            assert_eq!(routes.len(), 2);
            assert_eq!(routes[0], ("tool_call".into(), "tools".into()));
            assert_eq!(routes[1], ("final".into(), "END".into()));
        }
        other => panic!("expected conditional routing, got {other:?}"),
    }

    let tools = &bp.nodes[1];
    assert_eq!(tools.routing, Routing::Next("agent".into()));
}

#[test]
fn next_end_lowers_to_terminal() {
    let src = "graph g { start a node a { kind model next END } }";
    let bp = &compile(&parse_str(src).unwrap()).unwrap()[0];
    assert_eq!(bp.nodes[0].routing, Routing::Terminal);
}

#[test]
fn blueprint_round_trips_through_serde() {
    let bp = compile(&parse_str(SUPPORT_AGENT).unwrap())
        .unwrap()
        .remove(0);
    let json = serde_json::to_string(&bp).unwrap();
    let back: crate::language::types::Blueprint = serde_json::from_str(&json).unwrap();
    assert_eq!(bp, back);
}

#[test]
fn missing_start_is_a_compile_error() {
    let src = "graph g { node a { kind model } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(matches!(err, crate::error::RustAgentsError::Compile(_)));
    assert!(err.to_string().contains("no `start`"), "{err}");
}

#[test]
fn start_not_defined_is_a_compile_error() {
    let src = "graph g { start missing node a { kind model } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("is not defined"), "{err}");
}

#[test]
fn duplicate_node_is_a_compile_error() {
    let src = "graph g { start a node a { kind model } node a { kind model } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("duplicate node"), "{err}");
}

#[test]
fn unknown_route_target_is_a_compile_error() {
    let src = "graph g { start a node a { routes { go -> ghost } } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("route target"), "{err}");
}

#[test]
fn unknown_next_target_is_a_compile_error() {
    let src = "graph g { start a node a { next ghost } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("next target"), "{err}");
}

#[test]
fn mixing_next_and_routes_is_a_compile_error() {
    let src = "graph g { start a node a { next b routes { x -> b } } node b { } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("mixes static routing"), "{err}");
}

#[test]
fn mixing_edge_and_routes_is_a_compile_error() {
    let src = "graph g { start a node a { routes { x -> b } } node b { } a -> b }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("mixes static routing"), "{err}");
}

#[test]
fn duplicate_route_label_is_a_compile_error() {
    let src = "graph g { start a node a { routes { x -> b\n x -> b } } node b { } }";
    let err = compile(&parse_str(src).unwrap()).unwrap_err();
    assert!(err.to_string().contains("duplicate route label"), "{err}");
}

// ---------------------------------------------------------------------------
// Capability binding
// ---------------------------------------------------------------------------

#[test]
fn bind_capabilities_accepts_allowed_references() {
    let bp = compile(&parse_str(SUPPORT_AGENT).unwrap())
        .unwrap()
        .remove(0);
    let resolver = CapabilityResolver::from_lists(
        ["default".to_string()],
        ["lookup_user".to_string(), "create_ticket".to_string()],
    );
    bind_capabilities(&bp, &resolver).unwrap();
}

#[test]
fn bind_capabilities_rejects_unknown_model() {
    let bp = compile(&parse_str(SUPPORT_AGENT).unwrap())
        .unwrap()
        .remove(0);
    let resolver = CapabilityResolver::new()
        .allow_tool("lookup_user")
        .allow_tool("create_ticket");
    let err = bind_capabilities(&bp, &resolver).unwrap_err();
    assert!(matches!(err, crate::error::RustAgentsError::Capability(_)));
    assert!(err.to_string().contains("unknown model"), "{err}");
}

#[test]
fn bind_capabilities_rejects_unknown_tool() {
    let bp = compile(&parse_str(SUPPORT_AGENT).unwrap())
        .unwrap()
        .remove(0);
    let resolver = CapabilityResolver::new().allow_model("default");
    let err = bind_capabilities(&bp, &resolver).unwrap_err();
    assert!(err.to_string().contains("unknown tool"), "{err}");
}

// ---------------------------------------------------------------------------
// Graph materialisation + execution
// ---------------------------------------------------------------------------

#[derive(Clone, Debug, PartialEq, Eq)]
struct TestState {
    trail: Vec<String>,
    agent_visits: u32,
}

/// A trivial factory that materialises echo/route/end nodes purely from the
/// declarative [`NodeSpec`]. It demonstrates that runnable behaviour comes from
/// Rust, not the source: each node records its name; terminal nodes end, `next`
/// nodes continue, and conditional nodes loop once before terminating.
struct TestFactory;

impl NodeFactory<TestState> for TestFactory {
    fn make(&self, spec: &NodeSpec) -> crate::error::Result<Node<TestState>> {
        let name = spec.name.clone();
        let routing = spec.routing.clone();
        Ok(Node::new(name.clone(), move |mut state: TestState| {
            let name = name.clone();
            let routing = routing.clone();
            async move {
                state.trail.push(name.clone());
                let output = match &routing {
                    Routing::Terminal => NodeOutput::end(state),
                    Routing::Next(_) => NodeOutput::continue_with(state),
                    Routing::Conditional(_) => {
                        state.agent_visits += 1;
                        if state.agent_visits >= 2 {
                            // Take the `final -> END` route by ending.
                            NodeOutput::end(state)
                        } else {
                            NodeOutput::route(state, "tool_call")
                        }
                    }
                };
                Ok(output)
            }
        }))
    }
}

#[tokio::test]
async fn build_graph_runs_to_end() {
    let bp = compile(&parse_str(SUPPORT_AGENT).unwrap())
        .unwrap()
        .remove(0);
    let graph = build_graph(&bp, &TestFactory).unwrap();

    let run = graph
        .run(TestState {
            trail: Vec::new(),
            agent_visits: 0,
        })
        .await
        .unwrap();

    // agent -> tools -> agent (ends on second visit).
    assert_eq!(run.visited, vec!["agent", "tools", "agent"]);
    assert_eq!(run.state.trail, vec!["agent", "tools", "agent"]);
    assert_eq!(run.state.agent_visits, 2);
}

#[tokio::test]
async fn build_graph_handles_linear_terminal() {
    let src = "graph g { start a node a { kind model next b } node b { kind model next END } }";
    let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0);
    let graph = build_graph(&bp, &TestFactory).unwrap();
    let run = graph
        .run(TestState {
            trail: Vec::new(),
            agent_visits: 0,
        })
        .await
        .unwrap();
    assert_eq!(run.visited, vec!["a", "b"]);
}