capo-agent 0.5.0

//! Test-only helpers shared across capo crates. Enabled by the `test-support`
//! Cargo feature; off by default so production builds never pull these in.
//!
//! The motivating tool is `ScriptedLlm` — a deterministic `LlmClient` that
//! lets tests drive the agent through real tool calls (and through the
//! permission system) without needing a real LLM endpoint.

use std::collections::VecDeque;
use std::sync::Mutex;

use async_trait::async_trait;
use motosan_agent_loop::{
    ChatOutput, LlmClient, LlmResponse, Message, Result as MotosanResult, ToolCallItem,
};
use motosan_agent_tool::ToolDef;
use serde_json::Value;

/// A test `LlmClient` that replays a scripted sequence of responses. Each
/// call to `chat()` pops the next `LlmResponse` from the script; an
/// exhausted script panics so test bugs (an unexpected extra chat call)
/// surface loudly instead of silently hanging.
///
/// Construct via [`ScriptedLlm::builder`].
pub struct ScriptedLlm {
    script: Mutex<VecDeque<LlmResponse>>,
}

impl ScriptedLlm {
    pub fn builder() -> ScriptedLlmBuilder {
        ScriptedLlmBuilder {
            script: VecDeque::new(),
        }
    }
}

#[async_trait]
impl LlmClient for ScriptedLlm {
    async fn chat(&self, _messages: &[Message], _tools: &[ToolDef]) -> MotosanResult<ChatOutput> {
        let next = match self.script.lock() {
            Ok(mut s) => s.pop_front(),
            Err(poison) => poison.into_inner().pop_front(),
        };
        let response = next.unwrap_or_else(|| {
            panic!("ScriptedLlm: script exhausted (chat called more times than scripted)")
        });
        Ok(ChatOutput::new(response))
    }
}

/// Fluent builder for [`ScriptedLlm`].
pub struct ScriptedLlmBuilder {
    script: VecDeque<LlmResponse>,
}

impl ScriptedLlmBuilder {
    /// Queue an assistant-message response (no tool calls).
    pub fn then_message(mut self, text: impl Into<String>) -> Self {
        self.script.push_back(LlmResponse::Message(text.into()));
        self
    }

    /// Queue a one-tool-call response. `id` is the tool-call id (capo's
    /// `UiEvent::ToolCallStarted.id` will surface it), `name` is the tool
    /// (e.g. `"bash"`, `"read"`), `args` is the JSON argument object the
    /// tool's schema expects.
    pub fn then_tool_call(
        mut self,
        id: impl Into<String>,
        name: impl Into<String>,
        args: Value,
    ) -> Self {
        self.script
            .push_back(LlmResponse::ToolCalls(vec![ToolCallItem {
                id: id.into(),
                name: name.into(),
                args,
            }]));
        self
    }

    pub fn build(self) -> ScriptedLlm {
        ScriptedLlm {
            script: Mutex::new(self.script),
        }
    }
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use super::*;
    use motosan_agent_loop::LlmClient;

    #[tokio::test]
    async fn scripted_llm_replays_in_order() {
        let llm = ScriptedLlm::builder()
            .then_message("first")
            .then_message("second")
            .build();
        let a = llm.chat(&[], &[]).await.expect("chat 1");
        let b = llm.chat(&[], &[]).await.expect("chat 2");
        assert!(matches!(
            a.response,
            motosan_agent_loop::LlmResponse::Message(ref s) if s == "first"
        ));
        assert!(matches!(
            b.response,
            motosan_agent_loop::LlmResponse::Message(ref s) if s == "second"
        ));
    }

    #[tokio::test]
    async fn scripted_llm_then_tool_call_emits_tool_calls_variant() {
        let llm = ScriptedLlm::builder()
            .then_tool_call("t1", "bash", serde_json::json!({"command": "echo hi"}))
            .build();
        let out = llm.chat(&[], &[]).await.expect("chat");
        let calls = match out.response {
            motosan_agent_loop::LlmResponse::ToolCalls(c) => c,
            other => panic!("expected ToolCalls, got {other:?}"),
        };
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].id, "t1");
        assert_eq!(calls[0].name, "bash");
        assert_eq!(calls[0].args["command"], "echo hi");
    }

    #[tokio::test]
    #[should_panic(expected = "script exhausted")]
    async fn scripted_llm_panics_when_script_runs_out() {
        let llm = ScriptedLlm::builder().build();
        let _ = llm.chat(&[], &[]).await;
    }
}