diff --git a/crates/goose-cli/src/lib.rs b/crates/goose-cli/src/lib.rs index 68f2357f5ee0..ad0641b1817b 100644 --- a/crates/goose-cli/src/lib.rs +++ b/crates/goose-cli/src/lib.rs @@ -5,8 +5,10 @@ pub mod commands; pub mod logging; pub mod project_tracker; pub mod recipes; +pub mod scenario_tests; pub mod session; pub mod signal; + // Re-export commonly used types pub use session::Session; diff --git a/crates/goose-cli/src/scenario_tests/mod.rs b/crates/goose-cli/src/scenario_tests/mod.rs new file mode 100644 index 000000000000..036c07e3b761 --- /dev/null +++ b/crates/goose-cli/src/scenario_tests/mod.rs @@ -0,0 +1,106 @@ +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, + pub error: Option, +} + +impl ScenarioResult { + pub fn message_contents(&self) -> Vec { + 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 { + 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 + } + 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::("GOOSE_PROVIDER"), + config.get_param::("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 }) +} diff --git a/crates/goose-cli/src/scenario_tests/recordings/basic_greeting.json b/crates/goose-cli/src/scenario_tests/recordings/basic_greeting.json new file mode 100644 index 000000000000..a3fcbd60c7b3 --- /dev/null +++ b/crates/goose-cli/src/scenario_tests/recordings/basic_greeting.json @@ -0,0 +1,474 @@ +{ + "d2c95695a2c9ad5d95248955d850d147d76e209f537dab9bcc70fc815f2d8de7": { + "input": { + "system": "You are a general-purpose AI agent called Goose, created by Block, the parent company of Square, CashApp, and Tidal. Goose is being developed as an open-source software project.\n\nThe current date is 2025-07-22 15:02:47.\n\nGoose uses LLM providers with tool calling capability. You can be used with different language models (gpt-4o, claude-3.5-sonnet, o1, llama-3.2, deepseek-r1, etc).\nThese models have varying knowledge cut-off dates depending on when they were trained, but typically it's between 5-10 months prior to the current date.\n\n# Extensions\n\nExtensions allow other applications to provide context to Goose. Extensions connect Goose to different data sources and tools.\nYou are capable of dynamically plugging into new extensions and learning how to use them. You solve higher level problems using the tools in these extensions, and can interact with multiple at once.\nUse the search_available_extensions tool to find additional extensions to enable to help with your task. To enable extensions, use the enable_extension tool and provide the extension_name. You should only enable extensions found from the search_available_extensions tool.\n\n\nNo extensions are defined. You should let the user know that they should add extensions.\n\n\n\n# Suggestion\n\"\"\n\n\n\n\n# Response Guidelines\n\n- Use Markdown formatting for all responses.\n- Follow best practices for Markdown, including:\n - Using headers for organization.\n - Bullet points for lists.\n - Links formatted correctly, either as linked text (e.g., [this is linked text](https://example.com)) or automatic links using angle brackets (e.g., ).\n- For code examples, use fenced code blocks by placing triple backticks (` ``` `) before and after the code. Include the language identifier after the opening backticks (e.g., ` ```python `) to enable syntax highlighting.\n- Ensure clarity, conciseness, and proper formatting to enhance readability and usability.\n\n# Additional Instructions:\n\nRight now you are *NOT* in the chat only mode and have access to tool use and system.", + "messages": [ + { + "id": null, + "role": "user", + "created": 1753196567, + "content": [ + { + "type": "text", + "text": "hello" + } + ] + }, + { + "id": null, + "role": "assistant", + "created": 1753196578, + "content": [ + { + "type": "text", + "text": "Hello! I'm Goose, an AI agent created by Block. I'm here to help you with various tasks and questions.\n\nI notice that you don't currently have any extensions enabled. Extensions allow me to connect to different data sources and tools to help you accomplish more complex tasks. \n\nHere's what I can help you with right now:\n- **General conversation and questions**\n- **Task management** - I can create and execute dynamic tasks\n- **Scheduling** - I can manage scheduled recipe execution\n- **Extension management** - I can search for and enable additional capabilities\n\nTo get the most out of our interaction, you might want to:\n1. Let me search for available extensions that could help with your specific needs\n2. Tell me what kind of tasks you're looking to accomplish\n\nWhat would you like to do today?" + } + ] + }, + { + "id": null, + "role": "user", + "created": 1753196578, + "content": [ + { + "type": "text", + "text": "goodbye" + } + ] + } + ], + "tools": [ + { + "name": "platform__search_available_extensions", + "description": "Searches for additional extensions available to help complete tasks.\n Use this tool when you're unable to find a specific feature or functionality you need to complete your task, or when standard approaches aren't working.\n These extensions might provide the exact tools needed to solve your problem.\n If you find a relevant one, consider using your tools to enable it.", + "inputSchema": { + "properties": {}, + "required": [], + "type": "object" + }, + "annotations": { + "title": "Discover extensions", + "readOnlyHint": true, + "destructiveHint": false, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "platform__manage_extensions", + "description": "Tool to manage extensions and tools in goose context.\n Enable or disable extensions to help complete tasks.\n Enable or disable an extension by providing the extension name.\n ", + "inputSchema": { + "properties": { + "action": { + "description": "The action to perform", + "enum": [ + "enable", + "disable" + ], + "type": "string" + }, + "extension_name": { + "description": "The name of the extension to enable", + "type": "string" + } + }, + "required": [ + "action", + "extension_name" + ], + "type": "object" + }, + "annotations": { + "title": "Enable or disable an extension", + "readOnlyHint": false, + "destructiveHint": false, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "platform__manage_schedule", + "description": "Manage scheduled recipe execution for this Goose instance.\n\nActions:\n- \"list\": List all scheduled jobs\n- \"create\": Create a new scheduled job from a recipe file\n- \"run_now\": Execute a scheduled job immediately \n- \"pause\": Pause a scheduled job\n- \"unpause\": Resume a paused job\n- \"delete\": Remove a scheduled job\n- \"kill\": Terminate a currently running job\n- \"inspect\": Get details about a running job\n- \"sessions\": List execution history for a job\n- \"session_content\": Get the full content (messages) of a specific session\n", + "inputSchema": { + "properties": { + "action": { + "enum": [ + "list", + "create", + "run_now", + "pause", + "unpause", + "delete", + "kill", + "inspect", + "sessions", + "session_content" + ], + "type": "string" + }, + "cron_expression": { + "description": "A cron expression for create action. Supports both 5-field (minute hour day month weekday) and 6-field (second minute hour day month weekday) formats. 5-field expressions are automatically converted to 6-field by prepending '0' for seconds.", + "type": "string" + }, + "execution_mode": { + "default": "background", + "description": "Execution mode for create action: 'foreground' or 'background'", + "enum": [ + "foreground", + "background" + ], + "type": "string" + }, + "job_id": { + "description": "Job identifier for operations on existing jobs", + "type": "string" + }, + "limit": { + "default": 50, + "description": "Limit for sessions list", + "type": "integer" + }, + "recipe_path": { + "description": "Path to recipe file for create action", + "type": "string" + }, + "session_id": { + "description": "Session identifier for session_content action", + "type": "string" + } + }, + "required": [ + "action" + ], + "type": "object" + }, + "annotations": { + "title": "Manage scheduled recipes", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "dynamic_task__create_task", + "description": "Use this tool to create one or more dynamic tasks from a shared text instruction and varying parameters.How it works:\n - Provide a single text instruction\n - Use the 'task_parameters' field to pass an array of parameter sets\n - Each resulting task will use the same instruction with different parameter values\n This is useful when performing the same operation across many inputs (e.g., getting weather for multiple cities, searching multiple slack channels, iterating through various linear tickets, etc).\n Once created, these tasks should be passed to the 'subagent__execute_task' tool for execution. Tasks can run sequentially or in parallel.\n ---\n What is a 'subagent'?\n A 'subagent' is a stateless sub-process that executes a single task independently. Use subagents when:\n - You want to parallelize similar work across different inputs\n - You are not sure your search or operation will succeed on the first try\n Each subagent receives a task with a defined payload and returns a result, which is not visible to the user unless explicitly summarized by the system.\n ---\n Examples of 'task_parameters' for a single task:\n text_instruction: Search for the config file in the root directory.\n Examples of 'task_parameters' for multiple tasks:\n text_instruction: Get weather for Melbourne.\n timeout_seconds: 300\n text_instruction: Get weather for Los Angeles.\n timeout_seconds: 300\n text_instruction: Get weather for San Francisco.\n timeout_seconds: 300\n ", + "inputSchema": { + "properties": { + "task_parameters": { + "description": "Array of parameter sets for creating tasks. For a single task, provide an array with one element. For multiple tasks, provide an array with multiple elements, each with different parameter values. If there is no parameter set, provide an empty array.", + "items": { + "properties": { + "text_instruction": { + "description": "The text instruction to execute", + "type": "string" + }, + "timeout_seconds": { + "description": "Optional timeout for the task in seconds (default: 300)", + "minimum": 1, + "type": "integer" + } + }, + "required": [ + "text_instruction" + ], + "type": "object" + }, + "type": "array" + } + }, + "type": "object" + }, + "annotations": { + "title": "Dynamic Task Creation", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": true + } + }, + { + "name": "subagent__execute_task", + "description": "Only use the subagent__execute_task tool when you execute sub recipe task or dynamic task.\nEXECUTION STRATEGY DECISION:\n1. If the tasks are created with execution_mode, use the execution_mode.\n2. Execute tasks sequentially unless user explicitly requests parallel execution. PARALLEL: User uses keywords like 'parallel', 'simultaneously', 'at the same time', 'concurrently'\n\nIMPLEMENTATION:\n- Sequential execution: Call this tool multiple times, passing exactly ONE task per call\n- Parallel execution: Call this tool once, passing an ARRAY of all tasks\n\nEXAMPLES:\nUser Intent Based:\n- User: 'get weather and tell me a joke' → Sequential (2 separate tool calls, 1 task each)\n- User: 'get weather and joke in parallel' → Parallel (1 tool call with array of 2 tasks)\n- User: 'run these simultaneously' → Parallel (1 tool call with task array)\n- User: 'do task A then task B' → Sequential (2 separate tool calls)", + "inputSchema": { + "properties": { + "execution_mode": { + "default": "sequential", + "description": "Execution strategy for multiple tasks. Use 'sequential' (default) unless user explicitly requests parallel execution with words like 'parallel', 'simultaneously', 'at the same time', or 'concurrently'.", + "enum": [ + "sequential", + "parallel" + ], + "type": "string" + }, + "task_ids": { + "items": { + "description": "Unique identifier for the task", + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "task_ids" + ], + "type": "object" + }, + "annotations": { + "title": "Run tasks in parallel", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": true + } + } + ] + }, + "output": { + "message": { + "id": null, + "role": "assistant", + "created": 1753196580, + "content": [ + { + "type": "text", + "text": "Goodbye! It was nice chatting with you. Feel free to come back anytime if you need help with tasks, have questions, or want to explore what extensions might be useful for your work. Have a great day! 👋" + } + ] + }, + "usage": { + "model": "us.anthropic.claude-sonnet-4-20250514-v1:0", + "usage": { + "input_tokens": 2700, + "output_tokens": 52, + "total_tokens": 2752 + } + } + } + }, + "205407e3c76ac3acb35d9da9f560058217d3b267873f8e0a715a1946e2714ddd": { + "input": { + "system": "You are a general-purpose AI agent called Goose, created by Block, the parent company of Square, CashApp, and Tidal. Goose is being developed as an open-source software project.\n\nThe current date is 2025-07-22 15:02:47.\n\nGoose uses LLM providers with tool calling capability. You can be used with different language models (gpt-4o, claude-3.5-sonnet, o1, llama-3.2, deepseek-r1, etc).\nThese models have varying knowledge cut-off dates depending on when they were trained, but typically it's between 5-10 months prior to the current date.\n\n# Extensions\n\nExtensions allow other applications to provide context to Goose. Extensions connect Goose to different data sources and tools.\nYou are capable of dynamically plugging into new extensions and learning how to use them. You solve higher level problems using the tools in these extensions, and can interact with multiple at once.\nUse the search_available_extensions tool to find additional extensions to enable to help with your task. To enable extensions, use the enable_extension tool and provide the extension_name. You should only enable extensions found from the search_available_extensions tool.\n\n\nNo extensions are defined. You should let the user know that they should add extensions.\n\n\n\n# Suggestion\n\"\"\n\n\n\n\n# Response Guidelines\n\n- Use Markdown formatting for all responses.\n- Follow best practices for Markdown, including:\n - Using headers for organization.\n - Bullet points for lists.\n - Links formatted correctly, either as linked text (e.g., [this is linked text](https://example.com)) or automatic links using angle brackets (e.g., ).\n- For code examples, use fenced code blocks by placing triple backticks (` ``` `) before and after the code. Include the language identifier after the opening backticks (e.g., ` ```python `) to enable syntax highlighting.\n- Ensure clarity, conciseness, and proper formatting to enhance readability and usability.\n\n# Additional Instructions:\n\nRight now you are *NOT* in the chat only mode and have access to tool use and system.", + "messages": [ + { + "id": null, + "role": "user", + "created": 1753196567, + "content": [ + { + "type": "text", + "text": "hello" + } + ] + } + ], + "tools": [ + { + "name": "platform__search_available_extensions", + "description": "Searches for additional extensions available to help complete tasks.\n Use this tool when you're unable to find a specific feature or functionality you need to complete your task, or when standard approaches aren't working.\n These extensions might provide the exact tools needed to solve your problem.\n If you find a relevant one, consider using your tools to enable it.", + "inputSchema": { + "properties": {}, + "required": [], + "type": "object" + }, + "annotations": { + "title": "Discover extensions", + "readOnlyHint": true, + "destructiveHint": false, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "platform__manage_extensions", + "description": "Tool to manage extensions and tools in goose context.\n Enable or disable extensions to help complete tasks.\n Enable or disable an extension by providing the extension name.\n ", + "inputSchema": { + "properties": { + "action": { + "description": "The action to perform", + "enum": [ + "enable", + "disable" + ], + "type": "string" + }, + "extension_name": { + "description": "The name of the extension to enable", + "type": "string" + } + }, + "required": [ + "action", + "extension_name" + ], + "type": "object" + }, + "annotations": { + "title": "Enable or disable an extension", + "readOnlyHint": false, + "destructiveHint": false, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "platform__manage_schedule", + "description": "Manage scheduled recipe execution for this Goose instance.\n\nActions:\n- \"list\": List all scheduled jobs\n- \"create\": Create a new scheduled job from a recipe file\n- \"run_now\": Execute a scheduled job immediately \n- \"pause\": Pause a scheduled job\n- \"unpause\": Resume a paused job\n- \"delete\": Remove a scheduled job\n- \"kill\": Terminate a currently running job\n- \"inspect\": Get details about a running job\n- \"sessions\": List execution history for a job\n- \"session_content\": Get the full content (messages) of a specific session\n", + "inputSchema": { + "properties": { + "action": { + "enum": [ + "list", + "create", + "run_now", + "pause", + "unpause", + "delete", + "kill", + "inspect", + "sessions", + "session_content" + ], + "type": "string" + }, + "cron_expression": { + "description": "A cron expression for create action. Supports both 5-field (minute hour day month weekday) and 6-field (second minute hour day month weekday) formats. 5-field expressions are automatically converted to 6-field by prepending '0' for seconds.", + "type": "string" + }, + "execution_mode": { + "default": "background", + "description": "Execution mode for create action: 'foreground' or 'background'", + "enum": [ + "foreground", + "background" + ], + "type": "string" + }, + "job_id": { + "description": "Job identifier for operations on existing jobs", + "type": "string" + }, + "limit": { + "default": 50, + "description": "Limit for sessions list", + "type": "integer" + }, + "recipe_path": { + "description": "Path to recipe file for create action", + "type": "string" + }, + "session_id": { + "description": "Session identifier for session_content action", + "type": "string" + } + }, + "required": [ + "action" + ], + "type": "object" + }, + "annotations": { + "title": "Manage scheduled recipes", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": false + } + }, + { + "name": "dynamic_task__create_task", + "description": "Use this tool to create one or more dynamic tasks from a shared text instruction and varying parameters.How it works:\n - Provide a single text instruction\n - Use the 'task_parameters' field to pass an array of parameter sets\n - Each resulting task will use the same instruction with different parameter values\n This is useful when performing the same operation across many inputs (e.g., getting weather for multiple cities, searching multiple slack channels, iterating through various linear tickets, etc).\n Once created, these tasks should be passed to the 'subagent__execute_task' tool for execution. Tasks can run sequentially or in parallel.\n ---\n What is a 'subagent'?\n A 'subagent' is a stateless sub-process that executes a single task independently. Use subagents when:\n - You want to parallelize similar work across different inputs\n - You are not sure your search or operation will succeed on the first try\n Each subagent receives a task with a defined payload and returns a result, which is not visible to the user unless explicitly summarized by the system.\n ---\n Examples of 'task_parameters' for a single task:\n text_instruction: Search for the config file in the root directory.\n Examples of 'task_parameters' for multiple tasks:\n text_instruction: Get weather for Melbourne.\n timeout_seconds: 300\n text_instruction: Get weather for Los Angeles.\n timeout_seconds: 300\n text_instruction: Get weather for San Francisco.\n timeout_seconds: 300\n ", + "inputSchema": { + "properties": { + "task_parameters": { + "description": "Array of parameter sets for creating tasks. For a single task, provide an array with one element. For multiple tasks, provide an array with multiple elements, each with different parameter values. If there is no parameter set, provide an empty array.", + "items": { + "properties": { + "text_instruction": { + "description": "The text instruction to execute", + "type": "string" + }, + "timeout_seconds": { + "description": "Optional timeout for the task in seconds (default: 300)", + "minimum": 1, + "type": "integer" + } + }, + "required": [ + "text_instruction" + ], + "type": "object" + }, + "type": "array" + } + }, + "type": "object" + }, + "annotations": { + "title": "Dynamic Task Creation", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": true + } + }, + { + "name": "subagent__execute_task", + "description": "Only use the subagent__execute_task tool when you execute sub recipe task or dynamic task.\nEXECUTION STRATEGY DECISION:\n1. If the tasks are created with execution_mode, use the execution_mode.\n2. Execute tasks sequentially unless user explicitly requests parallel execution. PARALLEL: User uses keywords like 'parallel', 'simultaneously', 'at the same time', 'concurrently'\n\nIMPLEMENTATION:\n- Sequential execution: Call this tool multiple times, passing exactly ONE task per call\n- Parallel execution: Call this tool once, passing an ARRAY of all tasks\n\nEXAMPLES:\nUser Intent Based:\n- User: 'get weather and tell me a joke' → Sequential (2 separate tool calls, 1 task each)\n- User: 'get weather and joke in parallel' → Parallel (1 tool call with array of 2 tasks)\n- User: 'run these simultaneously' → Parallel (1 tool call with task array)\n- User: 'do task A then task B' → Sequential (2 separate tool calls)", + "inputSchema": { + "properties": { + "execution_mode": { + "default": "sequential", + "description": "Execution strategy for multiple tasks. Use 'sequential' (default) unless user explicitly requests parallel execution with words like 'parallel', 'simultaneously', 'at the same time', or 'concurrently'.", + "enum": [ + "sequential", + "parallel" + ], + "type": "string" + }, + "task_ids": { + "items": { + "description": "Unique identifier for the task", + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "task_ids" + ], + "type": "object" + }, + "annotations": { + "title": "Run tasks in parallel", + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": true + } + } + ] + }, + "output": { + "message": { + "id": null, + "role": "assistant", + "created": 1753196578, + "content": [ + { + "type": "text", + "text": "Hello! I'm Goose, an AI agent created by Block. I'm here to help you with various tasks and questions.\n\nI notice that you don't currently have any extensions enabled. Extensions allow me to connect to different data sources and tools to help you accomplish more complex tasks. \n\nHere's what I can help you with right now:\n- **General conversation and questions**\n- **Task management** - I can create and execute dynamic tasks\n- **Scheduling** - I can manage scheduled recipe execution\n- **Extension management** - I can search for and enable additional capabilities\n\nTo get the most out of our interaction, you might want to:\n1. Let me search for available extensions that could help with your specific needs\n2. Tell me what kind of tasks you're looking to accomplish\n\nWhat would you like to do today?" + } + ] + }, + "usage": { + "model": "us.anthropic.claude-sonnet-4-20250514-v1:0", + "usage": { + "input_tokens": 2517, + "output_tokens": 179, + "total_tokens": 2696 + } + } + } + } +} \ No newline at end of file diff --git a/crates/goose-cli/src/scenario_tests/scenarios.rs b/crates/goose-cli/src/scenario_tests/scenarios.rs new file mode 100644 index 000000000000..7a32df0a713a --- /dev/null +++ b/crates/goose-cli/src/scenario_tests/scenarios.rs @@ -0,0 +1,22 @@ +#[cfg(test)] +mod tests { + use crate::scenario_tests::run_test_scenario; + use anyhow::Result; + + #[tokio::test] + async fn test_basic_greeting() -> Result<()> { + let result = run_test_scenario("basic_greeting", &["hello", "goodbye"]).await?; + + assert!(result + .message_contents() + .iter() + .any(|msg| msg.contains("Hello"))); + assert!(result + .message_contents() + .iter() + .any(|msg| msg.contains("Goodbye"))); + assert!(result.error.is_none()); + + Ok(()) + } +} diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 4b155c69f277..38c810d4d171 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -24,6 +24,7 @@ pub mod openrouter; pub mod pricing; pub mod sagemaker_tgi; pub mod snowflake; +pub mod testprovider; pub mod toolshim; pub mod utils; pub mod utils_universal_openai_stream; diff --git a/crates/goose/src/providers/testprovider.rs b/crates/goose/src/providers/testprovider.rs new file mode 100644 index 000000000000..9016bdcf9612 --- /dev/null +++ b/crates/goose/src/providers/testprovider.rs @@ -0,0 +1,292 @@ +use 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}; + +use super::base::{Provider, ProviderMetadata, ProviderUsage}; +use super::errors::ProviderError; +use crate::message::Message; +use crate::model::ModelConfig; +use mcp_core::tool::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, +} + +impl TestProvider { + 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(), + } + } + + 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, + }) + } + + fn hash_input(messages: &[Message]) -> String { + let stable_messages: Vec<_> = messages + .iter() + .map(|msg| (msg.role.clone(), msg.content.clone())) + .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() + } +} + +#[async_trait] +impl Provider for TestProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "test", + "Test Provider", + "Provider for testing that can record/replay interactions", + "test-model", + vec!["test-model"], + "", + vec![], + ) + } + + async fn complete( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let hash = Self::hash_input(messages); + + if let Some(inner) = &self.inner { + // Recording mode + let (message, usage) = inner.complete(system, messages, tools).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((message, usage)) + } else { + // Replay mode + let records = self.records.lock().unwrap(); + if let Some(record) = records.get(&hash) { + Ok((record.output.message.clone(), record.output.usage.clone())) + } else { + Err(ProviderError::ExecutionError(format!( + "No recorded response found for input hash: {}", + hash + ))) + } + } + } + + fn get_model_config(&self) -> ModelConfig { + ModelConfig::new("test-model".to_string()) + } +} + +impl Drop for TestProvider { + fn drop(&mut self) { + if self.inner.is_some() { + if let Err(e) = self.save_records() { + eprintln!("Failed to save test records: {}", e); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::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 metadata() -> ProviderMetadata { + ProviderMetadata::new( + "mock", + "Mock Provider", + "Mock provider for testing", + "mock-model", + vec!["mock-model"], + "", + vec![], + ) + } + + async fn complete( + &self, + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + Ok(( + Message::new( + Role::Assistant, + Utc::now().timestamp(), + vec![MessageContent::Text(TextContent { + raw: RawTextContent { + text: self.response.clone(), + }, + annotations: None, + })], + ), + ProviderUsage::new("mock-model".to_string(), Usage::default()), + )) + } + + 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("mock-model".to_string()), + response: "Hello, world!".to_string(), + }); + + // Record phase + { + let test_provider = TestProvider::new_recording(mock, &temp_file); + + let result = test_provider.complete("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!"); + } + + test_provider.save_records().unwrap(); + assert_eq!(test_provider.get_record_count(), 1); + } + + // Replay phase + { + let replay_provider = TestProvider::new_replaying(&temp_file).unwrap(); + + let result = replay_provider.complete("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!"); + } + } + + // Cleanup + 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 result = replay_provider + .complete("Different system prompt", &[], &[]) + .await; + + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("No recorded response found")); + + // Cleanup + let _ = fs::remove_file(temp_file); + } +}