focus/crates/focus-core/tests/agent_loop.rs

//! Integration tests for the agent loop, driven by the mock provider.

use focus_core::event::VecSink;
use focus_core::mock::{make_text_response, make_tool_call_response, MockProvider};
use focus_core::model::*;
use focus_core::tool::{Tool, ToolEffects, ToolRegistry, ToolResult, ToolUpdateSink};
use focus_core::{Agent, AgentConfig};
use focus_json::JsonValue;

/// A simple echo tool that returns its arguments as text.
#[derive(Debug)]
struct EchoTool;

impl Tool for EchoTool {
    fn name(&self) -> &str {
        "echo"
    }
    fn description(&self) -> &str {
        "echoes the arguments back"
    }
    fn parameters(&self) -> JsonValue {
        let mut o = JsonValue::obj();
        let mut props = JsonValue::obj();
        props.insert("text", JsonValue::Str("the text".into())).ok();
        o.insert("type", "object".into()).ok();
        o.insert("properties", props).ok();
        o
    }
    fn effects(&self) -> ToolEffects {
        ToolEffects::READ
    }
    fn execute(
        &self,
        _id: &str,
        args: &JsonValue,
        _on_update: Option<&ToolUpdateSink>,
    ) -> Result<ToolResult, focus_core::CoreError> {
        let text = args.get_str("text").unwrap_or("(none)").to_string();
        Ok(ToolResult::text(format!("echo: {}", text)))
    }
}

fn agent_with(provider: MockProvider) -> Agent {
    let config = AgentConfig {
        model: "mock-model".into(),
        ..Default::default()
    };
    let tools = ToolRegistry::with(Box::new(EchoTool));
    Agent::new(config, Box::new(provider), tools)
}

#[test]
fn single_text_response_completes() {
    let provider = MockProvider::single_text("hello there", StopReason::Stop);
    let mut agent = agent_with(provider);
    let mut sink = VecSink::new();
    agent.prompt("hi", &mut sink).expect("run failed");

    assert_eq!(sink.count("agent_start"), 1);
    assert_eq!(sink.count("agent_end"), 1);
    assert_eq!(sink.count("message_end"), 2); // user + assistant

    // Transcript: user + assistant
    let msgs = agent.messages();
    assert_eq!(msgs.len(), 2);
    assert!(matches!(msgs[0], Message::User(_)));
    let assistant = match &msgs[1] {
        Message::Assistant(a) => a,
        _ => panic!("expected assistant message"),
    };
    let text = assistant
        .content
        .iter()
        .find_map(|c| c.as_text())
        .expect("assistant has text");
    assert_eq!(text.text, "hello there");
    assert_eq!(assistant.stop_reason, StopReason::Stop);
}

#[test]
fn tool_call_then_text_completes_two_turns() {
    // Turn 1: model calls "echo" tool. Turn 2: model gives final text.
    let mut args = JsonValue::obj();
    args.insert("text", "world".into()).ok();
    let provider = MockProvider::new(vec![
        make_tool_call_response(None, "echo", "call_1", args, StopReason::ToolUse),
        make_text_response("all done", StopReason::Stop),
    ]);
    let mut agent = agent_with(provider);
    let mut sink = VecSink::new();
    agent.prompt("run echo", &mut sink).expect("run failed");

    // Two turns.
    assert_eq!(sink.count("turn_end"), 2);
    assert_eq!(sink.count("tool_execution_end"), 1);
    assert_eq!(sink.count("agent_end"), 1);

    // Transcript: user, assistant(toolcall), toolResult, assistant(text)
    let msgs = agent.messages();
    assert_eq!(msgs.len(), 4);
    assert!(matches!(msgs[0], Message::User(_)));
    assert!(matches!(
        &msgs[1],
        Message::Assistant(a) if a.content.iter().any(|c| c.as_tool_call().is_some())
    ));
    assert!(matches!(
        &msgs[2],
        Message::ToolResult(tr) if tr.tool_name == "echo" && !tr.is_error
    ));
    let text = match &msgs[3] {
        Message::Assistant(a) => a
            .content
            .iter()
            .find_map(|c| c.as_text())
            .map(|t| t.text.clone()),
        _ => None,
    };
    assert_eq!(text.as_deref(), Some("all done"));

    // The tool result content should contain the echo.
    if let Message::ToolResult(tr) = &msgs[2] {
        let result_text = tr
            .content
            .iter()
            .find_map(|c| c.as_text())
            .map(|t| t.text.clone());
        assert_eq!(result_text.as_deref(), Some("echo: world"));
    }
}

#[test]
fn unknown_tool_produces_error_result() {
    let provider = MockProvider::new(vec![
        make_tool_call_response(
            None,
            "nonexistent",
            "call_1",
            JsonValue::obj(),
            StopReason::ToolUse,
        ),
        make_text_response("ok", StopReason::Stop),
    ]);
    let mut agent = agent_with(provider);
    let mut sink = VecSink::new();
    agent.prompt("go", &mut sink).expect("run failed");

    // The tool result should be an error.
    let msgs = agent.messages();
    assert!(matches!(msgs[2], Message::ToolResult(_)));
    if let Message::ToolResult(tr) = &msgs[2] {
        assert!(tr.is_error, "expected error for unknown tool");
    }
}

#[test]
fn error_stop_reason_ends_run_immediately() {
    let provider = MockProvider::new(vec![make_text_response("oops", StopReason::Error)]);
    let mut agent = agent_with(provider);
    let mut sink = VecSink::new();
    agent.prompt("hi", &mut sink).expect("run failed");

    assert_eq!(sink.count("turn_end"), 1);
    assert_eq!(sink.count("agent_end"), 1);
    let assistant = match agent.messages().last().unwrap() {
        Message::Assistant(a) => a,
        _ => panic!(),
    };
    assert_eq!(assistant.stop_reason, StopReason::Error);
}

#[test]
fn continue_requires_non_assistant_last_message() {
    let provider = MockProvider::single_text("hi", StopReason::Stop);
    let mut agent = agent_with(provider);
    let mut sink = VecSink::new();
    agent.prompt("hello", &mut sink).expect("run failed");
    // Last message is now assistant — continue must fail.
    let result = agent.r#continue(&mut sink);
    assert!(result.is_err());
}