recursive/llm/

mock.rs

//! Deterministic LLM for tests and offline development.
//!
//! `MockProvider` is fed a queue of pre-baked completions. The agent treats
//! it identically to a real provider, so the agent loop is fully testable
//! without network access.
// Test infrastructure: lock().unwrap() is acceptable because mock is only
// used in test contexts and mutex poison would indicate a test bug.
#![allow(clippy::unwrap_used, clippy::expect_used)]

use std::sync::Mutex;

use async_trait::async_trait;

use super::StructuredRequest;
use super::{ChatProvider, Completion, StreamSender, ToolSpec};
use crate::error::{Error, Result};
use crate::message::Message;
use tracing::Instrument;

#[derive(Default)]
pub struct MockProvider {
    scripted: Mutex<Vec<Completion>>,
    /// Queue of pre-baked errors to return before consuming scripted completions.
    /// Each call to `complete()` pops from this queue first; if non-empty,
    /// the error is returned instead of the next scripted completion.
    errors: Mutex<Vec<Error>>,
    /// Queue of pre-baked JSON responses for `complete_structured`.
    structured_responses: Mutex<Vec<Result<serde_json::Value>>>,
    calls: Mutex<Vec<Vec<Message>>>,
    /// Optional notifier — `notify_one()`'d after every `complete()` call,
    /// before returning. Used by tests that need to deterministically race
    /// some side action (e.g. cancellation) against agent progress.
    on_complete: Option<std::sync::Arc<tokio::sync::Notify>>,
    /// Optional synchronous hook — invoked after every `complete()` call,
    /// before returning. Prefer this over `on_complete` when the side
    /// effect must happen before the agent continues (e.g. cancellation).
    on_complete_fn: Option<std::sync::Arc<dyn Fn() + Send + Sync>>,
}

impl MockProvider {
    /// Alias for `new` — both accept completions that may have `usage` set.
    /// Provided for API symmetry when callers want to be explicit about usage.
    pub fn with_usage(scripted: Vec<Completion>) -> Self {
        Self::new(scripted)
    }

    /// Create a new MockProvider with the given scripted completions.
    pub fn new(scripted: Vec<Completion>) -> Self {
        Self {
            scripted: Mutex::new(scripted),
            calls: Mutex::new(Vec::new()),
            errors: Mutex::new(Vec::new()),
            structured_responses: Mutex::new(Vec::new()),
            on_complete: None,
            on_complete_fn: None,
        }
    }

    /// Attach a notifier that fires once per `complete()` call, after the
    /// completion is dequeued and immediately before it is returned. Tests
    /// use this to deterministically synchronise external side-effects
    /// (e.g. cancellation) with agent progress, avoiding wall-clock sleeps.
    pub fn with_on_complete(mut self, notify: std::sync::Arc<tokio::sync::Notify>) -> Self {
        self.on_complete = Some(notify);
        self
    }

    /// Attach a synchronous hook that runs once per `complete()` call, after
    /// the completion is dequeued and immediately before it is returned.
    pub fn with_on_complete_fn<F>(mut self, hook: F) -> Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        self.on_complete_fn = Some(std::sync::Arc::new(hook));
        self
    }

    /// Set a queue of errors to return from `complete()` before the scripted
    /// completions are consumed. Each call to `complete()` pops the next error;
    /// when the error queue is empty, scripted completions are returned as usual.
    pub fn with_errors(self, errors: Vec<Error>) -> Self {
        *self.errors.lock().unwrap() = errors;
        self
    }

    /// Snapshot of the transcripts the agent has sent to this provider.
    pub fn calls(&self) -> Vec<Vec<Message>> {
        self.calls.lock().unwrap().clone()
    }

    /// Set the queue of structured responses for `complete_structured`.
    /// Each call to `complete_structured` pops the next response.
    /// If the queue is empty, it returns an error (fallback path).
    pub fn with_structured_responses(self, responses: Vec<Result<serde_json::Value>>) -> Self {
        *self.structured_responses.lock().unwrap() = responses;
        self
    }
}

