Provider scenario tests

.gitignore

@@ -58,3 +58,4 @@ do_not_version/
 /temporal.db
 /ui/desktop/src/bin/goose-scheduler-executor
 /ui/desktop/src/bin/goose
+/.env

crates/goose-cli/Cargo.toml

@@ -25,6 +25,7 @@ rmcp = { workspace = true }
 clap = { version = "4.4", features = ["derive"] }
 cliclack = "0.3.5"
 console = "0.15.8"
+dotenvy = "0.15.7"
 bat = "0.24.0"
 anyhow = "1.0"
 serde_json = "1.0"
@@ -36,17 +37,17 @@ serde_yaml = "0.9"
 tempfile = "3"
 etcetera = "0.8.0"
 reqwest = { version = "0.12.9", features = [
-        "rustls-tls-native-roots",
-        "json",
-        "cookies",
-        "gzip",
-        "brotli",
-        "deflate",
-        "zstd",
-        "charset",
-        "http2",
-        "stream"
-    ], default-features = false }
+    "rustls-tls-native-roots",
+    "json",
+    "cookies",
+    "gzip",
+    "brotli",
+    "deflate",
+    "zstd",
+    "charset",
+    "http2",
+    "stream"
+], default-features = false }
 rand = "0.8.5"
 rustyline = "15.0.0"
 tracing = "0.1"

crates/goose-cli/src/scenario_tests/message_generator.rs

@@ -0,0 +1,33 @@
+//! Message generator for scenario tests with convenience methods to
+//! just generate an image or text message.
+
+use crate::scenario_tests::scenario_runner::SCENARIO_TESTS_DIR;
+use base64::engine::general_purpose;
+use base64::Engine;
+use goose::message::Message;
+use goose::providers::base::Provider;
+
+pub type MessageGenerator<'a> = Box<dyn Fn(&dyn Provider) -> Message + 'a>;
+
+pub fn text(text: &str) -> MessageGenerator<'static> {
+    let text = text.to_string();
+    Box::new(move |_provider| Message::user().with_text(&text))
+}
+
+pub fn image(text: &str, image_name: &str) -> MessageGenerator<'static> {
+    let text = text.to_string();
+    let image_name = image_name.to_string();
+    Box::new(move |_provider| {
+        let manifest_dir = env!("CARGO_MANIFEST_DIR");
+        let image_path = format!(
+            "{}/{}/test_data/{}.jpg",
+            manifest_dir, SCENARIO_TESTS_DIR, image_name
+        );
+
+        let image_data = std::fs::read(image_path).expect("Failed to read image");
+        let base64_data = general_purpose::STANDARD.encode(&image_data);
+        Message::user()
+            .with_text(&text)
+            .with_image(base64_data, "image/jpeg")
+    })
+}

crates/goose-cli/src/scenario_tests/mock_client.rs

