Angela Ning
1411 lines
48 KiB
Rust
1411 lines
48 KiB
Rust
use std::sync::Arc;
|
|
|
|
use anyhow::Result;
|
|
use futures::StreamExt;
|
|
use goose::agents::{Agent, AgentEvent};
|
|
use goose::config::extensions::{set_extension, ExtensionEntry};
|
|
use goose::conversation::message::Message;
|
|
use goose::conversation::Conversation;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::Provider;
|
|
use goose::providers::{
|
|
anthropic::AnthropicProvider, azure::AzureProvider, bedrock::BedrockProvider,
|
|
databricks::DatabricksProvider, gcpvertexai::GcpVertexAIProvider, google::GoogleProvider,
|
|
ollama::OllamaProvider, openai::OpenAiProvider, openrouter::OpenRouterProvider,
|
|
xai::XaiProvider,
|
|
};
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
enum ProviderType {
|
|
Azure,
|
|
OpenAi,
|
|
#[allow(dead_code)]
|
|
Anthropic,
|
|
Bedrock,
|
|
Databricks,
|
|
GcpVertexAI,
|
|
Google,
|
|
Ollama,
|
|
OpenRouter,
|
|
Xai,
|
|
}
|
|
|
|
impl ProviderType {
|
|
fn required_env(&self) -> &'static [&'static str] {
|
|
match self {
|
|
ProviderType::Azure => &[
|
|
"AZURE_OPENAI_API_KEY",
|
|
"AZURE_OPENAI_ENDPOINT",
|
|
"AZURE_OPENAI_DEPLOYMENT_NAME",
|
|
],
|
|
ProviderType::OpenAi => &["OPENAI_API_KEY"],
|
|
ProviderType::Anthropic => &["ANTHROPIC_API_KEY"],
|
|
ProviderType::Bedrock => &["AWS_PROFILE"],
|
|
ProviderType::Databricks => &["DATABRICKS_HOST"],
|
|
ProviderType::Google => &["GOOGLE_API_KEY"],
|
|
ProviderType::Ollama => &[],
|
|
ProviderType::OpenRouter => &["OPENROUTER_API_KEY"],
|
|
ProviderType::GcpVertexAI => &["GCP_PROJECT_ID", "GCP_LOCATION"],
|
|
ProviderType::Xai => &["XAI_API_KEY"],
|
|
}
|
|
}
|
|
|
|
fn pre_check(&self) -> Result<()> {
|
|
match self {
|
|
ProviderType::Ollama => {
|
|
// Check if the `ollama ls` CLI command works
|
|
use std::process::Command;
|
|
let output = Command::new("ollama").arg("ls").output();
|
|
if let Ok(output) = output {
|
|
if output.status.success() {
|
|
return Ok(()); // CLI is running
|
|
}
|
|
}
|
|
println!("Skipping Ollama tests - `ollama ls` command not found or failed");
|
|
Err(anyhow::anyhow!("Ollama CLI is not running"))
|
|
}
|
|
_ => Ok(()), // Other providers don't need special pre-checks
|
|
}
|
|
}
|
|
|
|
async fn create_provider(&self, model_config: ModelConfig) -> Result<Arc<dyn Provider>> {
|
|
Ok(match self {
|
|
ProviderType::Azure => Arc::new(AzureProvider::from_env(model_config).await?),
|
|
ProviderType::OpenAi => Arc::new(OpenAiProvider::from_env(model_config).await?),
|
|
ProviderType::Anthropic => Arc::new(AnthropicProvider::from_env(model_config).await?),
|
|
ProviderType::Bedrock => Arc::new(BedrockProvider::from_env(model_config).await?),
|
|
ProviderType::Databricks => Arc::new(DatabricksProvider::from_env(model_config).await?),
|
|
ProviderType::GcpVertexAI => {
|
|
Arc::new(GcpVertexAIProvider::from_env(model_config).await?)
|
|
}
|
|
ProviderType::Google => Arc::new(GoogleProvider::from_env(model_config).await?),
|
|
ProviderType::Ollama => Arc::new(OllamaProvider::from_env(model_config).await?),
|
|
ProviderType::OpenRouter => Arc::new(OpenRouterProvider::from_env(model_config).await?),
|
|
ProviderType::Xai => Arc::new(XaiProvider::from_env(model_config).await?),
|
|
})
|
|
}
|
|
}
|
|
|
|
pub fn check_required_env_vars(required_vars: &[&str]) -> Result<()> {
|
|
let missing_vars: Vec<&str> = required_vars
|
|
.iter()
|
|
.filter(|&&var| std::env::var(var).is_err())
|
|
.cloned()
|
|
.collect();
|
|
|
|
if !missing_vars.is_empty() {
|
|
println!(
|
|
"Skipping tests. Missing environment variables: {:?}",
|
|
missing_vars
|
|
);
|
|
return Err(anyhow::anyhow!("Required environment variables not set"));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn run_truncate_test(
|
|
provider_type: ProviderType,
|
|
model: &str,
|
|
context_window: usize,
|
|
) -> Result<()> {
|
|
let model_config = ModelConfig::new(model)
|
|
.unwrap()
|
|
.with_context_limit(Some(context_window))
|
|
.with_temperature(Some(0.0));
|
|
let provider = provider_type.create_provider(model_config).await?;
|
|
|
|
let agent = Agent::new();
|
|
agent.update_provider(provider).await?;
|
|
let repeat_count = context_window + 10_000;
|
|
let large_message_content = "hello ".repeat(repeat_count);
|
|
let conversation = Conversation::new(vec![
|
|
Message::user().with_text("hi there. what is 2 + 2?"),
|
|
Message::assistant().with_text("hey! I think it's 4."),
|
|
Message::user().with_text(&large_message_content),
|
|
Message::assistant().with_text("heyy!!"),
|
|
Message::user().with_text("what's the meaning of life?"),
|
|
Message::assistant().with_text("the meaning of life is 42"),
|
|
Message::user().with_text(
|
|
"did I ask you what's 2+2 in this message history? just respond with 'yes' or 'no'",
|
|
),
|
|
])
|
|
.unwrap();
|
|
|
|
let reply_stream = agent.reply(conversation, None, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
let mut responses = Vec::new();
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(response)) => responses.push(response),
|
|
Ok(AgentEvent::McpNotification(n)) => {
|
|
println!("MCP Notification: {n:?}");
|
|
}
|
|
Ok(AgentEvent::ModelChange { .. }) => {
|
|
// Model change events are informational, just continue
|
|
}
|
|
Ok(AgentEvent::HistoryReplaced(_updated_conversation)) => {
|
|
// Should update the conversation here, but we're not reading it
|
|
}
|
|
Err(e) => {
|
|
println!("Error: {:?}", e);
|
|
return Err(e);
|
|
}
|
|
}
|
|
}
|
|
|
|
println!("Responses: {responses:?}\n");
|
|
|
|
// Ollama and OpenRouter truncate by default even when the context window is exceeded
|
|
// We don't have control over the truncation behavior in these providers.
|
|
// Skip the strict assertions for these providers.
|
|
if provider_type == ProviderType::Ollama || provider_type == ProviderType::OpenRouter {
|
|
println!(
|
|
"WARNING: Skipping test for {:?} because it truncates by default when the context window is exceeded",
|
|
provider_type
|
|
);
|
|
return Ok(());
|
|
}
|
|
|
|
assert_eq!(responses.len(), 1);
|
|
|
|
assert_eq!(responses[0].content.len(), 1);
|
|
|
|
match responses[0].content[0] {
|
|
goose::conversation::message::MessageContent::Text(ref text_content) => {
|
|
assert!(text_content.text.to_lowercase().contains("no"));
|
|
assert!(!text_content.text.to_lowercase().contains("yes"));
|
|
}
|
|
_ => {
|
|
panic!(
|
|
"Unexpected message content type: {:?}",
|
|
responses[0].content[0]
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[derive(Debug)]
|
|
struct TestConfig {
|
|
provider_type: ProviderType,
|
|
model: &'static str,
|
|
context_window: usize,
|
|
}
|
|
|
|
async fn run_test_with_config(config: TestConfig) -> Result<()> {
|
|
println!("Starting test for {config:?}");
|
|
|
|
// Check for required environment variables
|
|
if check_required_env_vars(config.provider_type.required_env()).is_err() {
|
|
return Ok(()); // Skip test if env vars are missing
|
|
}
|
|
|
|
// Run provider-specific pre-checks
|
|
if config.provider_type.pre_check().is_err() {
|
|
return Ok(()); // Skip test if pre-check fails
|
|
}
|
|
|
|
// Run the truncate test
|
|
run_truncate_test(config.provider_type, config.model, config.context_window).await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_openai() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::OpenAi,
|
|
model: "o3-mini-low",
|
|
context_window: 200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_anthropic() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Anthropic,
|
|
model: "claude-sonnet-4",
|
|
context_window: 200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_azure() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Azure,
|
|
model: "gpt-4o-mini",
|
|
context_window: 128_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_bedrock() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Bedrock,
|
|
model: "anthropic.claude-sonnet-4-20250514:0",
|
|
context_window: 200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_databricks() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Databricks,
|
|
model: "databricks-meta-llama-3-3-70b-instruct",
|
|
context_window: 128_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_databricks_bedrock() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Databricks,
|
|
model: "claude-sonnet-4",
|
|
context_window: 200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_databricks_openai() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Databricks,
|
|
model: "gpt-4o-mini",
|
|
context_window: 128_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_google() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Google,
|
|
model: "gemini-2.0-flash-exp",
|
|
context_window: 1_200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_openrouter() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::OpenRouter,
|
|
model: "deepseek/deepseek-r1",
|
|
context_window: 130_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_ollama() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Ollama,
|
|
model: "llama3.2",
|
|
context_window: 128_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_gcpvertexai() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::GcpVertexAI,
|
|
model: "claude-sonnet-4-20250514",
|
|
context_window: 200_000,
|
|
})
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_agent_with_xai() -> Result<()> {
|
|
run_test_with_config(TestConfig {
|
|
provider_type: ProviderType::Xai,
|
|
model: "grok-3",
|
|
context_window: 9_000,
|
|
})
|
|
.await
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod schedule_tool_tests {
|
|
use super::*;
|
|
use async_trait::async_trait;
|
|
use chrono::{DateTime, Utc};
|
|
use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME;
|
|
use goose::scheduler::{ScheduledJob, SchedulerError};
|
|
use goose::scheduler_trait::SchedulerTrait;
|
|
use goose::session::Session;
|
|
use std::sync::Arc;
|
|
|
|
struct MockScheduler {
|
|
jobs: tokio::sync::Mutex<Vec<ScheduledJob>>,
|
|
}
|
|
|
|
impl MockScheduler {
|
|
fn new() -> Self {
|
|
Self {
|
|
jobs: tokio::sync::Mutex::new(Vec::new()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl SchedulerTrait for MockScheduler {
|
|
async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> {
|
|
let mut jobs = self.jobs.lock().await;
|
|
jobs.push(job);
|
|
Ok(())
|
|
}
|
|
|
|
async fn list_scheduled_jobs(&self) -> Result<Vec<ScheduledJob>, SchedulerError> {
|
|
let jobs = self.jobs.lock().await;
|
|
Ok(jobs.clone())
|
|
}
|
|
|
|
async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> {
|
|
let mut jobs = self.jobs.lock().await;
|
|
if let Some(pos) = jobs.iter().position(|job| job.id == id) {
|
|
jobs.remove(pos);
|
|
Ok(())
|
|
} else {
|
|
Err(SchedulerError::JobNotFound(id.to_string()))
|
|
}
|
|
}
|
|
|
|
async fn pause_schedule(&self, _id: &str) -> Result<(), SchedulerError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn unpause_schedule(&self, _id: &str) -> Result<(), SchedulerError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn run_now(&self, _id: &str) -> Result<String, SchedulerError> {
|
|
Ok("test_session_123".to_string())
|
|
}
|
|
|
|
async fn sessions(
|
|
&self,
|
|
_sched_id: &str,
|
|
_limit: usize,
|
|
) -> Result<Vec<(String, Session)>, SchedulerError> {
|
|
Ok(vec![])
|
|
}
|
|
|
|
async fn update_schedule(
|
|
&self,
|
|
_sched_id: &str,
|
|
_new_cron: String,
|
|
) -> Result<(), SchedulerError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn kill_running_job(&self, _sched_id: &str) -> Result<(), SchedulerError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn get_running_job_info(
|
|
&self,
|
|
_sched_id: &str,
|
|
) -> Result<Option<(String, DateTime<Utc>)>, SchedulerError> {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_schedule_management_tool_list() {
|
|
let agent = Agent::new();
|
|
let mock_scheduler = Arc::new(MockScheduler::new());
|
|
agent.set_scheduler(mock_scheduler.clone()).await;
|
|
|
|
// Test that the schedule management tool is available in the tools list
|
|
let tools = agent.list_tools(None).await;
|
|
let schedule_tool = tools
|
|
.iter()
|
|
.find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME);
|
|
assert!(schedule_tool.is_some());
|
|
|
|
let tool = schedule_tool.unwrap();
|
|
assert!(tool
|
|
.description
|
|
.clone()
|
|
.unwrap_or_default()
|
|
.contains("Manage scheduled recipe execution"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_schedule_management_tool_no_scheduler() {
|
|
let agent = Agent::new();
|
|
// Don't set scheduler - test that the tool still appears in the list
|
|
// but would fail if actually called (which we can't test directly through public API)
|
|
|
|
let tools = agent.list_tools(None).await;
|
|
let schedule_tool = tools
|
|
.iter()
|
|
.find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME);
|
|
assert!(schedule_tool.is_some());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_schedule_management_tool_in_platform_tools() {
|
|
let agent = Agent::new();
|
|
let tools = agent.list_tools(Some("platform".to_string())).await;
|
|
|
|
// Check that the schedule management tool is included in platform tools
|
|
let schedule_tool = tools
|
|
.iter()
|
|
.find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME);
|
|
assert!(schedule_tool.is_some());
|
|
|
|
let tool = schedule_tool.unwrap();
|
|
assert!(tool
|
|
.description
|
|
.clone()
|
|
.unwrap_or_default()
|
|
.contains("Manage scheduled recipe execution"));
|
|
|
|
// Verify the tool has the expected actions in its schema
|
|
if let Some(properties) = tool.input_schema.get("properties") {
|
|
if let Some(action_prop) = properties.get("action") {
|
|
if let Some(enum_values) = action_prop.get("enum") {
|
|
let actions: Vec<String> = enum_values
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.map(|v| v.as_str().unwrap().to_string())
|
|
.collect();
|
|
|
|
// Check that our session_content action is included
|
|
assert!(actions.contains(&"session_content".to_string()));
|
|
assert!(actions.contains(&"list".to_string()));
|
|
assert!(actions.contains(&"create".to_string()));
|
|
assert!(actions.contains(&"sessions".to_string()));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_schedule_management_tool_schema_validation() {
|
|
let agent = Agent::new();
|
|
let tools = agent.list_tools(None).await;
|
|
let schedule_tool = tools
|
|
.iter()
|
|
.find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME);
|
|
assert!(schedule_tool.is_some());
|
|
|
|
let tool = schedule_tool.unwrap();
|
|
|
|
// Verify the tool schema has the session_id parameter for session_content action
|
|
if let Some(properties) = tool.input_schema.get("properties") {
|
|
assert!(properties.get("session_id").is_some());
|
|
|
|
if let Some(session_id_prop) = properties.get("session_id") {
|
|
assert_eq!(
|
|
session_id_prop.get("type").unwrap().as_str().unwrap(),
|
|
"string"
|
|
);
|
|
assert!(session_id_prop
|
|
.get("description")
|
|
.unwrap()
|
|
.as_str()
|
|
.unwrap()
|
|
.contains("Session identifier for session_content action"));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod final_output_tool_tests {
|
|
use super::*;
|
|
use futures::stream;
|
|
use goose::agents::final_output_tool::FINAL_OUTPUT_TOOL_NAME;
|
|
use goose::conversation::Conversation;
|
|
use goose::providers::base::MessageStream;
|
|
use goose::recipe::Response;
|
|
use rmcp::model::CallToolRequestParam;
|
|
use rmcp::object;
|
|
|
|
#[tokio::test]
|
|
async fn test_final_output_assistant_message_in_reply() -> Result<()> {
|
|
use async_trait::async_trait;
|
|
use goose::conversation::message::Message;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::{Provider, ProviderUsage, Usage};
|
|
use goose::providers::errors::ProviderError;
|
|
use rmcp::model::Tool;
|
|
|
|
#[derive(Clone)]
|
|
struct MockProvider {
|
|
model_config: ModelConfig,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for MockProvider {
|
|
fn metadata() -> goose::providers::base::ProviderMetadata {
|
|
goose::providers::base::ProviderMetadata::empty()
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model_config.clone()
|
|
}
|
|
|
|
async fn complete(
|
|
&self,
|
|
_system: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
Ok((
|
|
Message::assistant().with_text("Task completed."),
|
|
ProviderUsage::new("mock".to_string(), Usage::default()),
|
|
))
|
|
}
|
|
|
|
async fn complete_with_model(
|
|
&self,
|
|
_model_config: &ModelConfig,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
self.complete(system, messages, tools).await
|
|
}
|
|
}
|
|
|
|
let agent = Agent::new();
|
|
|
|
let model_config = ModelConfig::new("test-model").unwrap();
|
|
let mock_provider = Arc::new(MockProvider { model_config });
|
|
agent.update_provider(mock_provider).await?;
|
|
|
|
let response = Response {
|
|
json_schema: Some(serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"result": {"type": "string"}
|
|
},
|
|
"required": ["result"]
|
|
})),
|
|
};
|
|
agent.add_final_output_tool(response).await;
|
|
|
|
// Simulate a final output tool call occurring.
|
|
let tool_call = CallToolRequestParam {
|
|
name: FINAL_OUTPUT_TOOL_NAME.into(),
|
|
arguments: Some(object!({
|
|
"result": "Test output"
|
|
})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(tool_call, "request_id".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "Tool call should succeed");
|
|
let final_result = result.unwrap().result.await;
|
|
assert!(final_result.is_ok(), "Tool execution should succeed");
|
|
|
|
let content = final_result.unwrap();
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("Final output successfully collected."),
|
|
"Tool result missing expected content: {}",
|
|
text.text
|
|
);
|
|
|
|
// Simulate the reply stream continuing after the final output tool call.
|
|
let reply_stream = agent.reply(Conversation::empty(), None, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
let mut responses = Vec::new();
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(response)) => responses.push(response),
|
|
Ok(_) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
assert!(!responses.is_empty(), "Should have received responses");
|
|
let last_message = responses.last().unwrap();
|
|
|
|
// Check that the last message is an assistant message with our final output
|
|
assert_eq!(last_message.role, rmcp::model::Role::Assistant);
|
|
let message_text = last_message.as_concat_text();
|
|
assert_eq!(message_text, r#"{"result":"Test output"}"#);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_when_final_output_not_called_in_reply() -> Result<()> {
|
|
use async_trait::async_trait;
|
|
use goose::agents::final_output_tool::FINAL_OUTPUT_CONTINUATION_MESSAGE;
|
|
use goose::conversation::message::Message;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::{Provider, ProviderUsage};
|
|
use goose::providers::errors::ProviderError;
|
|
use rmcp::model::Tool;
|
|
|
|
#[derive(Clone)]
|
|
struct MockProvider {
|
|
model_config: ModelConfig,
|
|
stream_round: std::sync::Arc<std::sync::Mutex<i32>>,
|
|
got_continuation_message: std::sync::Arc<std::sync::Mutex<bool>>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for MockProvider {
|
|
fn metadata() -> goose::providers::base::ProviderMetadata {
|
|
goose::providers::base::ProviderMetadata::empty()
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model_config.clone()
|
|
}
|
|
|
|
fn supports_streaming(&self) -> bool {
|
|
true
|
|
}
|
|
|
|
async fn stream(
|
|
&self,
|
|
_system: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
if let Some(last_msg) = _messages.last() {
|
|
for content in &last_msg.content {
|
|
if let goose::conversation::message::MessageContent::Text(text_content) =
|
|
content
|
|
{
|
|
if text_content.text == FINAL_OUTPUT_CONTINUATION_MESSAGE {
|
|
let mut got_continuation =
|
|
self.got_continuation_message.lock().unwrap();
|
|
*got_continuation = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut round = self.stream_round.lock().unwrap();
|
|
*round += 1;
|
|
|
|
let deltas = if *round == 1 {
|
|
vec![
|
|
Ok((Some(Message::assistant().with_text("Hello")), None)),
|
|
Ok((Some(Message::assistant().with_text("Hi!")), None)),
|
|
Ok((
|
|
Some(Message::assistant().with_text("What is the final output?")),
|
|
None,
|
|
)),
|
|
]
|
|
} else {
|
|
vec![Ok((
|
|
Some(Message::assistant().with_text("Additional random delta")),
|
|
None,
|
|
))]
|
|
};
|
|
|
|
let stream = stream::iter(deltas.into_iter());
|
|
Ok(Box::pin(stream))
|
|
}
|
|
|
|
async fn complete(
|
|
&self,
|
|
_system: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
Err(ProviderError::NotImplemented("Not implemented".to_string()))
|
|
}
|
|
|
|
async fn complete_with_model(
|
|
&self,
|
|
_model_config: &ModelConfig,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
self.complete(system, messages, tools).await
|
|
}
|
|
}
|
|
|
|
let agent = Agent::new();
|
|
|
|
let model_config = ModelConfig::new("test-model").unwrap();
|
|
let mock_provider = Arc::new(MockProvider {
|
|
model_config,
|
|
stream_round: std::sync::Arc::new(std::sync::Mutex::new(0)),
|
|
got_continuation_message: std::sync::Arc::new(std::sync::Mutex::new(false)),
|
|
});
|
|
let mock_provider_clone = mock_provider.clone();
|
|
agent.update_provider(mock_provider).await?;
|
|
|
|
let response = Response {
|
|
json_schema: Some(serde_json::json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"result": {"type": "string"}
|
|
},
|
|
"required": ["result"]
|
|
})),
|
|
};
|
|
agent.add_final_output_tool(response).await;
|
|
|
|
// Simulate the reply stream being called.
|
|
let reply_stream = agent.reply(Conversation::empty(), None, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
let mut responses = Vec::new();
|
|
let mut count = 0;
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(response)) => {
|
|
responses.push(response);
|
|
count += 1;
|
|
if count >= 4 {
|
|
// Limit to 4 messages to avoid infinite loop due to mock provider
|
|
break;
|
|
}
|
|
}
|
|
Ok(_) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
assert!(!responses.is_empty(), "Should have received responses");
|
|
let last_message = responses.last().unwrap();
|
|
|
|
// Check that the first 3 messages do not have FINAL_OUTPUT_CONTINUATION_MESSAGE
|
|
for (i, response) in responses.iter().take(3).enumerate() {
|
|
let message_text = response.as_concat_text();
|
|
assert_ne!(
|
|
message_text,
|
|
FINAL_OUTPUT_CONTINUATION_MESSAGE,
|
|
"Message {} should not be the continuation message, got: '{}'",
|
|
i + 1,
|
|
message_text
|
|
);
|
|
}
|
|
|
|
// Check that the last message after the llm stream is the message directing the agent to continue
|
|
assert_eq!(last_message.role, rmcp::model::Role::User);
|
|
let message_text = last_message.as_concat_text();
|
|
assert_eq!(message_text, FINAL_OUTPUT_CONTINUATION_MESSAGE);
|
|
|
|
// Continue streaming to consume any remaining content, this lets us verify the provider saw the continuation message
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(_response)) => {
|
|
break; // Stop after receiving the next message
|
|
}
|
|
Ok(_) => {}
|
|
Err(e) => {
|
|
println!("Error while streaming remaining content: {:?}", e);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Assert that the provider received the continuation message
|
|
let got_continuation = mock_provider_clone.got_continuation_message.lock().unwrap();
|
|
assert!(
|
|
*got_continuation,
|
|
"Provider should have received the FINAL_OUTPUT_CONTINUATION_MESSAGE"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod retry_tests {
|
|
use super::*;
|
|
use async_trait::async_trait;
|
|
use goose::agents::types::{RetryConfig, SuccessCheck};
|
|
use goose::conversation::message::Message;
|
|
use goose::conversation::Conversation;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::{Provider, ProviderUsage, Usage};
|
|
use goose::providers::errors::ProviderError;
|
|
use rmcp::model::Tool;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Clone)]
|
|
struct MockRetryProvider {
|
|
model_config: ModelConfig,
|
|
call_count: Arc<AtomicUsize>,
|
|
fail_until: usize,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for MockRetryProvider {
|
|
fn metadata() -> goose::providers::base::ProviderMetadata {
|
|
goose::providers::base::ProviderMetadata::empty()
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model_config.clone()
|
|
}
|
|
|
|
async fn complete(
|
|
&self,
|
|
_system: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
let count = self.call_count.fetch_add(1, Ordering::SeqCst);
|
|
|
|
if count < self.fail_until {
|
|
Ok((
|
|
Message::assistant().with_text("Task failed - will retry."),
|
|
ProviderUsage::new("mock".to_string(), Usage::default()),
|
|
))
|
|
} else {
|
|
Ok((
|
|
Message::assistant().with_text("Task completed successfully."),
|
|
ProviderUsage::new("mock".to_string(), Usage::default()),
|
|
))
|
|
}
|
|
}
|
|
|
|
async fn complete_with_model(
|
|
&self,
|
|
_model_config: &ModelConfig,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
self.complete(system, messages, tools).await
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_retry_config_validation_integration() -> Result<()> {
|
|
let agent = Agent::new();
|
|
|
|
let model_config = ModelConfig::new("test-model").unwrap();
|
|
let mock_provider = Arc::new(MockRetryProvider {
|
|
model_config,
|
|
call_count: Arc::new(AtomicUsize::new(0)),
|
|
fail_until: 0,
|
|
});
|
|
agent.update_provider(mock_provider.clone()).await?;
|
|
|
|
let retry_config = RetryConfig {
|
|
max_retries: 3,
|
|
checks: vec![SuccessCheck::Shell {
|
|
command: "echo 'success check'".to_string(),
|
|
}],
|
|
on_failure: Some("echo 'cleanup executed'".to_string()),
|
|
timeout_seconds: Some(30),
|
|
on_failure_timeout_seconds: Some(60),
|
|
};
|
|
|
|
assert!(
|
|
retry_config.validate().is_ok(),
|
|
"Valid config should pass validation"
|
|
);
|
|
|
|
let conversation =
|
|
Conversation::new(vec![Message::user().with_text("Complete this task")]).unwrap();
|
|
|
|
let reply_stream = agent.reply(conversation, None, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
let mut responses = Vec::new();
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(response)) => responses.push(response),
|
|
Ok(_) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
assert!(!responses.is_empty(), "Should have received responses");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_retry_success_check_execution() -> Result<()> {
|
|
use goose::agents::retry::execute_success_checks;
|
|
|
|
let retry_config = RetryConfig {
|
|
max_retries: 3,
|
|
checks: vec![],
|
|
on_failure: None,
|
|
timeout_seconds: Some(30),
|
|
on_failure_timeout_seconds: Some(60),
|
|
};
|
|
|
|
let success_checks = vec![SuccessCheck::Shell {
|
|
command: "echo 'test'".to_string(),
|
|
}];
|
|
|
|
let result = execute_success_checks(&success_checks, &retry_config).await;
|
|
assert!(result.is_ok(), "Success check should pass");
|
|
assert!(result.unwrap(), "Command should succeed");
|
|
|
|
let fail_checks = vec![SuccessCheck::Shell {
|
|
command: "false".to_string(),
|
|
}];
|
|
|
|
let result = execute_success_checks(&fail_checks, &retry_config).await;
|
|
assert!(result.is_ok(), "Success check execution should not error");
|
|
assert!(!result.unwrap(), "Command should fail");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_retry_logic_with_validation_errors() -> Result<()> {
|
|
let invalid_retry_config = RetryConfig {
|
|
max_retries: 0,
|
|
checks: vec![],
|
|
on_failure: None,
|
|
timeout_seconds: Some(0),
|
|
on_failure_timeout_seconds: None,
|
|
};
|
|
|
|
let validation_result = invalid_retry_config.validate();
|
|
assert!(
|
|
validation_result.is_err(),
|
|
"Should validate max_retries > 0"
|
|
);
|
|
assert!(validation_result
|
|
.unwrap_err()
|
|
.contains("max_retries must be greater than 0"));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_retry_attempts_counter_reset() -> Result<()> {
|
|
let agent = Agent::new();
|
|
|
|
agent.reset_retry_attempts().await;
|
|
let initial_attempts = agent.get_retry_attempts().await;
|
|
assert_eq!(initial_attempts, 0);
|
|
|
|
let new_attempts = agent.increment_retry_attempts().await;
|
|
assert_eq!(new_attempts, 1);
|
|
|
|
agent.reset_retry_attempts().await;
|
|
let reset_attempts = agent.get_retry_attempts().await;
|
|
assert_eq!(reset_attempts, 0);
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod max_turns_tests {
|
|
use super::*;
|
|
use async_trait::async_trait;
|
|
use goose::conversation::message::{Message, MessageContent};
|
|
use goose::conversation::Conversation;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
|
|
use goose::providers::errors::ProviderError;
|
|
use rmcp::model::{CallToolRequestParam, Tool};
|
|
use rmcp::object;
|
|
|
|
struct MockToolProvider {}
|
|
|
|
impl MockToolProvider {
|
|
fn new() -> Self {
|
|
Self {}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for MockToolProvider {
|
|
async fn complete(
|
|
&self,
|
|
_system_prompt: &str,
|
|
_messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
let tool_call = CallToolRequestParam {
|
|
name: "test_tool".into(),
|
|
arguments: Some(object!({"param": "value"})),
|
|
};
|
|
let message = Message::assistant().with_tool_request("call_123", Ok(tool_call));
|
|
|
|
let usage = ProviderUsage::new(
|
|
"mock-model".to_string(),
|
|
Usage::new(Some(10), Some(5), Some(15)),
|
|
);
|
|
|
|
Ok((message, usage))
|
|
}
|
|
|
|
async fn complete_with_model(
|
|
&self,
|
|
_model_config: &ModelConfig,
|
|
system_prompt: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
|
|
self.complete(system_prompt, messages, tools).await
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
ModelConfig::new("mock-model").unwrap()
|
|
}
|
|
|
|
fn metadata() -> ProviderMetadata {
|
|
ProviderMetadata {
|
|
name: "mock".to_string(),
|
|
display_name: "Mock Provider".to_string(),
|
|
description: "Mock provider for testing".to_string(),
|
|
default_model: "mock-model".to_string(),
|
|
known_models: vec![],
|
|
model_doc_link: "".to_string(),
|
|
config_keys: vec![],
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_max_turns_limit() -> Result<()> {
|
|
let agent = Agent::new();
|
|
let provider = Arc::new(MockToolProvider::new());
|
|
agent.update_provider(provider).await?;
|
|
// The mock provider will call a non-existent tool, which will fail and allow the loop to continue
|
|
let conversation = Conversation::new(vec![Message::user().with_text("Hello")]).unwrap();
|
|
|
|
let reply_stream = agent.reply(conversation, None, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
let mut responses = Vec::new();
|
|
while let Some(response_result) = reply_stream.next().await {
|
|
match response_result {
|
|
Ok(AgentEvent::Message(response)) => {
|
|
if let Some(MessageContent::ToolConfirmationRequest(ref req)) =
|
|
response.content.first()
|
|
{
|
|
agent.handle_confirmation(
|
|
req.id.clone(),
|
|
goose::permission::PermissionConfirmation {
|
|
principal_type: goose::permission::permission_confirmation::PrincipalType::Tool,
|
|
permission: goose::permission::Permission::AllowOnce,
|
|
}
|
|
).await;
|
|
}
|
|
responses.push(response);
|
|
}
|
|
Ok(AgentEvent::McpNotification(_)) => {}
|
|
Ok(AgentEvent::ModelChange { .. }) => {}
|
|
Ok(AgentEvent::HistoryReplaced(_updated_conversation)) => {
|
|
// We should update the conversation here, but we're not reading it
|
|
}
|
|
Err(e) => {
|
|
return Err(e);
|
|
}
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
!responses.is_empty(),
|
|
"Expected at least 1 response, got {}",
|
|
responses.len()
|
|
);
|
|
|
|
// Look for the max turns message as the last response
|
|
let last_response = responses.last().unwrap();
|
|
let last_content = last_response.content.first().unwrap();
|
|
if let MessageContent::Text(text_content) = last_content {
|
|
assert!(text_content.text.contains(
|
|
"I've reached the maximum number of actions I can do without user input"
|
|
));
|
|
} else {
|
|
panic!("Expected text content in last message");
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod extension_manager_tests {
|
|
use super::*;
|
|
use goose::agents::extension::{ExtensionConfig, PlatformExtensionContext};
|
|
use goose::agents::extension_manager_extension::{
|
|
MANAGE_EXTENSIONS_TOOL_NAME, SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
|
|
};
|
|
use rmcp::model::CallToolRequestParam;
|
|
use rmcp::object;
|
|
|
|
async fn setup_agent_with_extension_manager() -> Agent {
|
|
// Add the TODO extension to the config so it can be discovered by search_available_extensions
|
|
// Set it as disabled initially so tests can enable it
|
|
let todo_extension_entry = ExtensionEntry {
|
|
enabled: false,
|
|
config: ExtensionConfig::Platform {
|
|
name: "todo".to_string(),
|
|
description:
|
|
"Enable a todo list for Goose so it can keep track of what it is doing"
|
|
.to_string(),
|
|
bundled: Some(true),
|
|
available_tools: vec![],
|
|
},
|
|
};
|
|
set_extension(todo_extension_entry);
|
|
|
|
let agent = Agent::new();
|
|
|
|
agent
|
|
.extension_manager
|
|
.set_context(PlatformExtensionContext {
|
|
session_id: Some("test_session".to_string()),
|
|
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
|
|
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
|
|
})
|
|
.await;
|
|
|
|
// Now add the extension manager platform extension
|
|
let ext_config = ExtensionConfig::Platform {
|
|
name: "extensionmanager".to_string(),
|
|
description: "Extension Manager".to_string(),
|
|
bundled: Some(true),
|
|
available_tools: vec![],
|
|
};
|
|
|
|
agent
|
|
.add_extension(ext_config)
|
|
.await
|
|
.expect("Failed to add extension manager");
|
|
agent
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_extension_manager_tools_available() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
let tools = agent.list_tools(None).await;
|
|
|
|
// Note: Tool names are prefixed with the normalized extension name "extensionmanager"
|
|
// not the display name "Extension Manager"
|
|
let search_tool = tools.iter().find(|tool| {
|
|
tool.name == format!("extensionmanager__{SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME}")
|
|
});
|
|
assert!(
|
|
search_tool.is_some(),
|
|
"search_available_extensions tool should be available"
|
|
);
|
|
|
|
let manage_tool = tools
|
|
.iter()
|
|
.find(|tool| tool.name == format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}"));
|
|
assert!(
|
|
manage_tool.is_some(),
|
|
"manage_extensions tool should be available"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_search_available_extensions_tool_call() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
|
|
let tool_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(tool_call, "request_1".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "search_available_extensions should succeed");
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"search_available_extensions execution should succeed"
|
|
);
|
|
|
|
let content = call_result.unwrap();
|
|
assert!(!content.is_empty(), "Should return some content");
|
|
|
|
// Verify the content contains expected text
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("Extensions available to enable:"),
|
|
"Content should contain 'Extensions available to enable:'"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manage_extensions_enable_disable_success() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
|
|
let enable_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({
|
|
"action": "enable",
|
|
"extension_name": "todo"
|
|
})),
|
|
};
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(enable_call, "request_3a".to_string(), None, None)
|
|
.await;
|
|
assert!(result.is_ok());
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"manage_extensions enable execution should succeed"
|
|
);
|
|
|
|
let content = call_result.unwrap();
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("todo") && text.text.contains("installed successfully"),
|
|
"Response should indicate success, got: {}",
|
|
text.text
|
|
);
|
|
|
|
// Now disable it
|
|
let disable_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({
|
|
"action": "disable",
|
|
"extension_name": "todo"
|
|
})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(disable_call, "request_3b".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "manage_extensions disable should succeed");
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"manage_extensions disable execution should succeed"
|
|
);
|
|
|
|
let content = call_result.unwrap();
|
|
assert!(!content.is_empty(), "Should return confirmation message");
|
|
|
|
// Verify the message indicates success
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("successfully") || text.text.contains("disabled"),
|
|
"Response should indicate success, got: {}",
|
|
text.text
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manage_extensions_missing_parameters() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
|
|
// Call manage_extensions without action parameter
|
|
let tool_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({
|
|
"extension_name": "todo"
|
|
})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(tool_call, "request_4".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "Tool call should return a result");
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"Tool execution should return an error result"
|
|
);
|
|
|
|
let content = call_result.unwrap();
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("action") || text.text.contains("parameter"),
|
|
"Error should mention missing action parameter"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manage_extensions_invalid_action() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
|
|
let tool_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({
|
|
"action": "invalid_action",
|
|
"extension_name": "todo"
|
|
})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(tool_call, "request_6".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "Tool call should return a result");
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"Tool execution should return an error result"
|
|
);
|
|
|
|
let content = call_result.unwrap();
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("Invalid action")
|
|
|| text.text.contains("enable")
|
|
|| text.text.contains("disable"),
|
|
"Error should mention invalid action, got: {}",
|
|
text.text
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manage_extensions_extension_not_found() {
|
|
let agent = setup_agent_with_extension_manager().await;
|
|
|
|
// Try to enable a non-existent extension
|
|
let tool_call = CallToolRequestParam {
|
|
name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
|
|
arguments: Some(object!({
|
|
"action": "enable",
|
|
"extension_name": "nonexistent_extension_12345"
|
|
})),
|
|
};
|
|
|
|
let (_, result) = agent
|
|
.dispatch_tool_call(tool_call, "request_7".to_string(), None, None)
|
|
.await;
|
|
|
|
assert!(result.is_ok(), "Tool call should return a result");
|
|
let call_result = result.unwrap().result.await;
|
|
assert!(
|
|
call_result.is_ok(),
|
|
"Tool execution should return an error result"
|
|
);
|
|
|
|
// Check that the error message indicates extension not found
|
|
let content = call_result.unwrap();
|
|
let text = content.first().unwrap().as_text().unwrap();
|
|
assert!(
|
|
text.text.contains("not found") || text.text.contains("Extension"),
|
|
"Error should mention extension not found, got: {}",
|
|
text.text
|
|
);
|
|
}
|
|
}
|