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::StructuredRequest;
12use super::{Completion, LlmProvider, StreamSender, ToolSpec};
13use crate::error::{Error, Result};
14use crate::message::Message;
15use tracing::Instrument;
16
17#[derive(Default)]
18pub struct MockProvider {
19    scripted: Mutex<Vec<Completion>>,
20    /// Queue of pre-baked JSON responses for `complete_structured`.
21    structured_responses: Mutex<Vec<Result<serde_json::Value>>>,
22    calls: Mutex<Vec<Vec<Message>>>,
23}
24
25impl MockProvider {
26    /// Alias for `new` — both accept completions that may have `usage` set.
27    /// Provided for API symmetry when callers want to be explicit about usage.
28    pub fn with_usage(scripted: Vec<Completion>) -> Self {
29        Self::new(scripted)
30    }
31
32    /// Create a new MockProvider with the given scripted completions.
33    pub fn new(scripted: Vec<Completion>) -> Self {
34        Self {
35            scripted: Mutex::new(scripted),
36            calls: Mutex::new(Vec::new()),
37            structured_responses: Mutex::new(Vec::new()),
38        }
39    }
40
41    /// Snapshot of the transcripts the agent has sent to this provider.
42    pub fn calls(&self) -> Vec<Vec<Message>> {
43        self.calls.lock().unwrap().clone()
44    }
45
46    /// Set the queue of structured responses for `complete_structured`.
47    /// Each call to `complete_structured` pops the next response.
48    /// If the queue is empty, it returns an error (fallback path).
49    pub fn with_structured_responses(self, responses: Vec<Result<serde_json::Value>>) -> Self {
50        *self.structured_responses.lock().unwrap() = responses;
51        self
52    }
53}
54
55#[async_trait]
56impl LlmProvider for MockProvider {
57    async fn complete(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Completion> {
58        let span = tracing::info_span!("llm.complete", provider = "mock", model = "mock");
59        async move {
60            // Emit info log so tracing-test can capture the span
61            tracing::info!("mock llm call");
62            self.calls.lock().unwrap().push(messages.to_vec());
63            let mut queue = self.scripted.lock().unwrap();
64            if queue.is_empty() {
65                return Err(Error::Llm {
66                    provider: "mock".into(),
67                    message: "MockProvider: no scripted completions left".into(),
68                });
69            }
70            Ok(queue.remove(0))
71        }
72        .instrument(span)
73        .await
74    }
75
76    async fn complete_structured(&self, _req: StructuredRequest) -> Result<serde_json::Value> {
77        let mut queue = self.structured_responses.lock().unwrap();
78        if queue.is_empty() {
79            // Default: return error to trigger fallback
80            return Err(Error::Config {
81                message: "MockProvider: no structured responses configured".into(),
82            });
83        }
84        queue.remove(0)
85    }
86
87    async fn stream(
88        &self,
89        messages: &[Message],
90        tools: &[ToolSpec],
91        stream_tx: Option<StreamSender>,
92    ) -> Result<Completion> {
93        // MockProvider: just delegate to complete and emit the full content
94        let completion = self.complete(messages, tools).await?;
95        if let Some(tx) = stream_tx {
96            if !completion.content.is_empty() {
97                let _ = tx.send(completion.content.clone());
98            }
99        }
100        Ok(completion)
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[tokio::test]
109    async fn returns_scripted_in_order_and_records_calls() {
110        let provider = MockProvider::new(vec![
111            Completion {
112                content: "one".into(),
113                tool_calls: vec![],
114                finish_reason: Some("stop".into()),
115                usage: None,
116            },
117            Completion {
118                content: "two".into(),
119                tool_calls: vec![],
120                finish_reason: Some("stop".into()),
121                usage: None,
122            },
123        ]);
124
125        let m1 = vec![Message::user("hi")];
126        let r1 = provider.complete(&m1, &[]).await.unwrap();
127        assert_eq!(r1.content, "one");
128
129        let m2 = vec![Message::user("bye")];
130        let r2 = provider.complete(&m2, &[]).await.unwrap();
131        assert_eq!(r2.content, "two");
132
133        assert_eq!(provider.calls().len(), 2);
134        assert_eq!(provider.calls()[0][0].content, "hi");
135    }
136
137    #[tokio::test]
138    async fn errors_when_queue_drained() {
139        let provider = MockProvider::new(vec![]);
140        let err = provider.complete(&[], &[]).await.unwrap_err();
141        assert!(matches!(err, Error::Llm { .. }));
142    }
143}
144
145#[cfg(test)]
146mod tracing_tests {
147    use super::*;
148    use crate::llm::TokenUsage;
149    use tracing_test::traced_test;
150
151    #[traced_test]
152    #[tokio::test]
153    async fn llm_complete_records_token_fields() {
154        let usage = TokenUsage {
155            prompt_tokens: 100,
156            completion_tokens: 50,
157            total_tokens: 150,
158            cache_hit_tokens: 0,
159            cache_miss_tokens: 0,
160        };
161        let provider = MockProvider::new(vec![Completion {
162            content: "response".into(),
163            tool_calls: vec![],
164            finish_reason: Some("stop".into()),
165            usage: Some(usage),
166        }]);
167
168        provider.complete(&[], &[]).await.unwrap();
169
170        // Should have created an llm.complete span - check for span name in output
171        assert!(logs_contain("llm.complete"));
172    }
173}
recursive/llm/mock.rs

recursive/llm/
mock.rs