@@ -0,0 +1,165 @@
+//! MockClient is a mock implementation of the McpClientTrait for testing purposes.
+//! add a tool you want to have around and then add the client to the extension router
+
+use mcp_client::client::{ClientCapabilities, ClientInfo, Error, McpClientTrait};
+use mcp_core::protocol::{
+    CallToolResult, Implementation, InitializeResult, ListPromptsResult, ListResourcesResult,
+    ListToolsResult, ReadResourceResult, ServerCapabilities, ToolsCapability,
+};
+use mcp_core::{Tool, ToolError};
+use rmcp::model::{Content, GetPromptResult, ServerNotification};
+use serde_json::Value;
+use std::collections::HashMap;
+use tokio::sync::mpsc::{self, Receiver};
+
+pub struct MockClient {
+    tools: HashMap<String, Tool>,
+    handlers: HashMap<String, Box<dyn Fn(&Value) -> Result<Vec<Content>, ToolError> + Send + Sync>>,
+}
+
+impl MockClient {
+    pub(crate) fn new() -> Self {
+        Self {
+            tools: HashMap::new(),
+            handlers: HashMap::new(),
+        }
+    }
+
+    pub(crate) fn add_tool<F>(mut self, tool: Tool, handler: F) -> Self
+    where
+        F: Fn(&Value) -> Result<Vec<Content>, ToolError> + Send + Sync + 'static,
+    {
+        let tool_name = tool.name.to_string();
+        self.tools.insert(tool_name.clone(), tool);
+        self.handlers.insert(tool_name, Box::new(handler));
+        self
+    }
+}
+
+#[async_trait::async_trait]
+impl McpClientTrait for MockClient {
+    async fn initialize(
+        &mut self,
+        _: ClientInfo,
+        _: ClientCapabilities,
+    ) -> Result<InitializeResult, Error> {
+        Ok(InitializeResult {
+            protocol_version: "2024-11-05".to_string(),
+            capabilities: ServerCapabilities {
+                prompts: None,
+                resources: None,
+                tools: Some(ToolsCapability { list_changed: None }),
+            },
+            server_info: Implementation {
+                name: "MockClient".to_string(),
+                version: "1.0.0".to_string(),
+            },
+            instructions: None,
+        })
+    }
+
+    async fn list_resources(
+        &self,
+        _next_cursor: Option<String>,
+    ) -> Result<ListResourcesResult, Error> {
+        Ok(ListResourcesResult {
+            resources: vec![],
+            next_cursor: None,
+        })
+    }
+
+    async fn read_resource(&self, _uri: &str) -> Result<ReadResourceResult, Error> {
+        Err(Error::UnexpectedResponse(
+            "Resources not supported by mock client".to_string(),
+        ))
+    }
+
+    async fn list_tools(&self, _: Option<String>) -> Result<ListToolsResult, Error> {
+        let rmcp_tools: Vec<rmcp::model::Tool> = self
+            .tools
+            .values()
+            .map(|tool| {
+                let input_schema = if let serde_json::Value::Object(obj) = &tool.input_schema {
+                    std::sync::Arc::new(obj.clone())
+                } else {
+                    std::sync::Arc::new(serde_json::Map::new())
+                };
+
+                rmcp::model::Tool::new(
+                    tool.name.to_string(),
+                    tool.description.to_string(),
+                    input_schema,
+                )
+            })
+            .collect();
+
+        Ok(ListToolsResult {
+            tools: rmcp_tools,
+            next_cursor: None,
+        })
+    }
+
+    async fn call_tool(&self, name: &str, arguments: Value) -> Result<CallToolResult, Error> {
+        if let Some(handler) = self.handlers.get(name) {
+            match handler(&arguments) {
+                Ok(content) => Ok(CallToolResult {
+                    content,
+                    is_error: None,
+                }),
+                Err(e) => Err(Error::UnexpectedResponse(e.to_string())),
+            }
+        } else {
+            Err(Error::UnexpectedResponse(format!(
+                "Tool '{}' not found",
+                name
+            )))
+        }
+    }
+
+    async fn list_prompts(&self, _next_cursor: Option<String>) -> Result<ListPromptsResult, Error> {
+        Ok(ListPromptsResult { prompts: vec![] })
+    }
+
+    async fn get_prompt(&self, _name: &str, _arguments: Value) -> Result<GetPromptResult, Error> {
+        Err(Error::UnexpectedResponse(
+            "Prompts not supported by mock client".to_string(),
+        ))
+    }
+
+    async fn subscribe(&self) -> Receiver<ServerNotification> {
+        mpsc::channel(1).1
+    }
+}
+
+pub const WEATHER_TYPE: &str = "cloudy";
+
+pub fn weather_client() -> MockClient {
+    let weather_tool = Tool::new(
+        "get_weather",
+        "Get the weather for a location",
+        serde_json::json!({
+            "type": "object",
+            "required": ["location"],
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "description": "The city and state, e.g. San Francisco, CA"
+                }
+            }
+        }),
+        None, // ToolAnnotations
+    );
+
+    let mock_client = MockClient::new().add_tool(weather_tool, |args| {
+        let location = args
+            .get("location")
+            .and_then(|v| v.as_str())
+            .unwrap_or("unknown location");
+
+        Ok(vec![Content::text(format!(
+            "The weather in {} is {} and 18°C",
+            location, WEATHER_TYPE
+        ))])
+    });
+    mock_client
+}

