use anyhow::{anyhow, Result}; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use std::collections::HashMap; use std::fs; use std::path::Path; use std::sync::{Arc, Mutex}; #[cfg(test)] use super::base::stream_from_single_message; use super::base::{MessageStream, Provider, ProviderDef, ProviderMetadata, ProviderUsage}; use super::errors::ProviderError; use crate::conversation::message::{Message, ToolResponse}; use crate::model::ModelConfig; use futures::future::BoxFuture; use rmcp::model::{CallToolResult, Tool}; #[derive(Debug, Clone, Serialize, Deserialize)] struct TestInput { system: String, messages: Vec, tools: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] struct TestOutput { message: Message, usage: ProviderUsage, } #[derive(Debug, Clone, Serialize, Deserialize)] struct TestRecord { input: TestInput, output: TestOutput, } pub struct TestProvider { inner: Option>, records: Arc>>, file_path: String, name: String, } impl TestProvider { const PROVIDER_NAME: &str = "test"; pub fn new_recording(inner: Arc, file_path: impl Into) -> Self { Self { inner: Some(inner), records: Arc::new(Mutex::new(HashMap::new())), file_path: file_path.into(), name: Self::PROVIDER_NAME.to_string(), } } pub fn new_replaying(file_path: impl Into) -> Result { let file_path = file_path.into(); let records = Self::load_records(&file_path)?; Ok(Self { inner: None, records: Arc::new(Mutex::new(records)), file_path, name: Self::PROVIDER_NAME.to_string(), }) } pub fn finish_recording(self) -> Result<()> { if self.inner.is_some() { self.save_records()?; } Ok(()) } fn hash_input(messages: &[Message]) -> String { use crate::conversation::message::MessageContent; // Strip internal metadata (e.g. tool_meta/_meta) from content before hashing. // This metadata is used for internal routing (like goose_extension ownership) // and isn't part of the semantic input the LLM sees, so it shouldn't affect // replay matching. let stable_messages: Vec<_> = messages .iter() .map(|msg| { let mut cleaned_content: Vec<_> = msg.content.to_vec(); for content in &mut cleaned_content { match content { MessageContent::ToolRequest(ref mut req) => { req.tool_meta = None; } MessageContent::ToolResponse(ToolResponse { tool_result: Ok( ref mut result @ CallToolResult { is_error: Some(false), .. }, ), .. }) => { result.is_error = None; } _ => {} } } (msg.role.clone(), cleaned_content) }) .collect(); let serialized = serde_json::to_string(&stable_messages).unwrap_or_default(); let mut hasher = Sha256::new(); hasher.update(serialized.as_bytes()); format!("{:x}", hasher.finalize()) } fn load_records(file_path: &str) -> Result> { if !Path::new(file_path).exists() { return Ok(HashMap::new()); } let content = fs::read_to_string(file_path)?; let records: HashMap = serde_json::from_str(&content)?; Ok(records) } pub fn save_records(&self) -> Result<()> { let records = self.records.lock().unwrap(); let content = serde_json::to_string_pretty(&*records)?; fs::write(&self.file_path, content)?; Ok(()) } pub fn get_record_count(&self) -> usize { self.records.lock().unwrap().len() } } impl ProviderDef for TestProvider { type Provider = Self; fn metadata() -> ProviderMetadata { ProviderMetadata::new( Self::PROVIDER_NAME, "Test Provider", "Provider for testing that can record/replay interactions", "test-model", vec!["test-model"], "", vec![], ) } fn from_env( _model: ModelConfig, _extensions: Vec, ) -> BoxFuture<'static, Result> { Box::pin(async { Err(anyhow!("TestProvider must be constructed explicitly")) }) } } #[async_trait] impl Provider for TestProvider { fn get_name(&self) -> &str { &self.name } async fn stream( &self, model_config: &ModelConfig, session_id: &str, system: &str, messages: &[Message], tools: &[Tool], ) -> Result { let hash = Self::hash_input(messages); if let Some(inner) = &self.inner { // Call inner provider's stream and collect it let stream = inner .stream(model_config, session_id, system, messages, tools) .await?; let (message, usage) = super::base::collect_stream(stream).await?; let record = TestRecord { input: TestInput { system: system.to_string(), messages: messages.to_vec(), tools: tools.to_vec(), }, output: TestOutput { message: message.clone(), usage: usage.clone(), }, }; { let mut records = self.records.lock().unwrap(); records.insert(hash, record); } Ok(super::base::stream_from_single_message(message, usage)) } else { let records = self.records.lock().unwrap(); if let Some(record) = records.get(&hash) { let message = record.output.message.clone(); let usage = record.output.usage.clone(); Ok(super::base::stream_from_single_message(message, usage)) } else { Err(ProviderError::ExecutionError(format!( "No recorded response found for input hash: {}", hash ))) } } } fn get_model_config(&self) -> ModelConfig { ModelConfig::new_or_fail("test-model") } } #[cfg(test)] mod tests { use super::*; use crate::conversation::message::{Message, MessageContent}; use crate::providers::base::{ProviderUsage, Usage}; use chrono::Utc; use rmcp::model::{RawTextContent, Role, TextContent}; use std::env; #[derive(Clone)] struct MockProvider { model_config: ModelConfig, response: String, } #[async_trait] impl Provider for MockProvider { fn get_name(&self) -> &str { "mock-testprovider" } async fn stream( &self, _model_config: &ModelConfig, _session_id: &str, _system: &str, _messages: &[Message], _tools: &[Tool], ) -> Result { let message = Message::new( Role::Assistant, Utc::now().timestamp(), vec![MessageContent::Text(TextContent { raw: RawTextContent { text: self.response.clone(), meta: None, }, annotations: None, })], ); let usage = ProviderUsage::new("mock-model".to_string(), Usage::default()); Ok(stream_from_single_message(message, usage)) } fn get_model_config(&self) -> ModelConfig { self.model_config.clone() } } #[tokio::test] async fn test_record_and_replay() { let temp_file = format!( "{}/test_records_{}.json", env::temp_dir().display(), std::process::id() ); let mock = Arc::new(MockProvider { model_config: ModelConfig::new_or_fail("mock-model"), response: "Hello, world!".to_string(), }); { let test_provider = TestProvider::new_recording(mock, &temp_file); let model_config = test_provider.get_model_config(); let result = test_provider .complete( &model_config, "test-session-id", "You are helpful", &[], &[], ) .await; assert!(result.is_ok()); let (message, _) = result.unwrap(); if let MessageContent::Text(content) = &message.content[0] { assert_eq!(content.text, "Hello, world!"); } assert_eq!(test_provider.get_record_count(), 1); test_provider.finish_recording().unwrap(); } { let replay_provider = TestProvider::new_replaying(&temp_file).unwrap(); let model_config = replay_provider.get_model_config(); let result = replay_provider .complete( &model_config, "test-session-id", "You are helpful", &[], &[], ) .await; assert!(result.is_ok()); let (message, _) = result.unwrap(); if let MessageContent::Text(content) = &message.content[0] { assert_eq!(content.text, "Hello, world!"); } } let _ = fs::remove_file(temp_file); } #[tokio::test] async fn test_replay_missing_record() { let temp_file = format!( "{}/test_missing_{}.json", env::temp_dir().display(), std::process::id() ); let replay_provider = TestProvider::new_replaying(&temp_file).unwrap(); let model_config = replay_provider.get_model_config(); let result = replay_provider .complete( &model_config, "test-session-id", "Different system prompt", &[], &[], ) .await; assert!(result.is_err()); assert!(result .unwrap_err() .to_string() .contains("No recorded response found")); let _ = fs::remove_file(temp_file); } }