#[async_trait]
impl ChatProvider for MockProvider {
    async fn complete(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Completion> {
        let span = tracing::info_span!("llm.complete", provider = "mock", model = "mock");
        async move {
            // Emit info log so tracing-test can capture the span
            tracing::info!("mock llm call");
            self.calls.lock().unwrap().push(messages.to_vec());
            // Check error queue first — errors are injected before scripted completions.
            {
                let mut errors = self.errors.lock().unwrap();
                if !errors.is_empty() {
                    let err = errors.remove(0);
                    return Err(err);
                }
            }
            let mut queue = self.scripted.lock().unwrap();
            if queue.is_empty() {
                return Err(Error::Llm {
                    provider: "mock".into(),
                    message: "MockProvider: no scripted completions left".into(),
                });
            }
            let completion = queue.remove(0);
            drop(queue);
            if let Some(ref hook) = self.on_complete_fn {
                hook();
            }
            if let Some(ref notify) = self.on_complete {
                notify.notify_one();
            }
            Ok(completion)
        }
        .instrument(span)
        .await
    }

    async fn complete_structured(&self, _req: StructuredRequest) -> Result<serde_json::Value> {
        let mut queue = self.structured_responses.lock().unwrap();
        if queue.is_empty() {
            // Default: return error to trigger fallback
            return Err(Error::Config {
                message: "MockProvider: no structured responses configured".into(),
            });
        }
        queue.remove(0)
    }

    async fn stream(
        &self,
        messages: &[Message],
        tools: &[ToolSpec],
        stream_tx: Option<StreamSender>,
        _cancel_token: Option<tokio_util::sync::CancellationToken>,
    ) -> Result<Completion> {
        // MockProvider: just delegate to complete and emit the full content
        let completion = self.complete(messages, tools).await?;
        if let Some(tx) = stream_tx {
            if let Some(reasoning) = &completion.reasoning_content {
                if !reasoning.is_empty() {
                    let _ = tx.send(super::StreamChunk::Reasoning(reasoning.clone()));
                }
            }
            if !completion.content.is_empty() {
                let _ = tx.send(super::StreamChunk::Text(completion.content.clone()));
            }
        }
        Ok(completion)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn returns_scripted_in_order_and_records_calls() {
        let provider = MockProvider::new(vec![
            Completion {
                content: "one".into(),
                tool_calls: vec![],
                finish_reason: Some("stop".into()),
                usage: None,
                reasoning_content: None,
            },
            Completion {
                content: "two".into(),
                tool_calls: vec![],
                finish_reason: Some("stop".into()),
                usage: None,
                reasoning_content: None,
            },
        ]);

        let m1 = vec![Message::user("hi")];
        let r1 = provider.complete(&m1, &[]).await.unwrap();
        assert_eq!(r1.content, "one");

        let m2 = vec![Message::user("bye")];
        let r2 = provider.complete(&m2, &[]).await.unwrap();
        assert_eq!(r2.content, "two");

        assert_eq!(provider.calls().len(), 2);
        assert_eq!(provider.calls()[0][0].content, "hi");
    }

    #[tokio::test]
    async fn errors_when_queue_drained() {
        let provider = MockProvider::new(vec![]);
        let err = provider.complete(&[], &[]).await.unwrap_err();
        assert!(matches!(err, Error::Llm { .. }));
    }
}

#[cfg(test)]
mod tracing_tests {
    use super::*;
    use crate::llm::TokenUsage;
    use tracing_test::traced_test;

    #[tokio::test]
    async fn errors_are_injected_before_scripted_completions() {
        // kills `if !errors.is_empty()` guard removal mutation
        use crate::error::Error;
        let provider = MockProvider::new(vec![Completion {
            content: "fallback".into(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: None,
        }])
        .with_errors(vec![Error::Llm {
            provider: "mock".into(),
            message: "injected error".into(),
        }]);

        // First call should return the injected error.
        let r1 = provider.complete(&[], &[]).await;
        assert!(r1.is_err(), "first call must return injected error");
        assert!(r1.unwrap_err().to_string().contains("injected error"));

        // Second call should return the scripted completion.
        let r2 = provider.complete(&[], &[]).await;
        assert!(r2.is_ok(), "second call must return scripted completion");
        assert_eq!(r2.unwrap().content, "fallback");
    }

    #[tokio::test]
    async fn complete_structured_returns_scripted_value() {
        // kills `queue.remove(0)` → `Ok(Default::default())` and empty-queue mutations
        use crate::llm::StructuredRequest;
        use crate::message::Message;
        let provider = MockProvider::new(vec![])
            .with_structured_responses(vec![Ok(serde_json::json!({"key": "val"}))]);
        let req = StructuredRequest {
            messages: vec![Message::user("q")],
            schema: serde_json::json!({}),
            schema_name: "test".into(),
        };
        let result = provider.complete_structured(req).await.unwrap();
        assert_eq!(
            result["key"], "val",
            "must return the scripted structured value"
        );
    }

    #[tokio::test]
    async fn complete_structured_empty_queue_returns_error() {
        // kills `if queue.is_empty() { return Err(...) }` removal mutation
        use crate::llm::StructuredRequest;
        use crate::message::Message;
        let provider = MockProvider::new(vec![]);
        // No structured responses configured — must error.
        let req = StructuredRequest {
            messages: vec![Message::user("q")],
            schema: serde_json::json!({}),
            schema_name: "test".into(),
        };
        let result = provider.complete_structured(req).await;
        assert!(result.is_err(), "empty structured queue must return error");
    }

    #[tokio::test]
    async fn stream_sends_content_to_channel() {
        // kills `if !completion.content.is_empty()` guard removal mutation
        use crate::llm::StreamChunk;
        let provider = MockProvider::new(vec![Completion {
            content: "streamed text".into(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: None,
        }]);
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        provider.stream(&[], &[], Some(tx), None).await.unwrap();
        let chunk = rx.try_recv().expect("must receive a chunk");
        assert!(
            matches!(chunk, StreamChunk::Text(t) if t == "streamed text"),
            "must receive Text chunk with completion content"
        );
    }

    #[tokio::test]
    async fn stream_sends_reasoning_before_text() {
        // kills `if !reasoning.is_empty()` guard removal mutation
        use crate::llm::StreamChunk;
        let provider = MockProvider::new(vec![Completion {
            content: "answer".into(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: None,
            reasoning_content: Some("thinking step".into()),
        }]);
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        provider.stream(&[], &[], Some(tx), None).await.unwrap();

        let first = rx.try_recv().expect("must have first chunk");
        let second = rx.try_recv().expect("must have second chunk");
        assert!(
            matches!(first, StreamChunk::Reasoning(r) if r == "thinking step"),
            "first chunk must be Reasoning"
        );
        assert!(
            matches!(second, StreamChunk::Text(t) if t == "answer"),
            "second chunk must be Text"
        );
    }

    #[tokio::test]
    async fn calls_records_all_messages() {
        // kills `self.calls.lock().unwrap().push(messages.to_vec())` removal mutation
        let provider = MockProvider::new(vec![
            Completion {
                content: "a".into(),
                tool_calls: vec![],
                finish_reason: None,
                usage: None,
                reasoning_content: None,
            },
            Completion {
                content: "b".into(),
                tool_calls: vec![],
                finish_reason: None,
                usage: None,
                reasoning_content: None,
            },
        ]);
        let msgs1 = vec![Message::user("first call")];
        let msgs2 = vec![Message::user("second call")];
        provider.complete(&msgs1, &[]).await.unwrap();
        provider.complete(&msgs2, &[]).await.unwrap();

        let calls = provider.calls();
        assert_eq!(calls.len(), 2, "must record 2 calls");
        assert_eq!(calls[0][0].content, "first call");
        assert_eq!(calls[1][0].content, "second call");
    }

    #[traced_test]
    #[tokio::test]
    async fn llm_complete_records_token_fields() {
        let usage = TokenUsage {
            reasoning_tokens: 0,
            prompt_tokens: 100,
            completion_tokens: 50,
            total_tokens: 150,
            cache_hit_tokens: 0,
            cache_miss_tokens: 0,
            // Goal 273: default 0 for MockProvider (no reasoning).
        };
        let provider = MockProvider::new(vec![Completion {
            content: "response".into(),
            tool_calls: vec![],
            finish_reason: Some("stop".into()),
            usage: Some(usage),
            reasoning_content: None,
        }]);

        provider.complete(&[], &[]).await.unwrap();

        // Should have created an llm.complete span - check for span name in output
        assert!(logs_contain("llm.complete"));
    }
}