af-agent 0.4.0

//! Test doubles for the Agent seams: a scripted [`ChatModel`] that replays a
//! queue of responses and records every request it received.

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

use af_llm::{ChatMessage, CompletionRequest, CompletionResponse, LlmError};
use async_trait::async_trait;
use tokio::sync::mpsc::UnboundedSender;

use crate::ChatModel;

/// Replays queued responses in order and streams each response's text as one
/// delta. An exhausted script fails the request with a decode-style
/// [`LlmError`] so a test never hangs on a missing response.
#[derive(Default)]
pub struct ScriptedModel {
    responses: Mutex<VecDeque<Result<CompletionResponse, LlmError>>>,
    requests: Mutex<Vec<CompletionRequest>>,
}

impl ScriptedModel {
    /// Model that replays `responses` in order.
    pub fn new(responses: impl IntoIterator<Item = Result<CompletionResponse, LlmError>>) -> Self {
        Self {
            responses: Mutex::new(responses.into_iter().collect()),
            requests: Mutex::new(Vec::new()),
        }
    }

    /// Script one assistant reply per message, each finishing with `stop` (or
    /// `tool_calls` when the message carries tool calls).
    pub fn replies(messages: impl IntoIterator<Item = ChatMessage>) -> Self {
        Self::new(
            messages
                .into_iter()
                .map(|message| Ok(Self::response(message))),
        )
    }

    /// Build a canonical single-choice response for `message`.
    pub fn response(message: ChatMessage) -> CompletionResponse {
        let finish_reason = if message
            .tool_calls
            .as_ref()
            .is_some_and(|calls| !calls.is_empty())
        {
            af_llm::FinishReason::ToolCalls
        } else {
            af_llm::FinishReason::Stop
        };
        CompletionResponse {
            id: "scripted".into(),
            choices: vec![af_llm::Choice {
                index: 0,
                message,
                finish_reason: Some(finish_reason),
                output_blocks: Vec::new(),
            }],
            usage: Some(af_llm::Usage {
                prompt_tokens: 1,
                completion_tokens: 1,
                total_tokens: 2,
            }),
        }
    }

    /// Responses still queued.
    pub fn remaining(&self) -> usize {
        self.responses
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .len()
    }

    /// Every request the runtime sent, in order.
    pub fn requests(&self) -> Vec<CompletionRequest> {
        self.requests
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }
}

#[async_trait]
impl ChatModel for ScriptedModel {
    async fn complete_streaming(
        &self,
        request: &CompletionRequest,
        delta_tx: UnboundedSender<(String, bool)>,
    ) -> Result<CompletionResponse, LlmError> {
        self.requests
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push(request.clone());
        let next = self
            .responses
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .pop_front();
        let response = match next {
            Some(response) => response?,
            None => return Err(LlmError::StreamProtocol("scripted model exhausted".into())),
        };
        if let Some(content) = response.first_content() {
            let _ = delta_tx.send((content.into(), response.first_tool_calls().is_some()));
        }
        Ok(response)
    }
}

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

    #[tokio::test]
    async fn scripted_model_replays_records_and_exhausts_loudly() {
        let model = ScriptedModel::replies([ChatMessage::assistant("hi")]);
        assert_eq!(model.remaining(), 1);
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        let request = CompletionRequest::new("m", vec![ChatMessage::user("hello")]);
        let response = model.complete_streaming(&request, tx).await.unwrap();
        assert_eq!(response.first_content(), Some("hi"));
        assert_eq!(rx.recv().await.unwrap(), ("hi".to_string(), false));
        assert_eq!(model.requests().len(), 1);
        assert_eq!(model.remaining(), 0);
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        assert!(matches!(
            model.complete_streaming(&request, tx).await,
            Err(LlmError::StreamProtocol(_))
        ));
        let failing = ScriptedModel::new([Err(LlmError::CircuitOpen)]);
        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        assert!(matches!(
            failing.complete_streaming(&request, tx).await,
            Err(LlmError::CircuitOpen)
        ));
    }
}