Skip to main content

recursive/llm/
mock.rs

1//! Deterministic LLM for tests and offline development.
2//!
3//! `MockProvider` is fed a queue of pre-baked completions. The agent treats
4//! it identically to a real provider, so the agent loop is fully testable
5//! without network access.
6
7use std::sync::Mutex;
8
9use async_trait::async_trait;
10
11use super::{Completion, LlmProvider, ToolSpec};
12use crate::error::{Error, Result};
13use crate::message::Message;
14
15#[derive(Default)]
16pub struct MockProvider {
17    scripted: Mutex<Vec<Completion>>,
18    calls: Mutex<Vec<Vec<Message>>>,
19}
20
21impl MockProvider {
22    pub fn new(scripted: Vec<Completion>) -> Self {
23        Self {
24            scripted: Mutex::new(scripted),
25            calls: Mutex::new(Vec::new()),
26        }
27    }
28
29    /// Snapshot of the transcripts the agent has sent to this provider.
30    pub fn calls(&self) -> Vec<Vec<Message>> {
31        self.calls.lock().unwrap().clone()
32    }
33}
34
35#[async_trait]
36impl LlmProvider for MockProvider {
37    async fn complete(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Completion> {
38        self.calls.lock().unwrap().push(messages.to_vec());
39        let mut queue = self.scripted.lock().unwrap();
40        if queue.is_empty() {
41            return Err(Error::Llm(
42                "MockProvider: no scripted completions left".into(),
43            ));
44        }
45        Ok(queue.remove(0))
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[tokio::test]
54    async fn returns_scripted_in_order_and_records_calls() {
55        let provider = MockProvider::new(vec![
56            Completion {
57                content: "one".into(),
58                tool_calls: vec![],
59                finish_reason: Some("stop".into()),
60            },
61            Completion {
62                content: "two".into(),
63                tool_calls: vec![],
64                finish_reason: Some("stop".into()),
65            },
66        ]);
67
68        let m1 = vec![Message::user("hi")];
69        let r1 = provider.complete(&m1, &[]).await.unwrap();
70        assert_eq!(r1.content, "one");
71
72        let m2 = vec![Message::user("bye")];
73        let r2 = provider.complete(&m2, &[]).await.unwrap();
74        assert_eq!(r2.content, "two");
75
76        assert_eq!(provider.calls().len(), 2);
77        assert_eq!(provider.calls()[0][0].content, "hi");
78    }
79
80    #[tokio::test]
81    async fn errors_when_queue_drained() {
82        let provider = MockProvider::new(vec![]);
83        let err = provider.complete(&[], &[]).await.unwrap_err();
84        assert!(matches!(err, Error::Llm(_)));
85    }
86}