diff --git a/Cargo.lock b/Cargo.lock index 2c1e2340..b6d989de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5452,6 +5452,7 @@ dependencies = [ "mcp-core", "mcp-macros", "pin-project", + "rmcp", "schemars", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 653e34dc..cad9d62a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ description = "An AI agent" uninlined_format_args = "allow" [workspace.dependencies] -rmcp = "0.2.1" +rmcp = { version = "0.2.1", features = ["schemars"] } # Patch for Windows cross-compilation issue with crunchy [patch.crates-io] diff --git a/crates/goose-bench/src/eval_suites/core/developer_image/image.rs b/crates/goose-bench/src/eval_suites/core/developer_image/image.rs index 194605cc..2ac8a8ce 100644 --- a/crates/goose-bench/src/eval_suites/core/developer_image/image.rs +++ b/crates/goose-bench/src/eval_suites/core/developer_image/image.rs @@ -7,7 +7,6 @@ use crate::eval_suites::{ use crate::register_evaluation; use async_trait::async_trait; use goose::message::MessageContent; -use mcp_core::content::Content; use rmcp::model::Role; use serde_json::{self, Value}; @@ -68,7 +67,7 @@ impl Evaluation for DeveloperImage { if let Ok(result) = &tool_resp.tool_result { // Check each item in the result list for item in result { - if let Content::Image(image) = item { + if let Some(image) = item.as_image() { // Image content already contains mime_type and data if image.mime_type.starts_with("image/") && !image.data.is_empty() diff --git a/crates/goose-cli/src/session/export.rs b/crates/goose-cli/src/session/export.rs index eb336bd4..5c971402 100644 --- a/crates/goose-cli/src/session/export.rs +++ b/crates/goose-cli/src/session/export.rs @@ -1,8 +1,6 @@ use goose::message::{Message, MessageContent, ToolRequest, ToolResponse}; use goose::utils::safe_truncate; -use mcp_core::content::Content as McpContent; -use mcp_core::resource::ResourceContents; -use rmcp::model::Role; +use rmcp::model::{RawContent, ResourceContents, Role}; use serde_json::Value; const MAX_STRING_LENGTH_MD_EXPORT: usize = 4096; // Generous limit for export @@ -219,8 +217,8 @@ pub fn tool_response_to_markdown(resp: &ToolResponse, export_all_content: bool) } } - match content { - McpContent::Text(text_content) => { + match &content.raw { + RawContent::Text(text_content) => { let trimmed_text = text_content.text.trim(); if (trimmed_text.starts_with('{') && trimmed_text.ends_with('}')) || (trimmed_text.starts_with('[') && trimmed_text.ends_with(']')) @@ -236,7 +234,7 @@ pub fn tool_response_to_markdown(resp: &ToolResponse, export_all_content: bool) md.push_str("\n\n"); } } - McpContent::Image(image_content) => { + RawContent::Image(image_content) => { if image_content.mime_type.starts_with("image/") { // For actual images, provide a placeholder that indicates it's an image md.push_str(&format!( @@ -252,7 +250,7 @@ pub fn tool_response_to_markdown(resp: &ToolResponse, export_all_content: bool) )); } } - McpContent::Resource(resource) => { + RawContent::Resource(resource) => { match &resource.resource { ResourceContents::TextResourceContents { uri, @@ -299,6 +297,9 @@ pub fn tool_response_to_markdown(resp: &ToolResponse, export_all_content: bool) } } } + RawContent::Audio(_) => { + md.push_str("[audio content not displayed in Markdown export]\n\n") + } } } } @@ -360,8 +361,8 @@ pub fn message_to_markdown(message: &Message, export_all_content: bool) -> Strin mod tests { use super::*; use goose::message::{Message, ToolRequest, ToolResponse}; - use mcp_core::content::{Content as McpContent, TextContent}; use mcp_core::tool::ToolCall; + use rmcp::model::{Content, RawTextContent, TextContent}; use serde_json::json; #[test] @@ -521,12 +522,14 @@ mod tests { #[test] fn test_tool_response_to_markdown_text() { let text_content = TextContent { - text: "Command executed successfully".to_string(), + raw: RawTextContent { + text: "Command executed successfully".to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "test-id".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let result = tool_response_to_markdown(&tool_response, true); @@ -538,12 +541,14 @@ mod tests { fn test_tool_response_to_markdown_json() { let json_text = r#"{"status": "success", "data": "test"}"#; let text_content = TextContent { - text: json_text.to_string(), + raw: RawTextContent { + text: json_text.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "test-id".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let result = tool_response_to_markdown(&tool_response, true); @@ -640,12 +645,14 @@ if __name__ == "__main__": hello_world()"#; let text_content = TextContent { - text: python_code.to_string(), + raw: RawTextContent { + text: python_code.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "shell-cat".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); @@ -677,12 +684,14 @@ if __name__ == "__main__": let git_output = " M src/main.rs\n?? temp.txt\n A new_feature.rs"; let text_content = TextContent { - text: git_output.to_string(), + raw: RawTextContent { + text: git_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "git-status".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); @@ -722,12 +731,14 @@ warning: unused variable `x` Finished dev [unoptimized + debuginfo] target(s) in 2.45s"#; let text_content = TextContent { - text: build_output.to_string(), + raw: RawTextContent { + text: build_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "cargo-build".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let response_result = tool_response_to_markdown(&tool_response, true); @@ -765,12 +776,14 @@ warning: unused variable `x` }"#; let text_content = TextContent { - text: api_response.to_string(), + raw: RawTextContent { + text: api_response.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "curl-api".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let response_result = tool_response_to_markdown(&tool_response, true); @@ -797,12 +810,14 @@ warning: unused variable `x` }; let text_content = TextContent { - text: "File created successfully".to_string(), + raw: RawTextContent { + text: "File created successfully".to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "editor-write".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); @@ -850,12 +865,14 @@ def process_data(data: List[Dict]) -> List[Dict]: return [item for item in data if item.get('active', False)]"#; let text_content = TextContent { - text: python_code.to_string(), + raw: RawTextContent { + text: python_code.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "editor-view".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let response_result = tool_response_to_markdown(&tool_response, true); @@ -883,12 +900,14 @@ def process_data(data: List[Dict]) -> List[Dict]: Command failed with exit code 2"#; let text_content = TextContent { - text: error_output.to_string(), + raw: RawTextContent { + text: error_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "shell-error".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let response_result = tool_response_to_markdown(&tool_response, true); @@ -919,12 +938,14 @@ Command failed with exit code 2"#; 5^2 = 25"#; let text_content = TextContent { - text: script_output.to_string(), + raw: RawTextContent { + text: script_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "script-exec".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); @@ -962,12 +983,14 @@ drwx------ 3 user staff 96 Dec 6 16:20 com.apple.launchd.abc /tmp"#; let text_content = TextContent { - text: multi_output.to_string(), + raw: RawTextContent { + text: multi_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "multi-cmd".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&_tool_request, true); @@ -1001,12 +1024,14 @@ src/database.rs:23:async fn query_users(pool: &Pool) -> Result> { src/middleware.rs:12:async fn auth_middleware(req: Request, next: Next) -> Result {"#; let text_content = TextContent { - text: grep_output.to_string(), + raw: RawTextContent { + text: grep_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "grep-search".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); @@ -1037,12 +1062,14 @@ src/middleware.rs:12:async fn auth_middleware(req: Request, next: Next) -> Resul let json_output = r#"{"status": "success", "data": {"count": 42}}"#; let text_content = TextContent { - text: json_output.to_string(), + raw: RawTextContent { + text: json_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "json-test".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let response_result = tool_response_to_markdown(&tool_response, true); @@ -1074,12 +1101,14 @@ src/middleware.rs:12:async fn auth_middleware(req: Request, next: Next) -> Resul found 0 vulnerabilities"#; let text_content = TextContent { - text: npm_output.to_string(), + raw: RawTextContent { + text: npm_output.to_string(), + }, annotations: None, }; let tool_response = ToolResponse { id: "npm-install".to_string(), - tool_result: Ok(vec![McpContent::Text(text_content)]), + tool_result: Ok(vec![Content::text(text_content.raw.text)]), }; let request_result = tool_request_to_markdown(&tool_request, true); diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 3fcc50e4..a260e8f1 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -35,9 +35,9 @@ use goose::providers::pricing::initialize_pricing_cache; use goose::session; use input::InputResult; use mcp_core::handler::ToolError; -use mcp_core::prompt::PromptMessage; use mcp_core::protocol::JsonRpcMessage; use mcp_core::protocol::JsonRpcNotification; +use rmcp::model::PromptMessage; use rand::{distributions::Alphanumeric, Rng}; use rustyline::EditMode; @@ -359,7 +359,33 @@ impl Session { pub async fn get_prompt(&mut self, name: &str, arguments: Value) -> Result> { let result = self.agent.get_prompt(name, arguments).await?; - Ok(result.messages) + // Convert mcp_core::prompt::PromptMessage to rmcp::model::PromptMessage + let converted_messages = result + .messages + .into_iter() + .map(|msg| rmcp::model::PromptMessage { + role: match msg.role { + mcp_core::prompt::PromptMessageRole::User => { + rmcp::model::PromptMessageRole::User + } + mcp_core::prompt::PromptMessageRole::Assistant => { + rmcp::model::PromptMessageRole::Assistant + } + }, + content: match msg.content { + mcp_core::prompt::PromptMessageContent::Text { text } => { + rmcp::model::PromptMessageContent::Text { text } + } + mcp_core::prompt::PromptMessageContent::Image { image } => { + rmcp::model::PromptMessageContent::Image { image } + } + mcp_core::prompt::PromptMessageContent::Resource { resource } => { + rmcp::model::PromptMessageContent::Resource { resource } + } + }, + }) + .collect(); + Ok(converted_messages) } /// Process a single message and get the response diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index d2c5f003..673cb8c5 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -252,7 +252,7 @@ fn render_tool_response(resp: &ToolResponse, theme: Theme, debug: bool) { if debug { println!("{:#?}", content); - } else if let mcp_core::content::Content::Text(text) = content { + } else if let Some(text) = content.as_text() { print_markdown(&text.text, theme); } } diff --git a/crates/goose-mcp/src/computercontroller/docx_tool.rs b/crates/goose-mcp/src/computercontroller/docx_tool.rs index 5e349d1e..e88564e4 100644 --- a/crates/goose-mcp/src/computercontroller/docx_tool.rs +++ b/crates/goose-mcp/src/computercontroller/docx_tool.rs @@ -1,6 +1,7 @@ use docx_rs::*; use image::{self, ImageFormat}; -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; +use rmcp::model::Content; use std::{fs, io::Cursor}; #[derive(Debug)] @@ -568,9 +569,9 @@ mod tests { let content = result.unwrap(); assert!(!content.is_empty(), "Extracted text should not be empty"); let text = content[0].as_text().unwrap(); - println!("Extracted text:\n{}", text); + println!("Extracted text:\n{}", text.text); assert!( - !text.trim().is_empty(), + !text.text.trim().is_empty(), "Extracted text should not be empty" ); } @@ -609,11 +610,11 @@ mod tests { let content = result.unwrap(); let text = content[0].as_text().unwrap(); assert!( - text.contains("Test Heading"), + text.text.contains("Test Heading"), "Should contain written content" ); assert!( - text.contains("test paragraph"), + text.text.contains("test paragraph"), "Should contain written content" ); @@ -700,15 +701,15 @@ mod tests { let content = result.unwrap(); let text = content[0].as_text().unwrap(); assert!( - text.contains("New content here"), + text.text.contains("New content here"), "Should contain new content" ); assert!( - text.contains("Keep this text"), + text.text.contains("Keep this text"), "Should keep unmodified content" ); assert!( - !text.contains("This should be replaced"), + !text.text.contains("This should be replaced"), "Should not contain replaced text" ); @@ -846,22 +847,25 @@ mod tests { // Check for initial content assert!( - text.contains("Initial content"), + text.text.contains("Initial content"), "Should contain initial content" ); assert!( - text.contains("first paragraph"), + text.text.contains("first paragraph"), "Should contain first paragraph" ); assert!( - text.contains("should stay in the document"), + text.text.contains("should stay in the document"), "Should preserve existing content" ); // Check for new content - assert!(text.contains("New content"), "Should contain new content"); assert!( - text.contains("additional paragraph"), + text.text.contains("New content"), + "Should contain new content" + ); + assert!( + text.text.contains("additional paragraph"), "Should contain appended paragraph" ); diff --git a/crates/goose-mcp/src/computercontroller/mod.rs b/crates/goose-mcp/src/computercontroller/mod.rs index bdc9651e..ac0be9ac 100644 --- a/crates/goose-mcp/src/computercontroller/mod.rs +++ b/crates/goose-mcp/src/computercontroller/mod.rs @@ -17,10 +17,10 @@ use mcp_core::{ protocol::{JsonRpcMessage, ServerCapabilities}, resource::Resource, tool::{Tool, ToolAnnotations}, - Content, }; use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; +use rmcp::model::Content; mod docx_tool; mod pdf_tool; diff --git a/crates/goose-mcp/src/computercontroller/pdf_tool.rs b/crates/goose-mcp/src/computercontroller/pdf_tool.rs index f25dde64..424d1f9e 100644 --- a/crates/goose-mcp/src/computercontroller/pdf_tool.rs +++ b/crates/goose-mcp/src/computercontroller/pdf_tool.rs @@ -1,5 +1,6 @@ use lopdf::{content::Content as PdfContent, Document, Object}; -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; +use rmcp::model::Content; use std::{fs, path::Path}; pub async fn pdf_tool( @@ -341,10 +342,10 @@ mod tests { let content = result.unwrap(); assert!(!content.is_empty(), "Extracted text should not be empty"); let text = content[0].as_text().unwrap(); - println!("Extracted text:\n{}", text); - assert!(text.contains("Page 1"), "Should contain page marker"); + println!("Extracted text:\n{}", text.text); + assert!(text.text.contains("Page 1"), "Should contain page marker"); assert!( - text.contains("This is a test PDF"), + text.text.contains("This is a test PDF"), "Should contain expected test content" ); } @@ -373,18 +374,19 @@ mod tests { "Image extraction result should not be empty" ); let text = content[0].as_text().unwrap(); - println!("Extracted content: {}", text); + println!("Extracted content: {}", text.text); // Should either find images or explicitly state none were found assert!( - text.contains("Saved image to:") || text.contains("No images found"), + text.text.contains("Saved image to:") || text.text.contains("No images found"), "Should either save images or report none found" ); // If we found images, verify they exist - if text.contains("Saved image to:") { + if text.text.contains("Saved image to:") { // Extract the file path from the output let file_path = text + .text .lines() .find(|line| line.contains("Saved image to:")) .and_then(|line| line.split(": ").nth(1)) diff --git a/crates/goose-mcp/src/developer/mod.rs b/crates/goose-mcp/src/developer/mod.rs index b7851f02..eeb7006b 100644 --- a/crates/goose-mcp/src/developer/mod.rs +++ b/crates/goose-mcp/src/developer/mod.rs @@ -27,7 +27,6 @@ use mcp_core::{ protocol::{JsonRpcMessage, JsonRpcNotification, ServerCapabilities}, resource::Resource, tool::Tool, - Content, }; use mcp_core::{ prompt::{Prompt, PromptArgument, PromptTemplate}, @@ -35,6 +34,7 @@ use mcp_core::{ }; use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; +use rmcp::model::Content; use rmcp::model::Role; @@ -1886,7 +1886,7 @@ mod tests { .unwrap() .as_text() .unwrap(); - assert!(text.contains("Hello, world!")); + assert!(text.text.contains("Hello, world!")); temp_dir.close().unwrap(); } @@ -1940,7 +1940,9 @@ mod tests { .as_text() .unwrap(); - assert!(text.contains("has been edited, and the section now reads")); + assert!(text + .text + .contains("has been edited, and the section now reads")); // View the file to verify the change let view_result = router @@ -1968,9 +1970,9 @@ mod tests { // Check that the file has been modified and contains some form of "Rust" // The Editor API might transform the content differently than simple string replacement assert!( - text.contains("Rust") || text.contains("Hello, Rust!"), + text.text.contains("Rust") || text.text.contains("Hello, Rust!"), "Expected content to contain 'Rust', but got: {}", - text + text.text ); temp_dir.close().unwrap(); @@ -2029,7 +2031,7 @@ mod tests { .unwrap(); let text = undo_result.first().unwrap().as_text().unwrap(); - assert!(text.contains("Undid the last edit")); + assert!(text.text.contains("Undid the last edit")); // View the file to verify the undo let view_result = router @@ -2053,7 +2055,7 @@ mod tests { .unwrap() .as_text() .unwrap(); - assert!(text.contains("First line")); + assert!(text.text.contains("First line")); temp_dir.close().unwrap(); } @@ -2507,14 +2509,14 @@ mod tests { .unwrap(); // Should contain lines 3-6 with line numbers - assert!(text.contains("3: Line 3")); - assert!(text.contains("4: Line 4")); - assert!(text.contains("5: Line 5")); - assert!(text.contains("6: Line 6")); - assert!(text.contains("(lines 3-6)")); + assert!(text.text.contains("3: Line 3")); + assert!(text.text.contains("4: Line 4")); + assert!(text.text.contains("5: Line 5")); + assert!(text.text.contains("6: Line 6")); + assert!(text.text.contains("(lines 3-6)")); // Should not contain other lines - assert!(!text.contains("1: Line 1")); - assert!(!text.contains("7: Line 7")); + assert!(!text.text.contains("1: Line 1")); + assert!(!text.text.contains("7: Line 7")); temp_dir.close().unwrap(); } @@ -2569,13 +2571,13 @@ mod tests { .unwrap(); // Should contain lines 3 to end - assert!(text.contains("3: Line 3")); - assert!(text.contains("4: Line 4")); - assert!(text.contains("5: Line 5")); - assert!(text.contains("(lines 3-end)")); + assert!(text.text.contains("3: Line 3")); + assert!(text.text.contains("4: Line 4")); + assert!(text.text.contains("5: Line 5")); + assert!(text.text.contains("(lines 3-end)")); // Should not contain earlier lines - assert!(!text.contains("1: Line 1")); - assert!(!text.contains("2: Line 2")); + assert!(!text.text.contains("1: Line 1")); + assert!(!text.text.contains("2: Line 2")); temp_dir.close().unwrap(); } @@ -2695,7 +2697,7 @@ mod tests { .as_text() .unwrap(); - assert!(text.contains("Text has been inserted at line 1")); + assert!(text.text.contains("Text has been inserted at line 1")); // Verify the file content let view_result = router @@ -2720,10 +2722,10 @@ mod tests { .as_text() .unwrap(); - assert!(view_text.contains("1: Line 1")); - assert!(view_text.contains("2: Line 2")); - assert!(view_text.contains("3: Line 3")); - assert!(view_text.contains("4: Line 4")); + assert!(view_text.text.contains("1: Line 1")); + assert!(view_text.text.contains("2: Line 2")); + assert!(view_text.text.contains("3: Line 3")); + assert!(view_text.text.contains("4: Line 4")); temp_dir.close().unwrap(); } @@ -2778,7 +2780,7 @@ mod tests { .as_text() .unwrap(); - assert!(text.contains("Text has been inserted at line 3")); + assert!(text.text.contains("Text has been inserted at line 3")); // Verify the file content let view_result = router @@ -2803,11 +2805,11 @@ mod tests { .as_text() .unwrap(); - assert!(view_text.contains("1: Line 1")); - assert!(view_text.contains("2: Line 2")); - assert!(view_text.contains("3: Line 3")); - assert!(view_text.contains("4: Line 4")); - assert!(view_text.contains("5: Line 5")); + assert!(view_text.text.contains("1: Line 1")); + assert!(view_text.text.contains("2: Line 2")); + assert!(view_text.text.contains("3: Line 3")); + assert!(view_text.text.contains("4: Line 4")); + assert!(view_text.text.contains("5: Line 5")); temp_dir.close().unwrap(); } @@ -2862,7 +2864,7 @@ mod tests { .as_text() .unwrap(); - assert!(text.contains("Text has been inserted at line 4")); + assert!(text.text.contains("Text has been inserted at line 4")); // Verify the file content let view_result = router @@ -2887,10 +2889,10 @@ mod tests { .as_text() .unwrap(); - assert!(view_text.contains("1: Line 1")); - assert!(view_text.contains("2: Line 2")); - assert!(view_text.contains("3: Line 3")); - assert!(view_text.contains("4: Line 4")); + assert!(view_text.text.contains("1: Line 1")); + assert!(view_text.text.contains("2: Line 2")); + assert!(view_text.text.contains("3: Line 3")); + assert!(view_text.text.contains("4: Line 4")); temp_dir.close().unwrap(); } @@ -3059,7 +3061,7 @@ mod tests { .unwrap(); let text = undo_result.first().unwrap().as_text().unwrap(); - assert!(text.contains("Undid the last edit")); + assert!(text.text.contains("Undid the last edit")); // Verify the file is back to original content let view_result = router @@ -3084,9 +3086,9 @@ mod tests { .as_text() .unwrap(); - assert!(view_text.contains("1: Line 1")); - assert!(view_text.contains("2: Line 2")); - assert!(!view_text.contains("Inserted Line")); + assert!(view_text.text.contains("1: Line 1")); + assert!(view_text.text.contains("2: Line 2")); + assert!(!view_text.text.contains("Inserted Line")); temp_dir.close().unwrap(); } diff --git a/crates/goose-mcp/src/google_drive/mod.rs b/crates/goose-mcp/src/google_drive/mod.rs index 1f1aeae7..5e85ecc1 100644 --- a/crates/goose-mcp/src/google_drive/mod.rs +++ b/crates/goose-mcp/src/google_drive/mod.rs @@ -11,13 +11,13 @@ use mcp_core::protocol::JsonRpcMessage; use mcp_core::tool::ToolAnnotations; use oauth_pkce::PkceOAuth2Client; use regex::Regex; +use rmcp::model::Content; use serde_json::{json, Value}; use std::io::Cursor; use std::{env, fs, future::Future, path::Path, pin::Pin, sync::Arc}; use storage::CredentialsManager; use tokio::sync::mpsc; -use mcp_core::content::Content; use mcp_core::{ handler::{PromptError, ResourceError, ToolError}, prompt::Prompt, @@ -1845,7 +1845,12 @@ impl GoogleDriveRouter { .map(|contents| { contents .into_iter() - .map(|content| content.as_text().unwrap_or_default().to_string()) + .map(|content| { + content + .as_text() + .map(|text| text.text.clone()) + .unwrap_or_default() + }) .collect::>() .join("\n") }) diff --git a/crates/goose-mcp/src/jetbrains/mod.rs b/crates/goose-mcp/src/jetbrains/mod.rs index c015b9de..d6d71822 100644 --- a/crates/goose-mcp/src/jetbrains/mod.rs +++ b/crates/goose-mcp/src/jetbrains/mod.rs @@ -2,7 +2,6 @@ mod proxy; use anyhow::Result; use mcp_core::{ - content::Content, handler::{PromptError, ResourceError, ToolError}, prompt::Prompt, protocol::{JsonRpcMessage, ServerCapabilities}, @@ -12,6 +11,7 @@ use mcp_core::{ }; use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; +use rmcp::model::Content; use serde_json::Value; use std::future::Future; use std::pin::Pin; diff --git a/crates/goose-mcp/src/jetbrains/proxy.rs b/crates/goose-mcp/src/jetbrains/proxy.rs index 382f2714..b75b22c7 100644 --- a/crates/goose-mcp/src/jetbrains/proxy.rs +++ b/crates/goose-mcp/src/jetbrains/proxy.rs @@ -1,6 +1,7 @@ use anyhow::{anyhow, Result}; -use mcp_core::{Content, Tool}; +use mcp_core::Tool; use reqwest::Client; +use rmcp::model::Content; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::env; diff --git a/crates/goose-mcp/src/memory/mod.rs b/crates/goose-mcp/src/memory/mod.rs index eb5adaee..d1fee83c 100644 --- a/crates/goose-mcp/src/memory/mod.rs +++ b/crates/goose-mcp/src/memory/mod.rs @@ -18,10 +18,10 @@ use mcp_core::{ protocol::{JsonRpcMessage, ServerCapabilities}, resource::Resource, tool::{Tool, ToolAnnotations, ToolCall}, - Content, }; use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; +use rmcp::model::Content; // MemoryRouter implementation #[derive(Clone)] diff --git a/crates/goose-mcp/src/tutorial/mod.rs b/crates/goose-mcp/src/tutorial/mod.rs index ea9e32f0..966ab8fd 100644 --- a/crates/goose-mcp/src/tutorial/mod.rs +++ b/crates/goose-mcp/src/tutorial/mod.rs @@ -1,6 +1,7 @@ use anyhow::Result; use include_dir::{include_dir, Dir}; use indoc::formatdoc; +use rmcp::model::Content; use serde_json::{json, Value}; use std::{future::Future, pin::Pin}; use tokio::sync::mpsc; @@ -16,8 +17,6 @@ use mcp_core::{ use mcp_server::router::CapabilitiesBuilder; use mcp_server::Router; -use mcp_core::content::Content; - static TUTORIALS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/tutorial/tutorials"); pub struct TutorialRouter { diff --git a/crates/goose-server/Cargo.toml b/crates/goose-server/Cargo.toml index a1196065..22f2c711 100644 --- a/crates/goose-server/Cargo.toml +++ b/crates/goose-server/Cargo.toml @@ -53,4 +53,4 @@ path = "src/bin/generate_schema.rs" [dev-dependencies] tower = "0.5" -async-trait = "0.1" \ No newline at end of file +async-trait = "0.1" diff --git a/crates/goose-server/src/bin/generate_schema.rs b/crates/goose-server/src/bin/generate_schema.rs index 529be54c..8d2588c0 100644 --- a/crates/goose-server/src/bin/generate_schema.rs +++ b/crates/goose-server/src/bin/generate_schema.rs @@ -1,13 +1,18 @@ use goose_server::openapi; use std::env; use std::fs; +use std::path::PathBuf; fn main() { let schema = openapi::generate_schema(); - // Get the current working directory - let current_dir = env::current_dir().unwrap(); - let output_path = current_dir.join("ui").join("desktop").join("openapi.json"); + let package_dir = env::var("CARGO_MANIFEST_DIR").unwrap(); + let output_path = PathBuf::from(package_dir) + .join("..") + .join("..") + .join("ui") + .join("desktop") + .join("openapi.json"); // Ensure parent directory exists if let Some(parent) = output_path.parent() { diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index ceb210fd..b49f7cd5 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -11,40 +11,282 @@ use goose::permission::permission_confirmation::PrincipalType; use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata}; use goose::session::info::SessionInfo; use goose::session::SessionMetadata; -use mcp_core::content::{Annotations, Content, EmbeddedResource, ImageContent, TextContent}; use mcp_core::handler::ToolResultSchema; use mcp_core::resource::ResourceContents; use mcp_core::tool::{Tool, ToolAnnotations}; +use rmcp::model::{Annotations, Content, EmbeddedResource, ImageContent, Role, TextContent}; use utoipa::{OpenApi, ToSchema}; -/// A wrapper around rmcp::model::Role that implements ToSchema for utoipa -#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, ToSchema)] -#[serde(rename_all = "camelCase")] -pub enum Role { - /// A human user or client making a request - User, - /// An AI assistant or server providing a response - Assistant, +use rmcp::schemars::schema::{InstanceType, SchemaObject, SingleOrVec}; +use utoipa::openapi::schema::{ + AdditionalProperties, AnyOfBuilder, ArrayBuilder, ObjectBuilder, OneOfBuilder, Schema, + SchemaFormat, SchemaType, +}; +use utoipa::openapi::{AllOfBuilder, Ref, RefOr}; + +macro_rules! derive_utoipa { + ($inner_type:ident as $schema_name:ident) => { + struct $schema_name {} + + impl<'__s> ToSchema<'__s> for $schema_name { + fn schema() -> (&'__s str, utoipa::openapi::RefOr) { + let settings = rmcp::schemars::gen::SchemaSettings::openapi3(); + let generator = settings.into_generator(); + let schema = generator.into_root_schema_for::<$inner_type>(); + let schema = convert_schemars_to_utoipa(schema); + (stringify!($inner_type), schema) + } + + fn aliases() -> Vec<(&'__s str, utoipa::openapi::schema::Schema)> { + Vec::new() + } + } + }; } -impl From for Role { - fn from(role: rmcp::model::Role) -> Self { - match role { - rmcp::model::Role::User => Role::User, - rmcp::model::Role::Assistant => Role::Assistant, +fn convert_schemars_to_utoipa(schema: rmcp::schemars::schema::RootSchema) -> RefOr { + convert_schema_object(&rmcp::schemars::schema::Schema::Object( + schema.schema.clone(), + )) +} + +fn convert_schema_object(schema: &rmcp::schemars::schema::Schema) -> RefOr { + match schema { + rmcp::schemars::schema::Schema::Object(schema_object) => { + convert_schema_object_inner(schema_object) + } + rmcp::schemars::schema::Schema::Bool(true) => { + RefOr::T(Schema::Object(ObjectBuilder::new().build())) + } + rmcp::schemars::schema::Schema::Bool(false) => { + RefOr::T(Schema::Object(ObjectBuilder::new().build())) } } } -impl From for rmcp::model::Role { - fn from(role: Role) -> Self { - match role { - Role::User => rmcp::model::Role::User, - Role::Assistant => rmcp::model::Role::Assistant, +fn convert_schema_object_inner(schema: &SchemaObject) -> RefOr { + // Handle references first + if let Some(reference) = &schema.reference { + return RefOr::Ref(Ref::new(reference.clone())); + } + + // Handle subschemas (oneOf, allOf, anyOf) + if let Some(subschemas) = &schema.subschemas { + if let Some(one_of) = &subschemas.one_of { + let schemas: Vec> = one_of.iter().map(convert_schema_object).collect(); + let mut builder = OneOfBuilder::new(); + for schema in schemas { + builder = builder.item(schema); + } + return RefOr::T(Schema::OneOf(builder.build())); + } + if let Some(all_of) = &subschemas.all_of { + let schemas: Vec> = all_of.iter().map(convert_schema_object).collect(); + let mut all_of = AllOfBuilder::new(); + for schema in schemas { + all_of = all_of.item(schema); + } + return RefOr::T(Schema::AllOf(all_of.build())); + } + if let Some(any_of) = &subschemas.any_of { + let schemas: Vec> = any_of.iter().map(convert_schema_object).collect(); + let mut any_of = AnyOfBuilder::new(); + for schema in schemas { + any_of = any_of.item(schema); + } + return RefOr::T(Schema::AnyOf(any_of.build())); + } + } + + // Handle based on instance type + match &schema.instance_type { + Some(SingleOrVec::Single(instance_type)) => { + convert_single_instance_type(instance_type, schema) + } + Some(SingleOrVec::Vec(instance_types)) => { + // Multiple types - use AnyOf + let schemas: Vec> = instance_types + .iter() + .map(|instance_type| convert_single_instance_type(instance_type, schema)) + .collect(); + let mut any_of = AnyOfBuilder::new(); + for schema in schemas { + any_of = any_of.item(schema); + } + RefOr::T(Schema::AnyOf(any_of.build())) + } + None => { + // No type specified - create a generic schema + RefOr::T(Schema::Object(ObjectBuilder::new().build())) } } } +fn convert_single_instance_type( + instance_type: &InstanceType, + schema: &SchemaObject, +) -> RefOr { + match instance_type { + InstanceType::Object => { + let mut object_builder = ObjectBuilder::new(); + + if let Some(object_validation) = &schema.object { + // Add properties + for (name, prop_schema) in &object_validation.properties { + let prop = convert_schema_object(prop_schema); + object_builder = object_builder.property(name, prop); + } + + // Add required fields + for required_field in &object_validation.required { + object_builder = object_builder.required(required_field); + } + + // Handle additional properties + if let Some(additional) = &object_validation.additional_properties { + match &**additional { + rmcp::schemars::schema::Schema::Bool(false) => { + object_builder = object_builder + .additional_properties(Some(AdditionalProperties::FreeForm(false))); + } + rmcp::schemars::schema::Schema::Bool(true) => { + object_builder = object_builder + .additional_properties(Some(AdditionalProperties::FreeForm(true))); + } + rmcp::schemars::schema::Schema::Object(obj) => { + let schema = convert_schema_object( + &rmcp::schemars::schema::Schema::Object(obj.clone()), + ); + object_builder = object_builder + .additional_properties(Some(AdditionalProperties::RefOr(schema))); + } + } + } + } + + RefOr::T(Schema::Object(object_builder.build())) + } + InstanceType::Array => { + let mut array_builder = ArrayBuilder::new(); + + if let Some(array_validation) = &schema.array { + // Add items schema + if let Some(items) = &array_validation.items { + match items { + rmcp::schemars::schema::SingleOrVec::Single(item_schema) => { + let item_schema = convert_schema_object(item_schema); + array_builder = array_builder.items(item_schema); + } + rmcp::schemars::schema::SingleOrVec::Vec(item_schemas) => { + // Multiple item types - use AnyOf + let schemas: Vec> = + item_schemas.iter().map(convert_schema_object).collect(); + let mut any_of = AnyOfBuilder::new(); + for schema in schemas { + any_of = any_of.item(schema); + } + let any_of_schema = RefOr::T(Schema::AnyOf(any_of.build())); + array_builder = array_builder.items(any_of_schema); + } + } + } + + // Add constraints + if let Some(min_items) = array_validation.min_items { + array_builder = array_builder.min_items(Some(min_items as usize)); + } + if let Some(max_items) = array_validation.max_items { + array_builder = array_builder.max_items(Some(max_items as usize)); + } + } + + RefOr::T(Schema::Array(array_builder.build())) + } + InstanceType::String => { + let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::String); + + if let Some(string_validation) = &schema.string { + if let Some(min_length) = string_validation.min_length { + object_builder = object_builder.min_length(Some(min_length as usize)); + } + if let Some(max_length) = string_validation.max_length { + object_builder = object_builder.max_length(Some(max_length as usize)); + } + if let Some(pattern) = &string_validation.pattern { + object_builder = object_builder.pattern(Some(pattern.clone())); + } + } + + if let Some(format) = &schema.format { + object_builder = object_builder.format(Some(SchemaFormat::Custom(format.clone()))); + } + + RefOr::T(Schema::Object(object_builder.build())) + } + InstanceType::Number => { + let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::Number); + + if let Some(number_validation) = &schema.number { + if let Some(minimum) = number_validation.minimum { + object_builder = object_builder.minimum(Some(minimum)); + } + if let Some(maximum) = number_validation.maximum { + object_builder = object_builder.maximum(Some(maximum)); + } + if let Some(exclusive_minimum) = number_validation.exclusive_minimum { + object_builder = object_builder.exclusive_minimum(Some(exclusive_minimum)); + } + if let Some(exclusive_maximum) = number_validation.exclusive_maximum { + object_builder = object_builder.exclusive_maximum(Some(exclusive_maximum)); + } + if let Some(multiple_of) = number_validation.multiple_of { + object_builder = object_builder.multiple_of(Some(multiple_of)); + } + } + + RefOr::T(Schema::Object(object_builder.build())) + } + InstanceType::Integer => { + let mut object_builder = ObjectBuilder::new().schema_type(SchemaType::Integer); + + if let Some(number_validation) = &schema.number { + if let Some(minimum) = number_validation.minimum { + object_builder = object_builder.minimum(Some(minimum)); + } + if let Some(maximum) = number_validation.maximum { + object_builder = object_builder.maximum(Some(maximum)); + } + if let Some(exclusive_minimum) = number_validation.exclusive_minimum { + object_builder = object_builder.exclusive_minimum(Some(exclusive_minimum)); + } + if let Some(exclusive_maximum) = number_validation.exclusive_maximum { + object_builder = object_builder.exclusive_maximum(Some(exclusive_maximum)); + } + if let Some(multiple_of) = number_validation.multiple_of { + object_builder = object_builder.multiple_of(Some(multiple_of)); + } + } + + RefOr::T(Schema::Object(object_builder.build())) + } + InstanceType::Boolean => RefOr::T(Schema::Object( + ObjectBuilder::new() + .schema_type(SchemaType::Boolean) + .build(), + )), + InstanceType::Null => RefOr::T(Schema::Object( + ObjectBuilder::new().schema_type(SchemaType::String).build(), + )), + } +} + +derive_utoipa!(Role as RoleSchema); +derive_utoipa!(Content as ContentSchema); +derive_utoipa!(EmbeddedResource as EmbeddedResourceSchema); +derive_utoipa!(ImageContent as ImageContentSchema); +derive_utoipa!(TextContent as TextContentSchema); +derive_utoipa!(Annotations as AnnotationsSchema); + #[allow(dead_code)] // Used by utoipa for OpenAPI generation #[derive(OpenApi)] #[openapi( @@ -95,11 +337,11 @@ impl From for rmcp::model::Role { super::routes::session::SessionHistoryResponse, Message, MessageContent, - Content, - EmbeddedResource, - ImageContent, - Annotations, - TextContent, + ContentSchema, + EmbeddedResourceSchema, + ImageContentSchema, + AnnotationsSchema, + TextContentSchema, ToolResponse, ToolRequest, ToolResultSchema, @@ -110,7 +352,7 @@ impl From for rmcp::model::Role { ResourceContents, ContextLengthExceeded, SummarizationRequested, - Role, + RoleSchema, ProviderMetadata, ExtensionEntry, ExtensionConfig, diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 1f7f0588..705e88c7 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -18,7 +18,8 @@ use goose::{ permission::{Permission, PermissionConfirmation}, session, }; -use mcp_core::{protocol::JsonRpcMessage, role::Role, Content, ToolResult}; +use mcp_core::{protocol::JsonRpcMessage, role::Role, ToolResult}; +use rmcp::model::Content; use serde::{Deserialize, Serialize}; use serde_json::json; use serde_json::Value; diff --git a/crates/goose/examples/image_tool.rs b/crates/goose/examples/image_tool.rs index 24a75a74..3c860048 100644 --- a/crates/goose/examples/image_tool.rs +++ b/crates/goose/examples/image_tool.rs @@ -5,10 +5,8 @@ use goose::{ message::Message, providers::{bedrock::BedrockProvider, databricks::DatabricksProvider, openai::OpenAiProvider}, }; -use mcp_core::{ - content::Content, - tool::{Tool, ToolCall}, -}; +use mcp_core::tool::{Tool, ToolCall}; +use rmcp::model::Content; use serde_json::json; use std::fs; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 2cf03708..f7482347 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -47,9 +47,8 @@ use crate::agents::tool_router_index_manager::ToolRouterIndexManager; use crate::agents::tool_vectordb::generate_table_id; use crate::agents::types::SessionConfig; use crate::agents::types::{FrontendTool, ToolResultReceiver}; -use mcp_core::{ - prompt::Prompt, protocol::GetPromptResult, tool::Tool, Content, ToolError, ToolResult, -}; +use mcp_core::{prompt::Prompt, protocol::GetPromptResult, tool::Tool, ToolError, ToolResult}; +use rmcp::model::Content; use super::final_output_tool::FinalOutputTool; use super::platform_tools; diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index a418fe11..d03c3bff 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -19,7 +19,8 @@ use crate::config::{Config, ExtensionConfigManager}; use crate::prompt_template; use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait}; use mcp_client::transport::{SseTransport, StdioTransport, StreamableHttpTransport, Transport}; -use mcp_core::{prompt::Prompt, Content, Tool, ToolCall, ToolError}; +use mcp_core::{prompt::Prompt, Tool, ToolCall, ToolError}; +use rmcp::model::Content; use serde_json::Value; // By default, we set it to Jan 1, 2020 if the resource does not have a timestamp diff --git a/crates/goose/src/agents/final_output_tool.rs b/crates/goose/src/agents/final_output_tool.rs index 7059feb0..0c2e779b 100644 --- a/crates/goose/src/agents/final_output_tool.rs +++ b/crates/goose/src/agents/final_output_tool.rs @@ -3,8 +3,9 @@ use crate::recipe::Response; use indoc::formatdoc; use mcp_core::{ tool::{Tool, ToolAnnotations}, - Content, ToolCall, ToolError, + ToolCall, ToolError, }; +use rmcp::model::Content; use serde_json::Value; pub const FINAL_OUTPUT_TOOL_NAME: &str = "recipe__final_output"; diff --git a/crates/goose/src/agents/large_response_handler.rs b/crates/goose/src/agents/large_response_handler.rs index 0369bfa5..ff806621 100644 --- a/crates/goose/src/agents/large_response_handler.rs +++ b/crates/goose/src/agents/large_response_handler.rs @@ -1,5 +1,6 @@ use chrono::Utc; -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; +use rmcp::model::Content; use std::fs::File; use std::io::Write; @@ -14,8 +15,8 @@ pub fn process_tool_response( let mut processed_contents = Vec::new(); for content in contents { - match content { - Content::Text(text_content) => { + match content.as_text() { + Some(text_content) => { // Check if text exceeds threshold if text_content.text.chars().count() > LARGE_TEXT_THRESHOLD { // Write to temp file @@ -41,11 +42,13 @@ pub fn process_tool_response( } } else { // Keep original content for smaller texts - processed_contents.push(Content::Text(text_content)); + processed_contents.push(content); } } - // Pass through other content types unchanged - _ => processed_contents.push(content), + None => { + // Pass through other content types unchanged + processed_contents.push(content); + } } } @@ -76,7 +79,8 @@ fn write_large_text_to_file(content: &str) -> Result { #[cfg(test)] mod tests { use super::*; - use mcp_core::{Content, ImageContent, TextContent, ToolError}; + use mcp_core::ToolError; + use rmcp::model::Content; use std::fs; use std::path::Path; @@ -84,10 +88,7 @@ mod tests { fn test_small_text_response_passes_through() { // Create a small text response let small_text = "This is a small text response"; - let content = Content::Text(TextContent { - text: small_text.to_string(), - annotations: None, - }); + let content = Content::text(small_text.to_string()); let response = Ok(vec![content]); @@ -96,7 +97,7 @@ mod tests { // Verify the response is unchanged assert_eq!(processed.len(), 1); - if let Content::Text(text_content) = &processed[0] { + if let Some(text_content) = processed[0].as_text() { assert_eq!(text_content.text, small_text); } else { panic!("Expected text content"); @@ -107,10 +108,7 @@ mod tests { fn test_large_text_response_redirected_to_file() { // Create a text larger than the threshold let large_text = "a".repeat(LARGE_TEXT_THRESHOLD + 1000); - let content = Content::Text(TextContent { - text: large_text.clone(), - annotations: None, - }); + let content = Content::text(large_text.clone()); let response = Ok(vec![content]); @@ -119,7 +117,7 @@ mod tests { // Verify the response contains a message about the file assert_eq!(processed.len(), 1); - if let Content::Text(text_content) = &processed[0] { + if let Some(text_content) = processed[0].as_text() { assert!(text_content .text .contains("The response returned from the tool call was larger")); @@ -147,11 +145,7 @@ mod tests { #[test] fn test_image_content_passes_through() { // Create an image content - let image_content = Content::Image(ImageContent { - data: "base64data".to_string(), - mime_type: "image/png".to_string(), - annotations: None, - }); + let image_content = Content::image("base64data".to_string(), "image/png".to_string()); let response = Ok(vec![image_content]); @@ -160,12 +154,11 @@ mod tests { // Verify the response is unchanged assert_eq!(processed.len(), 1); - match &processed[0] { - Content::Image(img) => { - assert_eq!(img.data, "base64data"); - assert_eq!(img.mime_type, "image/png"); - } - _ => panic!("Expected image content"), + if let Some(img) = processed[0].as_image() { + assert_eq!(img.data, "base64data"); + assert_eq!(img.mime_type, "image/png"); + } else { + panic!("Expected image content"); } } @@ -173,15 +166,8 @@ mod tests { fn test_mixed_content_handled_correctly() { // Create a response with mixed content types let small_text = Content::text("Small text"); - let large_text = Content::Text(TextContent { - text: "a".repeat(LARGE_TEXT_THRESHOLD + 1000), - annotations: None, - }); - let image = Content::Image(ImageContent { - data: "image_data".to_string(), - mime_type: "image/jpeg".to_string(), - annotations: None, - }); + let large_text = Content::text("a".repeat(LARGE_TEXT_THRESHOLD + 1000)); + let image = Content::image("image_data".to_string(), "image/jpeg".to_string()); let response = Ok(vec![small_text, large_text, image]); @@ -192,14 +178,14 @@ mod tests { assert_eq!(processed.len(), 3); // First item should be unchanged small text - if let Content::Text(text_content) = &processed[0] { + if let Some(text_content) = processed[0].as_text() { assert_eq!(text_content.text, "Small text"); } else { panic!("Expected text content"); } // Second item should be a message about the file - if let Content::Text(text_content) = &processed[1] { + if let Some(text_content) = processed[1].as_text() { assert!(text_content .text .contains("The response returned from the tool call was larger")); @@ -216,12 +202,11 @@ mod tests { } // Third item should be unchanged image - match &processed[2] { - Content::Image(img) => { - assert_eq!(img.data, "image_data"); - assert_eq!(img.mime_type, "image/jpeg"); - } - _ => panic!("Expected image content"), + if let Some(img) = processed[2].as_image() { + assert_eq!(img.data, "image_data"); + assert_eq!(img.mime_type, "image/jpeg"); + } else { + panic!("Expected image content"); } } diff --git a/crates/goose/src/agents/recipe_tools/dynamic_task_tools.rs b/crates/goose/src/agents/recipe_tools/dynamic_task_tools.rs index 449ea04d..e4705e76 100644 --- a/crates/goose/src/agents/recipe_tools/dynamic_task_tools.rs +++ b/crates/goose/src/agents/recipe_tools/dynamic_task_tools.rs @@ -5,7 +5,8 @@ use crate::agents::subagent_execution_tool::tasks_manager::TasksManager; use crate::agents::subagent_execution_tool::{lib::ExecutionMode, task_types::Task}; use crate::agents::tool_execution::ToolCallResult; -use mcp_core::{tool::ToolAnnotations, Content, Tool, ToolError}; +use mcp_core::{tool::ToolAnnotations, Tool, ToolError}; +use rmcp::model::Content; use serde_json::{json, Value}; pub const DYNAMIC_TASK_TOOL_NAME_PREFIX: &str = "dynamic_task__create_task"; diff --git a/crates/goose/src/agents/router_tool_selector.rs b/crates/goose/src/agents/router_tool_selector.rs index 933316bc..52da661a 100644 --- a/crates/goose/src/agents/router_tool_selector.rs +++ b/crates/goose/src/agents/router_tool_selector.rs @@ -1,6 +1,6 @@ -use mcp_core::content::TextContent; use mcp_core::tool::Tool; -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; +use rmcp::model::Content; use anyhow::{Context, Result}; use async_trait::async_trait; @@ -115,10 +115,7 @@ impl RouterToolSelector for VectorToolSelector { "Tool: {}\nDescription: {}\nSchema: {}", tool.tool_name, tool.description, tool.schema ); - Content::Text(TextContent { - text, - annotations: None, - }) + Content::text(text) }) .collect(); @@ -292,12 +289,7 @@ impl RouterToolSelector for LLMToolSelector { let tool_entries: Vec = text .split("\n\n") .filter(|entry| entry.trim().starts_with("Tool:")) - .map(|entry| { - Content::Text(TextContent { - text: entry.trim().to_string(), - annotations: None, - }) - }) + .map(|entry| Content::text(entry.trim().to_string())) .collect(); Ok(tool_entries) diff --git a/crates/goose/src/agents/schedule_tool.rs b/crates/goose/src/agents/schedule_tool.rs index 04386637..30210544 100644 --- a/crates/goose/src/agents/schedule_tool.rs +++ b/crates/goose/src/agents/schedule_tool.rs @@ -6,7 +6,8 @@ use std::sync::Arc; use chrono::Utc; -use mcp_core::{Content, ToolError, ToolResult}; +use mcp_core::{ToolError, ToolResult}; +use rmcp::model::Content; use crate::recipe::Recipe; use crate::scheduler_trait::SchedulerTrait; diff --git a/crates/goose/src/agents/sub_recipe_manager.rs b/crates/goose/src/agents/sub_recipe_manager.rs index 891c3c9b..98431f5b 100644 --- a/crates/goose/src/agents/sub_recipe_manager.rs +++ b/crates/goose/src/agents/sub_recipe_manager.rs @@ -1,4 +1,5 @@ -use mcp_core::{Content, Tool, ToolError}; +use mcp_core::{Tool, ToolError}; +use rmcp::model::Content; use serde_json::Value; use std::collections::HashMap; diff --git a/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs b/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs index 73e3fa12..f3860253 100644 --- a/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs +++ b/crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs @@ -1,4 +1,5 @@ -use mcp_core::{tool::ToolAnnotations, Content, Tool, ToolError}; +use mcp_core::{tool::ToolAnnotations, Tool, ToolError}; +use rmcp::model::Content; use serde_json::Value; use crate::agents::subagent_task_config::TaskConfig; diff --git a/crates/goose/src/agents/subagent_execution_tool/tasks.rs b/crates/goose/src/agents/subagent_execution_tool/tasks.rs index 7ed93245..a330711e 100644 --- a/crates/goose/src/agents/subagent_execution_tool/tasks.rs +++ b/crates/goose/src/agents/subagent_execution_tool/tasks.rs @@ -1,4 +1,5 @@ use serde_json::Value; +use std::ops::Deref; use std::process::Stdio; use std::sync::Arc; use tokio::io::{AsyncBufReadExt, BufReader}; @@ -80,8 +81,10 @@ async fn handle_text_instruction_task( // Extract the text content from the result let result_text = contents .into_iter() - .filter_map(|content| match content { - mcp_core::Content::Text(text) => Some(text.text), + .filter_map(|content| match content.deref() { + rmcp::model::RawContent::Text(raw_text_content) => { + Some(raw_text_content.text.clone()) + } _ => None, }) .collect::>() diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index f40a34d4..6fadd247 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -1,7 +1,8 @@ use crate::agents::subagent::SubAgent; use crate::agents::subagent_task_config::TaskConfig; use anyhow::Result; -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; +use rmcp::model::Content; use serde_json::Value; /// Standalone function to run a complete subagent task diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index 446d1f58..ea997dfd 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -11,7 +11,8 @@ use crate::config::permission::PermissionLevel; use crate::config::PermissionManager; use crate::message::{Message, ToolRequest}; use crate::permission::Permission; -use mcp_core::{Content, ToolResult}; +use mcp_core::ToolResult; +use rmcp::model::Content; // ToolCallResult combines the result of a tool call with an optional notification stream that // can be used to receive notifications from the tool. diff --git a/crates/goose/src/agents/types.rs b/crates/goose/src/agents/types.rs index 41711cd0..4ad78f7f 100644 --- a/crates/goose/src/agents/types.rs +++ b/crates/goose/src/agents/types.rs @@ -1,5 +1,6 @@ use crate::session; -use mcp_core::{Content, Tool, ToolResult}; +use mcp_core::{Tool, ToolResult}; +use rmcp::model::Content; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::sync::Arc; diff --git a/crates/goose/src/context_mgmt/summarize.rs b/crates/goose/src/context_mgmt/summarize.rs index 17c8478a..5b7d049d 100644 --- a/crates/goose/src/context_mgmt/summarize.rs +++ b/crates/goose/src/context_mgmt/summarize.rs @@ -221,8 +221,9 @@ mod tests { use crate::providers::errors::ProviderError; use chrono::Utc; use mcp_core::tool::Tool; - use mcp_core::{Content, TextContent, ToolCall}; + use mcp_core::ToolCall; use rmcp::model::Role; + use rmcp::model::{AnnotateAble, Content, RawTextContent}; use serde_json::json; use std::sync::Arc; @@ -251,10 +252,12 @@ mod tests { Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: "Summarized content".to_string(), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: "Summarized content".to_string(), + } + .no_annotation(), + )], ), ProviderUsage::new("mock".to_string(), Usage::default()), )) @@ -448,10 +451,12 @@ mod tests { let summarized_messages = vec![Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: "Summary".to_string(), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: "Summary".to_string(), + } + .no_annotation(), + )], )]; let arguments = json!({ "param1": "value1" diff --git a/crates/goose/src/context_mgmt/truncate.rs b/crates/goose/src/context_mgmt/truncate.rs index e0e9d0bd..cb20500b 100644 --- a/crates/goose/src/context_mgmt/truncate.rs +++ b/crates/goose/src/context_mgmt/truncate.rs @@ -1,9 +1,9 @@ use crate::message::{Message, MessageContent}; use crate::utils::safe_truncate; use anyhow::{anyhow, Result}; -use mcp_core::{Content, ResourceContents}; -use rmcp::model::Role; +use rmcp::model::{RawContent, ResourceContents, Role}; use std::collections::HashSet; +use std::ops::DerefMut; use tracing::{debug, warn}; /// Maximum size for truncated content in characters @@ -90,7 +90,7 @@ fn truncate_message_content(message: &Message, max_content_size: usize) -> Resul MessageContent::ToolResponse(tool_response) => { if let Ok(ref mut result) = tool_response.tool_result { for content_item in result { - if let Content::Text(ref mut text_content) = content_item { + if let RawContent::Text(ref mut text_content) = content_item.deref_mut() { if text_content.text.chars().count() > max_content_size { let truncated = format!( "{}\n\n[... tool response truncated from {} to {} characters ...]", @@ -102,7 +102,9 @@ fn truncate_message_content(message: &Message, max_content_size: usize) -> Resul } } // Handle Resource content which might contain large text - else if let Content::Resource(ref mut resource_content) = content_item { + else if let RawContent::Resource(ref mut resource_content) = + content_item.deref_mut() + { if let ResourceContents::TextResourceContents { text, .. } = &mut resource_content.resource { @@ -140,19 +142,21 @@ fn estimate_message_tokens(message: &Message, estimate_fn: &dyn Fn(&str) -> usiz MessageContent::ToolResponse(tool_response) => { if let Ok(ref result) = tool_response.tool_result { for content_item in result { - match content_item { - Content::Text(text_content) => { + match &content_item.raw { + RawContent::Text(text_content) => { total_tokens += estimate_fn(&text_content.text); } - Content::Resource(resource_content) => { - match &resource_content.resource { + RawContent::Resource(resource) => { + match &resource.resource { ResourceContents::TextResourceContents { text, .. } => { total_tokens += estimate_fn(text); } _ => total_tokens += 5, // Small overhead for other resource types } } - _ => total_tokens += 5, // Small overhead for other content types + _ => { + total_tokens += 5; // Small overhead for other content types + } } } } @@ -376,8 +380,8 @@ mod tests { use super::*; use crate::message::Message; use anyhow::Result; - use mcp_core::content::Content; use mcp_core::tool::ToolCall; + use rmcp::model::Content; use serde_json::json; // Helper function to create a user text message with a specified token count diff --git a/crates/goose/src/message.rs b/crates/goose/src/message.rs index 3ff05420..699b67aa 100644 --- a/crates/goose/src/message.rs +++ b/crates/goose/src/message.rs @@ -8,12 +8,14 @@ use std::collections::HashSet; /// The content of the messages uses MCP types to avoid additional conversions /// when interacting with MCP servers. use chrono::Utc; -use mcp_core::content::{Content, ImageContent, TextContent}; use mcp_core::handler::ToolResult; -use mcp_core::prompt::{PromptMessage, PromptMessageContent, PromptMessageRole}; -use mcp_core::resource::ResourceContents; use mcp_core::tool::ToolCall; +use rmcp::model::ResourceContents; use rmcp::model::Role; +use rmcp::model::{ + AnnotateAble, Content, ImageContent, PromptMessage, PromptMessageContent, PromptMessageRole, + RawContent, RawImageContent, RawTextContent, TextContent, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; use utoipa::ToSchema; @@ -114,18 +116,17 @@ pub enum MessageContent { impl MessageContent { pub fn text>(text: S) -> Self { - MessageContent::Text(TextContent { - text: text.into(), - annotations: None, - }) + MessageContent::Text(RawTextContent { text: text.into() }.no_annotation()) } pub fn image, T: Into>(data: S, mime_type: T) -> Self { - MessageContent::Image(ImageContent { - data: data.into(), - mime_type: mime_type.into(), - annotations: None, - }) + MessageContent::Image( + RawImageContent { + data: data.into(), + mime_type: mime_type.into(), + } + .no_annotation(), + ) } pub fn tool_request>(id: S, tool_call: ToolResult) -> Self { @@ -220,7 +221,7 @@ impl MessageContent { if let Ok(contents) = &tool_response.tool_result { let texts: Vec = contents .iter() - .filter_map(|content| content.as_text().map(String::from)) + .filter_map(|content| content.as_text().map(|t| t.text.to_string())) .collect(); if !texts.is_empty() { return Some(texts.join("\n")); @@ -257,13 +258,25 @@ impl MessageContent { impl From for MessageContent { fn from(content: Content) -> Self { - match content { - Content::Text(text) => MessageContent::Text(text), - Content::Image(image) => MessageContent::Image(image), - Content::Resource(resource) => MessageContent::Text(TextContent { - text: resource.get_text(), - annotations: None, - }), + match content.raw { + RawContent::Text(text) => { + MessageContent::Text(text.optional_annotate(content.annotations)) + } + RawContent::Image(image) => { + MessageContent::Image(image.optional_annotate(content.annotations)) + } + RawContent::Resource(resource) => { + let text = match &resource.resource { + ResourceContents::TextResourceContents { text, .. } => text.clone(), + ResourceContents::BlobResourceContents { blob, .. } => { + format!("[Binary content: {}]", blob.clone()) + } + }; + MessageContent::text(text) + } + RawContent::Audio(_) => { + MessageContent::text("[Audio content: not supported]".to_string()) + } } } } @@ -280,16 +293,16 @@ impl From for Message { let content = match prompt_message.content { PromptMessageContent::Text { text } => MessageContent::text(text), PromptMessageContent::Image { image } => { - MessageContent::image(image.data, image.mime_type) + MessageContent::image(image.data.clone(), image.mime_type.clone()) } PromptMessageContent::Resource { resource } => { // For resources, convert to text content with the resource text - match resource.resource { + match &resource.resource { ResourceContents::TextResourceContents { text, .. } => { - MessageContent::text(text) + MessageContent::text(text.clone()) } ResourceContents::BlobResourceContents { blob, .. } => { - MessageContent::text(format!("[Binary content: {}]", blob)) + MessageContent::text(format!("[Binary content: {}]", blob.clone())) } } } @@ -512,10 +525,8 @@ impl Message { #[cfg(test)] mod tests { use super::*; - use mcp_core::content::EmbeddedResource; use mcp_core::handler::ToolError; - use mcp_core::prompt::PromptMessageContent; - use mcp_core::resource::ResourceContents; + use rmcp::model::{PromptMessage, PromptMessageContent, RawEmbeddedResource, ResourceContents}; use serde_json::{json, Value}; #[test] @@ -654,11 +665,11 @@ mod tests { #[test] fn test_from_prompt_message_image() { let prompt_content = PromptMessageContent::Image { - image: ImageContent { + image: RawImageContent { data: "base64data".to_string(), mime_type: "image/jpeg".to_string(), - annotations: None, - }, + } + .no_annotation(), }; let prompt_message = PromptMessage { @@ -685,10 +696,7 @@ mod tests { }; let prompt_content = PromptMessageContent::Resource { - resource: EmbeddedResource { - resource, - annotations: None, - }, + resource: RawEmbeddedResource { resource }.no_annotation(), }; let prompt_message = PromptMessage { @@ -714,10 +722,7 @@ mod tests { }; let prompt_content = PromptMessageContent::Resource { - resource: EmbeddedResource { - resource, - annotations: None, - }, + resource: RawEmbeddedResource { resource }.no_annotation(), }; let prompt_message = PromptMessage { diff --git a/crates/goose/src/permission/permission_judge.rs b/crates/goose/src/permission/permission_judge.rs index 40363f02..6a452e24 100644 --- a/crates/goose/src/permission/permission_judge.rs +++ b/crates/goose/src/permission/permission_judge.rs @@ -5,8 +5,8 @@ use crate::message::{Message, MessageContent, ToolRequest}; use crate::providers::base::Provider; use chrono::Utc; use indoc::indoc; +use mcp_core::tool::Tool; use mcp_core::tool::ToolAnnotations; -use mcp_core::{tool::Tool, TextContent}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use std::collections::HashSet; @@ -84,8 +84,7 @@ fn create_check_messages(tool_requests: Vec<&ToolRequest>) -> Vec { check_messages.push(Message::new( rmcp::model::Role::User, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: format!( + vec![MessageContent::text(format!( "Here are the tool requests: {:?}\n\nAnalyze the tool requests and list the tools that perform read-only operations. \ \n\nGuidelines for Read-Only Operations: \ \n- Read-only operations do not modify any data or state. \ @@ -93,9 +92,7 @@ fn create_check_messages(tool_requests: Vec<&ToolRequest>) -> Vec { \n- Write operations include INSERT, UPDATE, DELETE, and file writing. \ \n\nPlease provide a list of tool names that qualify as read-only:", tool_names.join(", "), - ), - annotations: None, - })], + ))], )); check_messages } diff --git a/crates/goose/src/providers/claude_code.rs b/crates/goose/src/providers/claude_code.rs index e210c1b3..563ea89d 100644 --- a/crates/goose/src/providers/claude_code.rs +++ b/crates/goose/src/providers/claude_code.rs @@ -10,7 +10,6 @@ use super::errors::ProviderError; use super::utils::emit_debug_trace; use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; -use mcp_core::content::TextContent; use mcp_core::tool::Tool; use rmcp::model::Role; @@ -98,7 +97,7 @@ impl ClaudeCodeProvider { // Convert tool result contents to text let content_text = tool_contents .iter() - .filter_map(|content| content.as_text()) + .filter_map(|content| content.as_text().map(|t| t.text.clone())) .collect::>() .join("\n"); @@ -214,10 +213,7 @@ impl ClaudeCodeProvider { )); } - let message_content = vec![MessageContent::Text(TextContent { - text: combined_text, - annotations: None, - })]; + let message_content = vec![MessageContent::text(combined_text)]; let response_message = Message::new( Role::Assistant, @@ -356,10 +352,7 @@ impl ClaudeCodeProvider { let message = Message::new( rmcp::model::Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContent::Text(mcp_core::content::TextContent { - text: description.clone(), - annotations: None, - })], + vec![MessageContent::text(description.clone())], ); let usage = Usage::default(); diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index e177c06b..c6123d0a 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -177,8 +177,7 @@ mod tests { use crate::message::{Message, MessageContent}; use crate::providers::base::{ProviderMetadata, ProviderUsage, Usage}; use chrono::Utc; - use mcp_core::content::TextContent; - use rmcp::model::Role; + use rmcp::model::{AnnotateAble, RawTextContent, Role}; use std::env; #[allow(dead_code)] @@ -216,13 +215,15 @@ mod tests { Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: format!( - "Response from {} with model {}", - self.name, self.model_config.model_name - ), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: format!( + "Response from {} with model {}", + self.name, self.model_config.model_name + ), + } + .no_annotation(), + )], ), ProviderUsage::new(self.model_config.model_name.clone(), Usage::default()), )) diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index e7907cf0..3001a08d 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -3,7 +3,6 @@ use crate::model::ModelConfig; use crate::providers::base::Usage; use crate::providers::errors::ProviderError; use anyhow::{anyhow, Result}; -use mcp_core::content::Content; use mcp_core::tool::{Tool, ToolCall}; use rmcp::model::Role; use serde_json::{json, Value}; @@ -69,10 +68,7 @@ pub fn format_messages(messages: &[Message]) -> Vec { Ok(result) => { let text = result .iter() - .filter_map(|c| match c { - Content::Text(t) => Some(t.text.clone()), - _ => None, - }) + .filter_map(|c| c.as_text().map(|t| t.text.clone())) .collect::>() .join("\n"); diff --git a/crates/goose/src/providers/formats/bedrock.rs b/crates/goose/src/providers/formats/bedrock.rs index 4b8ba9f4..ae8840f2 100644 --- a/crates/goose/src/providers/formats/bedrock.rs +++ b/crates/goose/src/providers/formats/bedrock.rs @@ -6,13 +6,12 @@ use aws_sdk_bedrockruntime::types as bedrock; use aws_smithy_types::{Document, Number}; use base64::Engine; use chrono::Utc; -use mcp_core::{Content, ResourceContents, Tool, ToolCall, ToolError, ToolResult}; -use rmcp::model::Role; +use mcp_core::{Tool, ToolCall, ToolError, ToolResult}; +use rmcp::model::{Content, RawContent, ResourceContents, Role}; use serde_json::Value; use super::super::base::Usage; use crate::message::{Message, MessageContent}; -use mcp_core::content::ImageContent; pub fn to_bedrock_message(message: &Message) -> Result { bedrock::Message::builder() @@ -34,7 +33,9 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result { bedrock::ContentBlock::Text("".to_string()) } - MessageContent::Image(image) => bedrock::ContentBlock::Image(to_bedrock_image(image)?), + MessageContent::Image(image) => { + bedrock::ContentBlock::Image(to_bedrock_image(&image.data, &image.mime_type)?) + } MessageContent::Thinking(_) => { // Thinking blocks are not supported in Bedrock - skip bedrock::ContentBlock::Text("".to_string()) @@ -89,7 +90,7 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result>()?, ), Err(_) => None, @@ -115,12 +116,14 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result Result { - Ok(match content { - Content::Text(text) => bedrock::ToolResultContentBlock::Text(text.text.to_string()), - Content::Image(image) => bedrock::ToolResultContentBlock::Image(to_bedrock_image(image)?), - Content::Resource(resource) => match &resource.resource { + Ok(match content.raw { + RawContent::Text(text) => bedrock::ToolResultContentBlock::Text(text.text), + RawContent::Image(image) => { + bedrock::ToolResultContentBlock::Image(to_bedrock_image(&image.data, &image.mime_type)?) + } + RawContent::Resource(resource) => match &resource.resource { ResourceContents::TextResourceContents { text, .. } => { match to_bedrock_document(tool_use_id, &resource.resource)? { Some(doc) => bedrock::ToolResultContentBlock::Document(doc), @@ -131,6 +134,7 @@ pub fn to_bedrock_tool_result_content_block( bail!("Blob resource content is not supported by Bedrock provider yet") } }, + RawContent::Audio(..) => bail!("Audio is not not supported by Bedrock provider"), }) } @@ -141,23 +145,23 @@ pub fn to_bedrock_role(role: &Role) -> bedrock::ConversationRole { } } -pub fn to_bedrock_image(image: &ImageContent) -> Result { +pub fn to_bedrock_image(data: &String, mime_type: &String) -> Result { // Extract format from MIME type - let format = match image.mime_type.as_str() { + let format = match mime_type.as_str() { "image/png" => bedrock::ImageFormat::Png, "image/jpeg" | "image/jpg" => bedrock::ImageFormat::Jpeg, "image/gif" => bedrock::ImageFormat::Gif, "image/webp" => bedrock::ImageFormat::Webp, _ => bail!( "Unsupported image format: {}. Bedrock supports png, jpeg, gif, webp", - image.mime_type + mime_type ), }; // Create image source with base64 data let source = bedrock::ImageSource::Bytes(aws_smithy_types::Blob::new( base64::prelude::BASE64_STANDARD - .decode(&image.data) + .decode(data) .map_err(|e| anyhow!("Failed to decode base64 image data: {}", e))?, )); @@ -348,7 +352,7 @@ pub fn from_bedrock_json(document: &Document) -> Result { mod tests { use super::*; use anyhow::Result; - use mcp_core::content::ImageContent; + use rmcp::model::{AnnotateAble, RawImageContent}; // Base64 encoded 1x1 PNG image for testing const TEST_IMAGE_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="; @@ -364,13 +368,13 @@ mod tests { ]; for mime_type in supported_formats { - let image = ImageContent { + let image = RawImageContent { data: TEST_IMAGE_BASE64.to_string(), mime_type: mime_type.to_string(), - annotations: None, - }; + } + .no_annotation(); - let result = to_bedrock_image(&image); + let result = to_bedrock_image(&image.data, &image.mime_type); assert!(result.is_ok(), "Failed to convert {} format", mime_type); } @@ -379,13 +383,13 @@ mod tests { #[test] fn test_to_bedrock_image_unsupported_format() { - let image = ImageContent { + let image = RawImageContent { data: TEST_IMAGE_BASE64.to_string(), mime_type: "image/bmp".to_string(), - annotations: None, - }; + } + .no_annotation(); - let result = to_bedrock_image(&image); + let result = to_bedrock_image(&image.data, &image.mime_type); assert!(result.is_err()); let error_msg = result.unwrap_err().to_string(); assert!(error_msg.contains("Unsupported image format: image/bmp")); @@ -394,13 +398,13 @@ mod tests { #[test] fn test_to_bedrock_image_invalid_base64() { - let image = ImageContent { + let image = RawImageContent { data: "invalid_base64_data!!!".to_string(), mime_type: "image/png".to_string(), - annotations: None, - }; + } + .no_annotation(); - let result = to_bedrock_image(&image); + let result = to_bedrock_image(&image.data, &image.mime_type); assert!(result.is_err()); let error_msg = result.unwrap_err().to_string(); assert!(error_msg.contains("Failed to decode base64 image data")); @@ -408,11 +412,11 @@ mod tests { #[test] fn test_to_bedrock_message_content_image() -> Result<()> { - let image = ImageContent { + let image = RawImageContent { data: TEST_IMAGE_BASE64.to_string(), mime_type: "image/png".to_string(), - annotations: None, - }; + } + .no_annotation(); let message_content = MessageContent::Image(image); let result = to_bedrock_message_content(&message_content)?; @@ -425,14 +429,8 @@ mod tests { #[test] fn test_to_bedrock_tool_result_content_block_image() -> Result<()> { - let image = ImageContent { - data: TEST_IMAGE_BASE64.to_string(), - mime_type: "image/png".to_string(), - annotations: None, - }; - - let content = Content::Image(image); - let result = to_bedrock_tool_result_content_block("test_id", &content)?; + let content = Content::image(TEST_IMAGE_BASE64.to_string(), "image/png".to_string()); + let result = to_bedrock_tool_result_content_block("test_id", content)?; // Verify the wrapper correctly converts Content::Image to ToolResultContentBlock::Image assert!(matches!(result, bedrock::ToolResultContentBlock::Image(_))); diff --git a/crates/goose/src/providers/formats/databricks.rs b/crates/goose/src/providers/formats/databricks.rs index 89cae498..10a59cc1 100644 --- a/crates/goose/src/providers/formats/databricks.rs +++ b/crates/goose/src/providers/formats/databricks.rs @@ -6,8 +6,9 @@ use crate::providers::utils::{ }; use anyhow::{anyhow, Error}; use mcp_core::ToolError; -use mcp_core::{Content, Tool, ToolCall}; +use mcp_core::{Tool, ToolCall}; use rmcp::model::Role; +use rmcp::model::{AnnotateAble, Content, RawContent, ResourceContents}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -127,7 +128,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< .audience() .is_none_or(|audience| audience.contains(&Role::Assistant)) }) - .map(|content| content.unannotated()) + .map(|content| content.raw.clone()) .collect(); // Process all content, replacing images with placeholder text @@ -136,30 +137,33 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< for content in abridged { match content { - Content::Image(image) => { + RawContent::Image(image) => { // Add placeholder text in the tool response tool_content.push(Content::text("This tool result included an image that is uploaded in the next message.")); // Create a separate image message image_messages.push(json!({ "role": "user", - "content": [convert_image(&image, image_format)] + "content": [convert_image(&image.no_annotation(), image_format)] })); } - Content::Resource(resource) => { - tool_content.push(Content::text(resource.get_text())); + RawContent::Resource(resource) => { + let text = match &resource.resource { + ResourceContents::TextResourceContents { + text, .. + } => text.clone(), + _ => String::new(), + }; + tool_content.push(Content::text(text)); } _ => { - tool_content.push(content); + tool_content.push(content.no_annotation()); } } } let tool_response_content: Value = json!(tool_content .iter() - .map(|content| match content { - Content::Text(text) => text.text.clone(), - _ => String::new(), - }) + .filter_map(|content| content.as_text().map(|t| t.text.clone())) .collect::>() .join(" ")); @@ -585,7 +589,6 @@ pub fn create_request( #[cfg(test)] mod tests { use super::*; - use mcp_core::content::Content; use serde_json::json; #[test] diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index 46806d3c..2d0c4871 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -4,11 +4,11 @@ use crate::providers::base::Usage; use crate::providers::errors::ProviderError; use crate::providers::utils::{is_valid_function_name, sanitize_function_name}; use anyhow::Result; -use mcp_core::content::Content; use mcp_core::tool::{Tool, ToolCall}; use rand::{distributions::Alphanumeric, Rng}; -use rmcp::model::Role; +use rmcp::model::{AnnotateAble, RawContent, Role}; use serde_json::{json, Map, Value}; +use std::ops::Deref; /// Convert internal Message format to Google's API message specification pub fn format_messages(messages: &[Message]) -> Vec { @@ -66,13 +66,13 @@ pub fn format_messages(messages: &[Message]) -> Vec { audience.contains(&Role::Assistant) }) }) - .map(|content| content.unannotated()) + .map(|content| content.raw.clone()) .collect(); let mut tool_content = Vec::new(); for content in abridged { match content { - Content::Image(image) => { + RawContent::Image(image) => { parts.push(json!({ "inline_data": { "mime_type": image.mime_type, @@ -81,15 +81,20 @@ pub fn format_messages(messages: &[Message]) -> Vec { })); } _ => { - tool_content.push(content); + tool_content.push(content.no_annotation()); } } } let mut text = tool_content .iter() - .filter_map(|c| match c { - Content::Text(t) => Some(t.text.clone()), - Content::Resource(r) => Some(r.get_text()), + .filter_map(|c| match c.deref() { + RawContent::Text(t) => Some(t.text.clone()), + RawContent::Resource(raw_embedded_resource) => Some( + raw_embedded_resource + .clone() + .no_annotation() + .get_text(), + ), _ => None, }) .collect::>() @@ -313,6 +318,7 @@ pub fn create_request( #[cfg(test)] mod tests { use super::*; + use rmcp::model::Content; use serde_json::json; fn set_up_text_message(text: &str, role: Role) -> Message { diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 68db8be9..d6b62933 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -9,10 +9,12 @@ use anyhow::{anyhow, Error}; use async_stream::try_stream; use futures::Stream; use mcp_core::ToolError; -use mcp_core::{Content, Tool, ToolCall}; +use mcp_core::{Tool, ToolCall}; use rmcp::model::Role; +use rmcp::model::{AnnotateAble, Content, RawContent, ResourceContents}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +use std::ops::Deref; #[derive(Serialize, Deserialize, Debug)] struct DeltaToolCallFunction { @@ -135,7 +137,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< .audience() .is_none_or(|audience| audience.contains(&Role::Assistant)) }) - .map(|content| content.unannotated()) + .cloned() .collect(); // Process all content, replacing images with placeholder text @@ -143,19 +145,25 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< let mut image_messages = Vec::new(); for content in abridged { - match content { - Content::Image(image) => { + match content.deref() { + RawContent::Image(image) => { // Add placeholder text in the tool response tool_content.push(Content::text("This tool result included an image that is uploaded in the next message.")); // Create a separate image message image_messages.push(json!({ "role": "user", - "content": [convert_image(&image, image_format)] + "content": [convert_image(&image.clone().no_annotation(), image_format)] })); } - Content::Resource(resource) => { - tool_content.push(Content::text(resource.get_text())); + RawContent::Resource(resource) => { + let text = match &resource.resource { + ResourceContents::TextResourceContents { + text, .. + } => text.clone(), + _ => String::new(), + }; + tool_content.push(Content::text(text)); } _ => { tool_content.push(content); @@ -164,8 +172,8 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec< } let tool_response_content: Value = json!(tool_content .iter() - .map(|content| match content { - Content::Text(text) => text.text.clone(), + .map(|content| match content.deref() { + RawContent::Text(text) => text.text.clone(), _ => String::new(), }) .collect::>() @@ -600,7 +608,6 @@ pub fn create_request( #[cfg(test)] mod tests { use super::*; - use mcp_core::content::Content; use serde_json::json; #[test] diff --git a/crates/goose/src/providers/formats/snowflake.rs b/crates/goose/src/providers/formats/snowflake.rs index 9e7fb6d1..d29e8448 100644 --- a/crates/goose/src/providers/formats/snowflake.rs +++ b/crates/goose/src/providers/formats/snowflake.rs @@ -3,7 +3,6 @@ use crate::model::ModelConfig; use crate::providers::base::Usage; use crate::providers::errors::ProviderError; use anyhow::{anyhow, Result}; -use mcp_core::content::Content; use mcp_core::tool::{Tool, ToolCall}; use rmcp::model::Role; use serde_json::{json, Value}; @@ -39,10 +38,7 @@ pub fn format_messages(messages: &[Message]) -> Vec { if let Ok(result) = &tool_response.tool_result { let text = result .iter() - .filter_map(|c| match c { - Content::Text(t) => Some(t.text.clone()), - _ => None, - }) + .filter_map(|c| c.as_text().map(|t| t.text.clone())) .collect::>() .join("\n"); diff --git a/crates/goose/src/providers/gemini_cli.rs b/crates/goose/src/providers/gemini_cli.rs index 9eb1bce9..f3175b7d 100644 --- a/crates/goose/src/providers/gemini_cli.rs +++ b/crates/goose/src/providers/gemini_cli.rs @@ -10,7 +10,6 @@ use super::errors::ProviderError; use super::utils::emit_debug_trace; use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; -use mcp_core::content::TextContent; use mcp_core::tool::Tool; use rmcp::model::Role; @@ -172,10 +171,7 @@ impl GeminiCliProvider { let message = Message::new( Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: response_text, - annotations: None, - })], + vec![MessageContent::text(response_text)], ); let usage = Usage::default(); // No usage info available for gemini CLI @@ -217,10 +213,7 @@ impl GeminiCliProvider { let message = Message::new( Role::Assistant, chrono::Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: description.clone(), - annotations: None, - })], + vec![MessageContent::text(description.clone())], ); let usage = Usage::default(); diff --git a/crates/goose/src/providers/lead_worker.rs b/crates/goose/src/providers/lead_worker.rs index ef33694b..5d993b52 100644 --- a/crates/goose/src/providers/lead_worker.rs +++ b/crates/goose/src/providers/lead_worker.rs @@ -1,5 +1,6 @@ use anyhow::Result; use async_trait::async_trait; +use std::ops::Deref; use std::sync::Arc; use tokio::sync::Mutex; @@ -7,7 +8,8 @@ use super::base::{LeadWorkerProviderTrait, Provider, ProviderMetadata, ProviderU use super::errors::ProviderError; use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; -use mcp_core::{tool::Tool, Content}; +use mcp_core::tool::Tool; +use rmcp::model::{Content, RawContent}; /// A provider that switches between a lead model and a worker model based on turn count /// and can fallback to lead model on consecutive failures @@ -239,7 +241,7 @@ impl LeadWorkerProvider { /// Check if tool output contains error indicators fn contains_error_indicators(&self, contents: &[Content]) -> bool { for content in contents { - if let Content::Text(text_content) = content { + if let RawContent::Text(text_content) = content.deref() { let text_lower = text_content.text.to_lowercase(); // Common error patterns in tool outputs @@ -455,8 +457,7 @@ mod tests { use crate::message::MessageContent; use crate::providers::base::{ProviderMetadata, ProviderUsage, Usage}; use chrono::Utc; - use mcp_core::content::TextContent; - use rmcp::model::Role; + use rmcp::model::{AnnotateAble, RawTextContent, Role}; #[derive(Clone)] struct MockProvider { @@ -484,10 +485,12 @@ mod tests { Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: format!("Response from {}", self.name), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: format!("Response from {}", self.name), + } + .no_annotation(), + )], ), ProviderUsage::new(self.name.clone(), Usage::default()), )) @@ -647,10 +650,12 @@ mod tests { Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: format!("Response from {}", self.name), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: format!("Response from {}", self.name), + } + .no_annotation(), + )], ), ProviderUsage::new(self.name.clone(), Usage::default()), )) diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs index 2ee13e8b..d5da1058 100644 --- a/crates/goose/src/providers/sagemaker_tgi.rs +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -16,7 +16,6 @@ use super::utils::emit_debug_trace; use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; use chrono::Utc; -use mcp_core::content::TextContent; use rmcp::model::Role; pub const SAGEMAKER_TGI_DOC_LINK: &str = @@ -206,10 +205,7 @@ impl SageMakerTgiProvider { Ok(Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: clean_text, - annotations: None, - })], + vec![MessageContent::text(clean_text)], )) } diff --git a/crates/goose/src/providers/toolshim.rs b/crates/goose/src/providers/toolshim.rs index 0647d0a0..f07f6556 100644 --- a/crates/goose/src/providers/toolshim.rs +++ b/crates/goose/src/providers/toolshim.rs @@ -38,9 +38,10 @@ use crate::model::ModelConfig; use crate::providers::formats::openai::create_request; use anyhow::Result; use mcp_core::tool::{Tool, ToolCall}; -use mcp_core::Content; use reqwest::Client; +use rmcp::model::RawContent; use serde_json::{json, Value}; +use std::ops::Deref; use std::time::Duration; use uuid::Uuid; @@ -340,8 +341,8 @@ pub fn convert_tool_messages_to_text(messages: &[Message]) -> Vec { Ok(contents) => { let text_contents: Vec = contents .iter() - .filter_map(|c| match c { - Content::Text(t) => Some(t.text.clone()), + .filter_map(|c| match c.deref() { + RawContent::Text(t) => Some(t.text.clone()), _ => None, }) .collect(); diff --git a/crates/goose/src/providers/utils.rs b/crates/goose/src/providers/utils.rs index 9c4ec8c6..c76cdfcb 100644 --- a/crates/goose/src/providers/utils.rs +++ b/crates/goose/src/providers/utils.rs @@ -5,13 +5,13 @@ use anyhow::Result; use base64::Engine; use regex::Regex; use reqwest::{Response, StatusCode}; +use rmcp::model::{AnnotateAble, ImageContent, RawImageContent}; use serde::{Deserialize, Serialize}; use serde_json::{from_value, json, Map, Value}; use std::io::Read; use std::path::Path; use crate::providers::errors::{OpenAIError, ProviderError}; -use mcp_core::content::ImageContent; #[derive(serde::Deserialize)] struct OpenAIErrorResponse { @@ -292,11 +292,11 @@ pub fn load_image_file(path: &str) -> Result { // Convert to base64 let data = base64::prelude::BASE64_STANDARD.encode(&bytes); - Ok(ImageContent { + Ok(RawImageContent { mime_type: mime_type.to_string(), data, - annotations: None, - }) + } + .no_annotation()) } pub fn unescape_json_values(value: &Value) -> Value { diff --git a/crates/goose/src/scheduler.rs b/crates/goose/src/scheduler.rs index 584a6815..ef1dee41 100644 --- a/crates/goose/src/scheduler.rs +++ b/crates/goose/src/scheduler.rs @@ -1331,8 +1331,8 @@ mod tests { providers::base::{ProviderMetadata, ProviderUsage, Usage}, providers::errors::ProviderError, }; - use mcp_core::{content::TextContent, tool::Tool}; - use rmcp::model::Role; + use mcp_core::tool::Tool; + use rmcp::model::{AnnotateAble, RawTextContent, Role}; // Removed: use crate::session::storage::{get_most_recent_session, read_metadata}; // `read_metadata` is still used by the test itself, so keep it or its module. use crate::session::storage::read_metadata; @@ -1375,10 +1375,12 @@ mod tests { Message::new( Role::Assistant, Utc::now().timestamp(), - vec![MessageContent::Text(TextContent { - text: "Mocked scheduled response".to_string(), - annotations: None, - })], + vec![MessageContent::Text( + RawTextContent { + text: "Mocked scheduled response".to_string(), + } + .no_annotation(), + )], ), ProviderUsage::new("mock-scheduler-test".to_string(), Usage::default()), )) diff --git a/crates/goose/src/session/storage.rs b/crates/goose/src/session/storage.rs index ceaa45c6..4f7557f3 100644 --- a/crates/goose/src/session/storage.rs +++ b/crates/goose/src/session/storage.rs @@ -15,6 +15,7 @@ use regex::Regex; use serde::{Deserialize, Serialize}; use std::fs; use std::io::{self, BufRead, Write}; +use std::ops::DerefMut; use std::path::{Path, PathBuf}; use std::sync::Arc; use utoipa::ToSchema; @@ -679,7 +680,7 @@ fn parse_message_with_truncation( /// Truncate content within a message in place fn truncate_message_content_in_place(message: &mut Message, max_content_size: usize) { use crate::message::MessageContent; - use mcp_core::{Content, ResourceContents}; + use rmcp::model::{RawContent, ResourceContents}; for content in &mut message.content { match content { @@ -697,8 +698,8 @@ fn truncate_message_content_in_place(message: &mut Message, max_content_size: us MessageContent::ToolResponse(tool_response) => { if let Ok(ref mut result) = tool_response.tool_result { for content_item in result { - match content_item { - Content::Text(ref mut text_content) => { + match content_item.deref_mut() { + RawContent::Text(ref mut text_content) => { if text_content.text.chars().count() > max_content_size { let truncated = format!( "{}\n\n[... tool response truncated during session loading from {} to {} characters ...]", @@ -709,7 +710,7 @@ fn truncate_message_content_in_place(message: &mut Message, max_content_size: us text_content.text = truncated; } } - Content::Resource(ref mut resource_content) => { + RawContent::Resource(ref mut resource_content) => { if let ResourceContents::TextResourceContents { text, .. } = &mut resource_content.resource { diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 7b951be2..ad8ffe5d 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -613,9 +613,9 @@ mod final_output_tool_tests { let content = final_result.unwrap(); let text = content.first().unwrap().as_text().unwrap(); assert!( - text.contains("Final output successfully collected."), + text.text.contains("Final output successfully collected."), "Tool result missing expected content: {}", - text + text.text ); // Simulate the reply stream continuing after the final output tool call. diff --git a/crates/goose/tests/private_tests.rs b/crates/goose/tests/private_tests.rs index 78e7cdd7..d2ec7a06 100644 --- a/crates/goose/tests/private_tests.rs +++ b/crates/goose/tests/private_tests.rs @@ -1,6 +1,6 @@ #![cfg(test)] -use mcp_core::{Content, ToolError}; +use mcp_core::ToolError; use serde_json::json; use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; @@ -33,7 +33,7 @@ async fn test_schedule_tool_list_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content.text.contains("Scheduled Jobs:")); assert!(text_content.text.contains("job1")); assert!(text_content.text.contains("job2")); @@ -63,7 +63,7 @@ async fn test_schedule_tool_list_action_empty() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content.text.contains("Scheduled Jobs:")); } @@ -127,7 +127,7 @@ async fn test_schedule_tool_create_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Successfully created scheduled job")); @@ -286,7 +286,7 @@ async fn test_schedule_tool_run_now_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Successfully started job 'job1'")); @@ -370,7 +370,7 @@ async fn test_schedule_tool_pause_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content.text.contains("Successfully paused job 'job1'")); } @@ -455,7 +455,7 @@ async fn test_schedule_tool_unpause_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Successfully unpaused job 'job1'")); @@ -487,7 +487,7 @@ async fn test_schedule_tool_delete_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Successfully deleted job 'job1'")); @@ -521,7 +521,7 @@ async fn test_schedule_tool_kill_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Successfully killed running job 'job1'")); @@ -585,7 +585,7 @@ async fn test_schedule_tool_inspect_action_running() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Job 'job1' is currently running")); @@ -617,7 +617,7 @@ async fn test_schedule_tool_inspect_action_not_running() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Job 'job1' is not currently running")); @@ -663,7 +663,7 @@ async fn test_schedule_tool_sessions_action() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content.text.contains("Sessions for job 'job1'")); assert!(text_content.text.contains("session1")); assert!(text_content.text.contains("session2")); @@ -738,7 +738,7 @@ async fn test_schedule_tool_sessions_action_empty() { let content = result.unwrap(); assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("No sessions found for job 'job1'")); @@ -809,7 +809,7 @@ async fn test_schedule_tool_session_content_action_with_real_session() { if let Ok(content) = result { assert_eq!(content.len(), 1); - if let mcp_core::Content::Text(text_content) = &content[0] { + if let Some(text_content) = content[0].as_text() { assert!(text_content .text .contains("Session 'test_session_real' Content:")); diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index c4884b7c..5cebdf3d 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -6,8 +6,8 @@ use goose::providers::errors::ProviderError; use goose::providers::{ anthropic, azure, bedrock, databricks, google, groq, ollama, openai, openrouter, snowflake, xai, }; -use mcp_core::content::Content; use mcp_core::tool::Tool; +use rmcp::model::{AnnotateAble, Content, RawImageContent}; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; @@ -256,7 +256,6 @@ impl ProviderTester { async fn test_image_content_support(&self) -> Result<()> { use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; - use mcp_core::content::ImageContent; use std::fs; // Try to read the test image @@ -273,11 +272,11 @@ impl ProviderTester { }; let base64_image = BASE64.encode(image_data); - let image_content = ImageContent { + let image_content = RawImageContent { data: base64_image, mime_type: "image/png".to_string(), - annotations: None, - }; + } + .no_annotation(); // Test 1: Direct image message let message_with_image = @@ -324,8 +323,13 @@ impl ProviderTester { serde_json::json!({}), )), ); - let tool_response = - Message::user().with_tool_response("test_id", Ok(vec![Content::Image(image_content)])); + let tool_response = Message::user().with_tool_response( + "test_id", + Ok(vec![Content::image( + image_content.data.clone(), + image_content.mime_type.clone(), + )]), + ); let result2 = self .provider diff --git a/crates/mcp-core/src/content.rs b/crates/mcp-core/src/content.rs deleted file mode 100644 index da728ac1..00000000 --- a/crates/mcp-core/src/content.rs +++ /dev/null @@ -1,313 +0,0 @@ -/// Content sent around agents, extensions, and LLMs -/// The various content types can be display to humans but also understood by models -/// They include optional annotations used to help inform agent usage -use super::Role; -use crate::resource::ResourceContents; -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use utoipa::ToSchema; - -#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct Annotations { - #[serde(skip_serializing_if = "Option::is_none")] - pub audience: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub priority: Option, - #[serde(skip_serializing_if = "Option::is_none")] - #[schema(value_type = String, format = "date-time", example = "2023-01-01T00:00:00Z")] - // for openapi - pub timestamp: Option>, -} - -impl Annotations { - /// Creates a new Annotations instance specifically for resources - /// optional priority, and a timestamp (defaults to now if None) - pub fn for_resource(priority: f32, timestamp: DateTime) -> Self { - assert!( - (0.0..=1.0).contains(&priority), - "Priority {priority} must be between 0.0 and 1.0" - ); - Annotations { - priority: Some(priority), - timestamp: Some(timestamp), - audience: None, - } - } -} - -#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct TextContent { - pub text: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub annotations: Option, -} - -#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ImageContent { - pub data: String, - pub mime_type: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub annotations: Option, -} - -#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct EmbeddedResource { - pub resource: ResourceContents, - #[serde(skip_serializing_if = "Option::is_none")] - pub annotations: Option, -} - -impl EmbeddedResource { - pub fn get_text(&self) -> String { - match &self.resource { - ResourceContents::TextResourceContents { text, .. } => text.clone(), - _ => String::new(), - } - } -} - -#[derive(ToSchema, Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(tag = "type", rename_all = "camelCase")] -pub enum Content { - Text(TextContent), - Image(ImageContent), - Resource(EmbeddedResource), -} - -impl Content { - pub fn text>(text: S) -> Self { - Content::Text(TextContent { - text: text.into(), - annotations: None, - }) - } - - pub fn image, T: Into>(data: S, mime_type: T) -> Self { - Content::Image(ImageContent { - data: data.into(), - mime_type: mime_type.into(), - annotations: None, - }) - } - - pub fn resource(resource: ResourceContents) -> Self { - Content::Resource(EmbeddedResource { - resource, - annotations: None, - }) - } - - pub fn embedded_text, T: Into>(uri: S, content: T) -> Self { - Content::Resource(EmbeddedResource { - resource: ResourceContents::TextResourceContents { - uri: uri.into(), - mime_type: Some("text".to_string()), - text: content.into(), - }, - annotations: None, - }) - } - - /// Get the text content if this is a TextContent variant - pub fn as_text(&self) -> Option<&str> { - match self { - Content::Text(text) => Some(&text.text), - _ => None, - } - } - - /// Get the image content if this is an ImageContent variant - pub fn as_image(&self) -> Option<(&str, &str)> { - match self { - Content::Image(image) => Some((&image.data, &image.mime_type)), - _ => None, - } - } - - /// Set the audience for the content - pub fn with_audience(mut self, audience: Vec) -> Self { - let annotations = match &mut self { - Content::Text(text) => &mut text.annotations, - Content::Image(image) => &mut image.annotations, - Content::Resource(resource) => &mut resource.annotations, - }; - *annotations = Some(match annotations.take() { - Some(mut a) => { - a.audience = Some(audience); - a - } - None => Annotations { - audience: Some(audience), - priority: None, - timestamp: None, - }, - }); - self - } - - /// Set the priority for the content - /// # Panics - /// Panics if priority is not between 0.0 and 1.0 inclusive - pub fn with_priority(mut self, priority: f32) -> Self { - if !(0.0..=1.0).contains(&priority) { - panic!("Priority must be between 0.0 and 1.0"); - } - let annotations = match &mut self { - Content::Text(text) => &mut text.annotations, - Content::Image(image) => &mut image.annotations, - Content::Resource(resource) => &mut resource.annotations, - }; - *annotations = Some(match annotations.take() { - Some(mut a) => { - a.priority = Some(priority); - a - } - None => Annotations { - audience: None, - priority: Some(priority), - timestamp: None, - }, - }); - self - } - - /// Get the audience if set - pub fn audience(&self) -> Option<&Vec> { - match self { - Content::Text(text) => text.annotations.as_ref().and_then(|a| a.audience.as_ref()), - Content::Image(image) => image.annotations.as_ref().and_then(|a| a.audience.as_ref()), - Content::Resource(resource) => resource - .annotations - .as_ref() - .and_then(|a| a.audience.as_ref()), - } - } - - /// Get the priority if set - pub fn priority(&self) -> Option { - match self { - Content::Text(text) => text.annotations.as_ref().and_then(|a| a.priority), - Content::Image(image) => image.annotations.as_ref().and_then(|a| a.priority), - Content::Resource(resource) => resource.annotations.as_ref().and_then(|a| a.priority), - } - } - - pub fn unannotated(&self) -> Self { - match self { - Content::Text(text) => Content::text(text.text.clone()), - Content::Image(image) => Content::image(image.data.clone(), image.mime_type.clone()), - Content::Resource(resource) => Content::resource(resource.resource.clone()), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_content_text() { - let content = Content::text("hello"); - assert_eq!(content.as_text(), Some("hello")); - assert_eq!(content.as_image(), None); - } - - #[test] - fn test_content_image() { - let content = Content::image("data", "image/png"); - assert_eq!(content.as_text(), None); - assert_eq!(content.as_image(), Some(("data", "image/png"))); - } - - #[test] - fn test_content_annotations_basic() { - let content = Content::text("hello") - .with_audience(vec![Role::User]) - .with_priority(0.5); - assert_eq!(content.audience(), Some(&vec![Role::User])); - assert_eq!(content.priority(), Some(0.5)); - } - - #[test] - fn test_content_annotations_order_independence() { - let content1 = Content::text("hello") - .with_audience(vec![Role::User]) - .with_priority(0.5); - let content2 = Content::text("hello") - .with_priority(0.5) - .with_audience(vec![Role::User]); - - assert_eq!(content1.audience(), content2.audience()); - assert_eq!(content1.priority(), content2.priority()); - } - - #[test] - fn test_content_annotations_overwrite() { - let content = Content::text("hello") - .with_audience(vec![Role::User]) - .with_priority(0.5) - .with_audience(vec![Role::Assistant]) - .with_priority(0.8); - - assert_eq!(content.audience(), Some(&vec![Role::Assistant])); - assert_eq!(content.priority(), Some(0.8)); - } - - #[test] - fn test_content_annotations_image() { - let content = Content::image("data", "image/png") - .with_audience(vec![Role::User]) - .with_priority(0.5); - - assert_eq!(content.audience(), Some(&vec![Role::User])); - assert_eq!(content.priority(), Some(0.5)); - } - - #[test] - fn test_content_annotations_preservation() { - let text_content = Content::text("hello") - .with_audience(vec![Role::User]) - .with_priority(0.5); - - match &text_content { - Content::Text(TextContent { annotations, .. }) => { - assert!(annotations.is_some()); - let ann = annotations.as_ref().unwrap(); - assert_eq!(ann.audience, Some(vec![Role::User])); - assert_eq!(ann.priority, Some(0.5)); - } - _ => panic!("Expected Text content"), - } - } - - #[test] - #[should_panic(expected = "Priority must be between 0.0 and 1.0")] - fn test_invalid_priority() { - Content::text("hello").with_priority(1.5); - } - - #[test] - fn test_unannotated() { - let content = Content::text("hello") - .with_audience(vec![Role::User]) - .with_priority(0.5); - let unannotated = content.unannotated(); - assert_eq!(unannotated.audience(), None); - assert_eq!(unannotated.priority(), None); - } - - #[test] - fn test_partial_annotations() { - let content = Content::text("hello").with_priority(0.5); - assert_eq!(content.audience(), None); - assert_eq!(content.priority(), Some(0.5)); - - let content = Content::text("hello").with_audience(vec![Role::User]); - assert_eq!(content.audience(), Some(&vec![Role::User])); - assert_eq!(content.priority(), None); - } -} diff --git a/crates/mcp-core/src/lib.rs b/crates/mcp-core/src/lib.rs index 5a37ceea..4bd7d1ad 100644 --- a/crates/mcp-core/src/lib.rs +++ b/crates/mcp-core/src/lib.rs @@ -1,5 +1,3 @@ -pub mod content; -pub use content::{Annotations, Content, ImageContent, TextContent}; pub mod handler; pub mod role; pub use role::Role; diff --git a/crates/mcp-core/src/prompt.rs b/crates/mcp-core/src/prompt.rs index 4a0106e3..9b50396f 100644 --- a/crates/mcp-core/src/prompt.rs +++ b/crates/mcp-core/src/prompt.rs @@ -1,7 +1,6 @@ -use crate::content::{Annotations, EmbeddedResource, ImageContent}; use crate::handler::PromptError; -use crate::resource::ResourceContents; use base64::engine::{general_purpose::STANDARD as BASE64_STANDARD, Engine}; +use rmcp::model::{Annotations, EmbeddedResource, ImageContent}; use serde::{Deserialize, Serialize}; /// A prompt that can be used to generate text from a model @@ -113,8 +112,7 @@ impl PromptMessage { role, content: PromptMessageContent::Image { image: ImageContent { - data, - mime_type, + raw: rmcp::model::RawImageContent { data, mime_type }, annotations, }, }, @@ -129,7 +127,7 @@ impl PromptMessage { text: Option, annotations: Option, ) -> Self { - let resource_contents = ResourceContents::TextResourceContents { + let resource_contents = rmcp::model::ResourceContents::TextResourceContents { uri, mime_type: Some(mime_type), text: text.unwrap_or_default(), @@ -139,7 +137,9 @@ impl PromptMessage { role, content: PromptMessageContent::Resource { resource: EmbeddedResource { - resource: resource_contents, + raw: rmcp::model::RawEmbeddedResource { + resource: resource_contents, + }, annotations, }, }, diff --git a/crates/mcp-core/src/protocol.rs b/crates/mcp-core/src/protocol.rs index 202d514d..e7d1c12f 100644 --- a/crates/mcp-core/src/protocol.rs +++ b/crates/mcp-core/src/protocol.rs @@ -1,11 +1,11 @@ /// The protocol messages exchanged between client and server use crate::{ - content::Content, prompt::{Prompt, PromptMessage}, resource::Resource, resource::ResourceContents, tool::Tool, }; +use rmcp::model::Content; use serde::{Deserialize, Serialize}; use serde_json::Value; diff --git a/crates/mcp-core/src/resource.rs b/crates/mcp-core/src/resource.rs index ae9c0689..7ad35d71 100644 --- a/crates/mcp-core/src/resource.rs +++ b/crates/mcp-core/src/resource.rs @@ -1,7 +1,7 @@ -use crate::content::Annotations; /// Resources that servers provide to clients use anyhow::{anyhow, Result}; use chrono::{DateTime, Utc}; +use rmcp::model::Annotations; use serde::{Deserialize, Serialize}; use url::Url; use utoipa::ToSchema; diff --git a/crates/mcp-server/Cargo.toml b/crates/mcp-server/Cargo.toml index fbaf90ba..c6392b75 100644 --- a/crates/mcp-server/Cargo.toml +++ b/crates/mcp-server/Cargo.toml @@ -11,6 +11,7 @@ anyhow = "1.0.94" thiserror = "1.0" mcp-core = { path = "../mcp-core" } mcp-macros = { path = "../mcp-macros" } +rmcp = { workspace = true } serde = { version = "1.0.216", features = ["derive"] } serde_json = "1.0.133" schemars = "0.8" diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index ad23cd51..2d7edfe0 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -1,5 +1,4 @@ use anyhow::Result; -use mcp_core::content::Content; use mcp_core::handler::{PromptError, ResourceError}; use mcp_core::prompt::{Prompt, PromptArgument}; use mcp_core::protocol::JsonRpcMessage; @@ -7,6 +6,7 @@ use mcp_core::tool::ToolAnnotations; use mcp_core::{handler::ToolError, protocol::ServerCapabilities, resource::Resource, tool::Tool}; use mcp_server::router::{CapabilitiesBuilder, RouterService}; use mcp_server::{ByteTransport, Router, Server}; +use rmcp::model::Content; use serde_json::Value; use std::{future::Future, pin::Pin, sync::Arc}; use tokio::sync::mpsc; diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index 6370bd1f..06c294a6 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -7,7 +7,6 @@ use std::{ type PromptFuture = Pin> + Send + 'static>>; use mcp_core::{ - content::Content, handler::{PromptError, ResourceError, ToolError}, prompt::{Prompt, PromptMessage, PromptMessageRole}, protocol::{ @@ -18,6 +17,7 @@ use mcp_core::{ }, ResourceContents, }; +use rmcp::model::Content; use serde_json::Value; use tokio::sync::mpsc; use tower_service::Service; diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 59a29554..d891d653 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -934,18 +934,14 @@ "type": "array", "items": { "$ref": "#/components/schemas/Role" - }, - "nullable": true + } }, "priority": { - "type": "number", - "format": "float", - "nullable": true + "type": "number" }, "timestamp": { "type": "string", - "format": "date-time", - "example": "2023-01-01T00:00:00Z" + "format": "date-time" } } }, @@ -1002,72 +998,81 @@ "Content": { "oneOf": [ { - "allOf": [ - { - "$ref": "#/components/schemas/TextContent" + "type": "object", + "required": [ + "text", + "type" + ], + "properties": { + "text": { + "type": "string" }, - { - "type": "object", - "required": [ - "type" - ], - "properties": { - "type": { - "type": "string", - "enum": [ - "text" - ] - } - } + "type": { + "type": "string" } - ] + } }, { - "allOf": [ - { - "$ref": "#/components/schemas/ImageContent" + "type": "object", + "required": [ + "data", + "mimeType", + "type" + ], + "properties": { + "data": { + "type": "string" }, - { - "type": "object", - "required": [ - "type" - ], - "properties": { - "type": { - "type": "string", - "enum": [ - "image" - ] - } - } + "mimeType": { + "type": "string" + }, + "type": { + "type": "string" } - ] + } }, { - "allOf": [ - { - "$ref": "#/components/schemas/EmbeddedResource" + "type": "object", + "required": [ + "resource", + "type" + ], + "properties": { + "resource": { + "$ref": "#/components/schemas/ResourceContents" }, - { - "type": "object", - "required": [ - "type" - ], - "properties": { - "type": { - "type": "string", - "enum": [ - "resource" - ] - } - } + "type": { + "type": "string" } - ] + } + }, + { + "type": "object", + "required": [ + "data", + "mimeType", + "type" + ], + "properties": { + "annotations": { + "allOf": [ + { + "$ref": "#/components/schemas/Annotations" + } + ] + }, + "data": { + "type": "string" + }, + "mimeType": { + "type": "string" + }, + "type": { + "type": "string" + } + } } - ], - "discriminator": { - "propertyName": "type" - } + ] }, "ContextLengthExceeded": { "type": "object", @@ -1160,8 +1165,7 @@ { "$ref": "#/components/schemas/Annotations" } - ], - "nullable": true + ] }, "resource": { "$ref": "#/components/schemas/ResourceContents" @@ -1491,8 +1495,7 @@ { "$ref": "#/components/schemas/Annotations" } - ], - "nullable": true + ] }, "data": { "type": "string" @@ -1997,11 +2000,13 @@ ] }, "Role": { - "type": "string", - "description": "A wrapper around rmcp::model::Role that implements ToSchema for utoipa", - "enum": [ - "user", - "assistant" + "oneOf": [ + { + "type": "string" + }, + { + "type": "string" + } ] }, "RunNowResponse": { @@ -2285,8 +2290,7 @@ { "$ref": "#/components/schemas/Annotations" } - ], - "nullable": true + ] }, "text": { "type": "string" diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 9238f246..071c7bfd 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -1,8 +1,8 @@ // This file is auto-generated by @hey-api/openapi-ts export type Annotations = { - audience?: Array | null; - priority?: number | null; + audience?: Array; + priority?: number; timestamp?: string; }; @@ -22,13 +22,22 @@ export type ConfigResponse = { config: {}; }; -export type Content = (TextContent & { - type: 'text'; -}) | (ImageContent & { - type: 'image'; -}) | (EmbeddedResource & { - type: 'resource'; -}); +export type Content = { + text: string; + type: string; +} | { + data: string; + mimeType: string; + type: string; +} | { + resource: ResourceContents; + type: string; +} | { + annotations?: Annotations; + data: string; + mimeType: string; + type: string; +}; export type ContextLengthExceeded = { msg: string; @@ -70,7 +79,7 @@ export type CreateScheduleRequest = { }; export type EmbeddedResource = { - annotations?: Annotations | null; + annotations?: Annotations; resource: ResourceContents; }; @@ -186,7 +195,7 @@ export type FrontendToolRequest = { }; export type ImageContent = { - annotations?: Annotations | null; + annotations?: Annotations; data: string; mimeType: string; }; @@ -338,10 +347,7 @@ export type ResourceContents = { uri: string; }; -/** - * A wrapper around rmcp::model::Role that implements ToSchema for utoipa - */ -export type Role = 'user' | 'assistant'; +export type Role = string; export type RunNowResponse = { session_id: string; @@ -459,7 +465,7 @@ export type SummarizationRequested = { }; export type TextContent = { - annotations?: Annotations | null; + annotations?: Annotations; text: string; }; diff --git a/ui/desktop/src/components/context_management/index.ts b/ui/desktop/src/components/context_management/index.ts index 547b6652..6af637a0 100644 --- a/ui/desktop/src/components/context_management/index.ts +++ b/ui/desktop/src/components/context_management/index.ts @@ -59,7 +59,7 @@ export function convertApiMessageToFrontendMessage( display: display ?? true, sendToLLM: sendToLLM ?? true, id: generateId(), - role: apiMessage.role, + role: apiMessage.role as Role, created: apiMessage.created, content: apiMessage.content .map((apiContent) => mapApiContentToFrontendMessageContent(apiContent))