crates/goose-cli/src/scenario_tests/mod.rs

@@ -1,106 +1,10 @@
+#[cfg(test)]
+mod message_generator;
+#[cfg(test)]
+mod mock_client;
+#[cfg(test)]
+mod provider_configs;
+#[cfg(test)]
+mod scenario_runner;
+#[cfg(test)]
 mod scenarios;
-
-use crate::session::Session;
-use anyhow::Result;
-use goose::agents::Agent;
-use goose::config::Config;
-use goose::message::Message;
-use goose::model::ModelConfig;
-use goose::providers::{create, testprovider::TestProvider};
-use std::path::Path;
-use std::sync::Arc;
-
-#[derive(Debug, Clone)]
-pub struct ScenarioResult {
-    pub messages: Vec<Message>,
-    pub error: Option<String>,
-}
-
-impl ScenarioResult {
-    pub fn message_contents(&self) -> Vec<String> {
-        self.messages
-            .iter()
-            .flat_map(|msg| &msg.content)
-            .map(|content| content.as_text().unwrap_or("").to_string())
-            .collect()
-    }
-}
-
-pub async fn run_test_scenario(test_name: &str, inputs: &[&str]) -> Result<ScenarioResult> {
-    let manifest_dir = env!("CARGO_MANIFEST_DIR");
-    let file_path = format!(
-        "{}/src/scenario_tests/recordings/{}.json",
-        manifest_dir, test_name
-    );
-
-    if let Some(parent) = Path::new(&file_path).parent() {
-        std::fs::create_dir_all(parent)?;
-    }
-
-    let replay_mode = Path::new(&file_path).exists();
-    let provider = if replay_mode {
-        match TestProvider::new_replaying(&file_path) {
-            Ok(test_provider) => {
-                Arc::new(test_provider) as Arc<dyn goose::providers::base::Provider>
-            }
-            Err(e) => {
-                let _ = std::fs::remove_file(&file_path);
-                return Err(anyhow::anyhow!(
-                    "Test replay failed for '{}': {}. File deleted - re-run test to record fresh data.",
-                    test_name, e
-                ));
-            }
-        }
-    } else {
-        if std::env::var("GITHUB_ACTIONS").is_ok() {
-            panic!(
-                "Test recording is not supported on CI. \
-            Did you forget to add the file {} to the repository and were expecting that to replay?",
-                file_path
-            );
-        }
-        let config = Config::global();
-
-        let (provider_name, model_name): (String, String) = match (
-            config.get_param::<String>("GOOSE_PROVIDER"),
-            config.get_param::<String>("GOOSE_MODEL"),
-        ) {
-            (Ok(provider), Ok(model)) => (provider, model),
-            _ => {
-                panic!("Provider or model not configured. Run 'goose configure' first");
-            }
-        };
-
-        let model_config = ModelConfig::new(model_name);
-
-        let inner_provider = create(&provider_name, model_config)?;
-        Arc::new(TestProvider::new_recording(inner_provider, &file_path))
-    };
-
-    let agent = Agent::new();
-    agent.update_provider(provider).await?;
-
-    let mut session = Session::new(agent, None, false, None, None, None, None);
-
-    let mut error = None;
-    for input in inputs {
-        if let Err(e) = session.headless(input.to_string()).await {
-            error = Some(e.to_string());
-            break;
-        }
-    }
-
-    let messages = session.message_history().to_vec();
-
-    if let Some(ref err_msg) = error {
-        if err_msg.contains("No recorded response found") {
-            let _ = std::fs::remove_file(&file_path);
-            return Err(anyhow::anyhow!(
-                "Test replay failed for '{}' - missing recorded interaction: {}. File deleted - re-run test to record fresh data.",
-                test_name, err_msg
-            ));
-        }
-    }
-
-    Ok(ScenarioResult { messages, error })
-}