awesome-bookmarks/goose
3
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-17 12:56:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: adds default tools that can use MCP resources (#619)

Browse Source 
* adds 2 default "platform" tools to list and read MCP resources

* dynamically added when any of the connected extensions have resource capabilities

Co-authored-by: Kalvin Chau <kalvin@squareup.com>
This commit is contained in:
Salman Mohammed
2025-01-20 18:30:11 -05:00
committed by GitHub
parent b827a7e148
commit a22daca0c7
5 changed files with 259 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-mcp/src/developer2/mod.rs
+1 -4
View File
@@ -620,10 +620,7 @@ impl Router for Developer2Router {
}
fn capabilities(&self) -> ServerCapabilities {
CapabilitiesBuilder::new()
.with_tools(false)
.with_resources(false, false)
.build()
CapabilitiesBuilder::new().with_tools(false).build()
}
fn list_tools(&self) -> Vec<Tool> {
crates/goose/src/agents/capabilities.rs
+201 -16
View File
@@ -1,6 +1,7 @@
use chrono::{DateTime, TimeZone, Utc};
use futures::stream::{FuturesUnordered, StreamExt};
use mcp_client::McpService;
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::LazyLock;
use std::time::Duration;
@@ -13,6 +14,7 @@ use crate::providers::base::{Provider, ProviderUsage};
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport};
use mcp_core::{Content, Tool, ToolCall, ToolError, ToolResult};
use serde_json::Value;
// By default, we set it to Jan 1, 2020 if the resource does not have a timestamp
// This is to ensure that the resource is considered less important than resources with a more recent timestamp
@@ -23,6 +25,7 @@ static DEFAULT_TIMESTAMP: LazyLock<DateTime<Utc>> =
pub struct Capabilities {
clients: HashMap<String, Arc<Mutex<Box<dyn McpClientTrait>>>>,
instructions: HashMap<String, String>,
resource_capable_systems: HashSet<String>,
provider: Box<dyn Provider>,
provider_usage: Mutex<Vec<ProviderUsage>>,
}
@@ -80,11 +83,16 @@ impl Capabilities {
Self {
clients: HashMap::new(),
instructions: HashMap::new(),
resource_capable_systems: HashSet::new(),
provider,
provider_usage: Mutex::new(Vec::new()),
}
}
pub fn supports_resources(&self) -> bool {
!self.resource_capable_systems.is_empty()
}
/// Add a new MCP system based on the provided client type
// TODO IMPORTANT need to ensure this times out if the system command is broken!
pub async fn add_system(&mut self, config: SystemConfig) -> SystemResult<()> {
@@ -141,6 +149,12 @@ impl Capabilities {
.insert(init_result.server_info.name.clone(), instructions);
}
// if the server is capable if resources we track it
if init_result.capabilities.resources.is_some() {
self.resource_capable_systems
.insert(sanitize(init_result.server_info.name.clone()));
}
// Store the client
self.clients.insert(
sanitize(init_result.server_info.name.clone()),
@@ -278,7 +292,8 @@ impl Capabilities {
.keys()
.map(|name| {
let instructions = self.instructions.get(name).cloned().unwrap_or_default();
SystemInfo::new(name, "", &instructions)
let has_resources = self.resource_capable_systems.contains(name);
SystemInfo::new(name, &instructions, has_resources)
})
.collect();
@@ -297,25 +312,195 @@ impl Capabilities {
.map(Arc::clone)
}
// Function that gets executed for read_resource tool
async fn read_resource(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let uri = params
.get("uri")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("Missing 'uri' parameter".to_string()))?;
let system_name = params.get("system_name").and_then(|v| v.as_str());
// If system name is provided, we can just look it up
if system_name.is_some() {
let result = self
.read_resource_from_system(uri, system_name.unwrap())
.await?;
return Ok(result);
}
// If system name is not provided, we need to search for the resource across all systems
// Loop through each system and try to read the resource, don't raise an error if the resource is not found
// TODO: do we want to find if a provided uri is in multiple systems?
// currently it will reutrn the first match and skip any systems
for system_name in self.resource_capable_systems.iter() {
let result = self.read_resource_from_system(uri, system_name).await;
match result {
Ok(result) => return Ok(result),
Err(_) => continue,
}
}
// None of the systems had the resource so we raise an error
let available_systems = self
.clients
.keys()
.map(|s| s.as_str())
.collect::<Vec<&str>>()
.join(", ");
let error_msg = format!(
"Resource with uri '{}' not found. Here are the available systems: {}",
uri, available_systems
);
Err(ToolError::InvalidParameters(error_msg))
}
async fn read_resource_from_system(
&self,
uri: &str,
system_name: &str,
) -> Result<Vec<Content>, ToolError> {
let available_systems = self
.clients
.keys()
.map(|s| s.as_str())
.collect::<Vec<&str>>()
.join(", ");
let error_msg = format!(
"System '{}' not found. Here are the available systems: {}",
system_name, available_systems
);
let client = self
.clients
.get(system_name)
.ok_or(ToolError::InvalidParameters(error_msg))?;
let client_guard = client.lock().await;
let read_result = client_guard.read_resource(uri).await.map_err(|_| {
ToolError::ExecutionError(format!("Could not read resource with uri: {}", uri))
})?;
let mut result = Vec::new();
for content in read_result.contents {
// Only reading the text resource content; skipping the blob content cause it's too long
if let mcp_core::resource::ResourceContents::TextResourceContents { text, .. } = content
{
let content_str = format!("{}\n\n{}", uri, text);
result.push(Content::text(content_str));
}
}
Ok(result)
}
async fn list_resources_from_system(
&self,
system_name: &str,
) -> Result<Vec<Content>, ToolError> {
let client = self.clients.get(system_name).ok_or_else(|| {
ToolError::InvalidParameters(format!("System {} is not valid", system_name))
})?;
let client_guard = client.lock().await;
client_guard
.list_resources(None)
.await
.map_err(|e| {
ToolError::ExecutionError(format!(
"Unable to list resources for {}, {:?}",
system_name, e
))
})
.map(|lr| {
let resource_list = lr
.resources
.into_iter()
.map(|r| format!("{} - {}, uri: ({})", system_name, r.name, r.uri))
.collect::<Vec<String>>()
.join("\n");
vec![Content::text(resource_list)]
})
}
async fn list_resources(&self, params: Value) -> Result<Vec<Content>, ToolError> {
let system = params.get("system").and_then(|v| v.as_str());
match system {
Some(system_name) => {
// Handle single system case
self.list_resources_from_system(system_name).await
}
None => {
// Handle all systems case using FuturesUnordered
let mut futures = FuturesUnordered::new();
// Create futures for each resource_capable_system
for system_name in &self.resource_capable_systems {
futures.push(async move { self.list_resources_from_system(system_name).await });
}
let mut all_resources = Vec::new();
let mut errors = Vec::new();
// Process results as they complete
while let Some(result) = futures.next().await {
match result {
Ok(content) => {
all_resources.extend(content);
}
Err(tool_error) => {
errors.push(tool_error);
}
}
}
// Log any errors that occurred
if !errors.is_empty() {
tracing::error!(
errors = ?errors
.into_iter()
.map(|e| format!("{:?}", e))
.collect::<Vec<_>>(),
"errors from listing resources"
);
}
Ok(all_resources)
}
}
}
/// Dispatch a single tool call to the appropriate client
#[instrument(skip(self, tool_call), fields(input, output))]
pub async fn dispatch_tool_call(&self, tool_call: ToolCall) -> ToolResult<Vec<Content>> {
let client = self
.get_client_for_tool(&tool_call.name)
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
let result = if tool_call.name == "platform__read_resource" {
// Check if the tool is read_resource and handle it separately
self.read_resource(tool_call.arguments.clone()).await
} else if tool_call.name == "platform__list_resources" {
self.list_resources(tool_call.arguments.clone()).await
} else {
// Else, dispatch tool call based on the prefix naming convention
let client = self
.get_client_for_tool(&tool_call.name)
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
let tool_name = tool_call
.name
.split("__")
.nth(1)
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
let tool_name = tool_call
.name
.split("__")
.nth(1)
.ok_or_else(|| ToolError::NotFound(tool_call.name.clone()))?;
let client_guard = client.lock().await;
let result = client_guard
.call_tool(tool_name, tool_call.clone().arguments)
.await
.map(|result| result.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()));
let client_guard = client.lock().await;
client_guard
.call_tool(tool_name, tool_call.clone().arguments)
.await
.map(|result| result.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()))
};
debug!(
"input" = serde_json::to_string(&tool_call).unwrap(),
crates/goose/src/agents/reference.rs
+50 -2
View File
@@ -13,7 +13,10 @@ use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
use crate::register_agent;
use crate::token_counter::TokenCounter;
use serde_json::Value;
use indoc::indoc;
use mcp_core::tool::Tool;
use serde_json::{json, Value};
/// Reference implementation of an Agent
pub struct ReferenceAgent {
capabilities: Mutex<Capabilities>,
@@ -66,7 +69,52 @@ impl Agent for ReferenceAgent {
let mut messages = messages.to_vec();
let reply_span = tracing::Span::current();
let mut capabilities = self.capabilities.lock().await;
let tools = capabilities.get_prefixed_tools().await?;
let mut tools = capabilities.get_prefixed_tools().await?;
// we add in the read_resource tool by default
// TODO: make sure there is no collision with another system's tool name
let read_resource_tool = Tool::new(
"platform__read_resource".to_string(),
indoc! {r#"
Read a resource from a system.
Resources allow systems to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool searches for the
resource URI in the provided system, and reads in the resource content. If no system
is provided, the tool will search all systems for the resource.
"#}.to_string(),
json!({
"type": "object",
"required": ["uri"],
"properties": {
"uri": {"type": "string", "description": "Resource URI"},
"system_name": {"type": "string", "description": "Optional system name"}
}
}),
);
let list_resources_tool = Tool::new(
"platform__list_resources".to_string(),
indoc! {r#"
List resources from a system(s).
Resources allow systems to share data that provide context to LLMs, such as
files, database schemas, or application-specific information. This tool lists resources
in the provided system, and returns a list for the user to browse. If no system
is provided, the tool will search all systems for the resource.
"#}.to_string(),
json!({
"type": "object",
"properties": {
"system_name": {"type": "string", "description": "Optional system name"}
}
}),
);
if capabilities.supports_resources() {
tools.push(read_resource_tool);
tools.push(list_resources_tool);
}
let system_prompt = capabilities.get_system_prompt().await;
let _estimated_limit = capabilities
.provider()
crates/goose/src/agents/system.rs
+3 -3
View File
@@ -114,16 +114,16 @@ impl std::fmt::Display for SystemConfig {
#[derive(Clone, Debug, Serialize)]
pub struct SystemInfo {
name: String,
description: String,
instructions: String,
has_resources: bool,
}
impl SystemInfo {
pub fn new(name: &str, description: &str, instructions: &str) -> Self {
pub fn new(name: &str, instructions: &str, has_resources: bool) -> Self {
Self {
name: name.to_string(),
description: description.to_string(),
instructions: instructions.to_string(),
has_resources,
}
}
}
crates/goose/src/prompts/system.md
+4 -2
View File
@@ -14,8 +14,10 @@ in your tool specification.
{% for system in systems %}
## {{system.name}}
{{system.description}}
{% if system.has_resources %}
{{system.name}} supports resources, you can use platform__read_resource,
and platform__list_resources on this system.
{% endif %}
{% if system.instructions %}### Instructions
{{system.instructions}}{% endif %}
{% endfor %}