tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Tool reply meta (#6074)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-12-11 19:43:36 -05:00
committed by GitHub
parent 7dd244eff6
commit cdab2fc3a6
52 changed files with 1852 additions and 1083 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+25 -8
View File
@@ -58,7 +58,7 @@ use crate::tool_monitor::RepetitionInspector;
use crate::utils::is_token_cancelled;
use regex::Regex;
use rmcp::model::{
CallToolRequestParam, Content, ErrorCode, ErrorData, GetPromptResult, Prompt,
CallToolRequestParam, CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Prompt,
ServerNotification, Tool,
};
use serde_json::Value;
@@ -100,7 +100,7 @@ pub struct Agent {
pub(super) prompt_manager: Mutex<PromptManager>,
pub(super) confirmation_tx: mpsc::Sender<(String, PermissionConfirmation)>,
pub(super) confirmation_rx: Mutex<mpsc::Receiver<(String, PermissionConfirmation)>>,
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<CallToolResult>)>,
pub(super) tool_result_rx: ToolResultReceiver,
pub tool_route_manager: Arc<ToolRouteManager>,
@@ -128,7 +128,8 @@ pub enum ToolStreamItem<T> {
Result(T),
}
pub type ToolStream = Pin<Box<dyn Stream<Item = ToolStreamItem<ToolResult<Vec<Content>>>> + Send>>;
pub type ToolStream =
Pin<Box<dyn Stream<Item = ToolStreamItem<ToolResult<CallToolResult>>> + Send>>;
// tool_stream combines a stream of ServerNotifications with a future representing the
// final result of the tool call. MCP notifications are not request-scoped, but
@@ -137,7 +138,7 @@ pub type ToolStream = Pin<Box<dyn Stream<Item = ToolStreamItem<ToolResult<Vec<Co
pub fn tool_stream<S, F>(rx: S, done: F) -> ToolStream
where
S: Stream<Item = ServerNotification> + Send + Unpin + 'static,
F: Future<Output = ToolResult<Vec<Content>>> + Send + 'static,
F: Future<Output = ToolResult<CallToolResult>> + Send + 'static,
{
Box::pin(async_stream::stream! {
tokio::pin!(done);
@@ -360,7 +361,12 @@ impl Agent {
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response(
request.id.clone(),
Ok(vec![rmcp::model::Content::text(DECLINED_RESPONSE)]),
Ok(CallToolResult {
content: vec![rmcp::model::Content::text(DECLINED_RESPONSE)],
structured_content: None,
is_error: Some(true),
meta: None,
}),
);
}
}
@@ -454,7 +460,13 @@ impl Agent {
let result = self
.handle_schedule_management(arguments, request_id.clone())
.await;
return (request_id, Ok(ToolCallResult::from(result)));
let wrapped_result = result.map(|content| CallToolResult {
content,
structured_content: None,
is_error: Some(false),
meta: None,
});
return (request_id, Ok(ToolCallResult::from(wrapped_result)));
}
if tool_call.name == FINAL_OUTPUT_TOOL_NAME {
@@ -1090,7 +1102,12 @@ impl Agent {
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response(
request.id.clone(),
Ok(vec![Content::text(CHAT_MODE_TOOL_SKIPPED_RESPONSE)]),
Ok(CallToolResult {
content: vec![Content::text(CHAT_MODE_TOOL_SKIPPED_RESPONSE)],
structured_content: None,
is_error: Some(false),
meta: None,
}),
);
}
}
@@ -1421,7 +1438,7 @@ impl Agent {
Ok(plan_prompt)
}
pub async fn handle_tool_result(&self, id: String, result: ToolResult<Vec<Content>>) {
pub async fn handle_tool_result(&self, id: String, result: ToolResult<CallToolResult>) {
if let Err(e) = self.tool_result_tx.send((id, result)).await {
error!("Failed to send tool result: {}", e);
}
crates/goose/src/agents/extension_manager.rs
+74 -8
View File
@@ -41,8 +41,8 @@ use crate::oauth::oauth_flow;
use crate::prompt_template;
use crate::subprocess::configure_command_no_window;
use rmcp::model::{
CallToolRequestParam, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, ResourceContents,
ServerInfo, Tool,
CallToolRequestParam, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, RawContent,
Resource, ResourceContents, ServerInfo, Tool,
};
use rmcp::transport::auth::AuthClient;
use schemars::_private::NoSerialize;
@@ -758,6 +758,7 @@ impl ExtensionManager {
uri,
extension_name.unwrap(),
cancellation_token.clone(),
true,
)
.await?;
return Ok(result);
@@ -773,7 +774,12 @@ impl ExtensionManager {
for extension_name in extension_names {
let result = self
.read_resource_from_extension(uri, &extension_name, cancellation_token.clone())
.read_resource_from_extension(
uri,
&extension_name,
cancellation_token.clone(),
true,
)
.await;
match result {
Ok(result) => return Ok(result),
@@ -807,6 +813,7 @@ impl ExtensionManager {
uri: &str,
extension_name: &str,
cancellation_token: CancellationToken,
format_with_uri: bool,
) -> Result<Vec<Content>, ErrorData> {
let available_extensions = self
.extensions
@@ -840,9 +847,12 @@ impl ExtensionManager {
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 ResourceContents::TextResourceContents { text, .. } = content {
let content_str = format!("{}\n\n{}", uri, text);
let content_str = if format_with_uri {
format!("{}\n\n{}", uri, text)
} else {
text
};
result.push(Content::text(content_str));
}
}
@@ -850,6 +860,65 @@ impl ExtensionManager {
Ok(result)
}
pub async fn get_ui_resources(&self) -> Result<Vec<(String, Resource)>, ErrorData> {
let mut ui_resources = Vec::new();
let extensions_to_check: Vec<(String, McpClientBox)> = {
let extensions = self.extensions.lock().await;
extensions
.iter()
.map(|(name, ext)| (name.clone(), ext.get_client()))
.collect()
};
for (extension_name, client) in extensions_to_check {
let client_guard = client.lock().await;
match client_guard
.list_resources(None, CancellationToken::default())
.await
{
Ok(list_response) => {
for resource in list_response.resources {
if resource.uri.starts_with("ui://") {
ui_resources.push((extension_name.clone(), resource));
}
}
}
Err(e) => {
warn!("Failed to list resources for {}: {:?}", extension_name, e);
}
}
}
Ok(ui_resources)
}
pub async fn read_ui_resource(
&self,
uri: &str,
extension_name: &str,
cancellation_token: CancellationToken,
) -> Result<String, ErrorData> {
let contents = self
.read_resource_from_extension(uri, extension_name, cancellation_token, false)
.await?;
contents
.into_iter()
.find_map(|c| match c.raw {
RawContent::Text(text_content) => Some(text_content.text),
_ => None,
})
.ok_or_else(|| {
ErrorData::new(
ErrorCode::RESOURCE_NOT_FOUND,
format!("No text content in resource '{}'", uri),
None,
)
})
}
async fn list_resources_from_extension(
&self,
extension_name: &str,
@@ -936,7 +1005,6 @@ impl ExtensionManager {
}
}
// Log any errors that occurred
if !errors.is_empty() {
tracing::error!(
errors = ?errors
@@ -998,7 +1066,6 @@ impl ExtensionManager {
client_guard
.call_tool(&tool_name, arguments, cancellation_token)
.await
.map(|call| call.content)
.map_err(|e| match e {
ServiceError::McpError(error_data) => error_data,
_ => {
@@ -1077,7 +1144,6 @@ impl ExtensionManager {
}
}
// Log any errors that occurred
if !errors.is_empty() {
tracing::debug!(
errors = ?errors
crates/goose/src/agents/final_output_tool.rs
+8 -3
View File
@@ -123,9 +123,14 @@ impl FinalOutputTool {
match result {
Ok(parsed_value) => {
self.final_output = Some(Self::parsed_final_output_string(parsed_value));
ToolCallResult::from(Ok(vec![Content::text(
"Final output successfully collected.".to_string(),
)]))
ToolCallResult::from(Ok(rmcp::model::CallToolResult {
content: vec![Content::text(
"Final output successfully collected.".to_string(),
)],
structured_content: None,
is_error: Some(false),
meta: None,
}))
}
Err(error) => ToolCallResult::from(Err(ErrorData {
code: ErrorCode::INVALID_PARAMS,
crates/goose/src/agents/large_response_handler.rs
+42 -21
View File
@@ -1,5 +1,5 @@
use chrono::Utc;
use rmcp::model::{Content, ErrorData};
use rmcp::model::{CallToolResult, Content, ErrorData};
use std::fs::File;
use std::io::Write;
@@ -7,13 +7,13 @@ const LARGE_TEXT_THRESHOLD: usize = 200_000;
/// Process tool response and handle large text content
pub fn process_tool_response(
response: Result<Vec<Content>, ErrorData>,
) -> Result<Vec<Content>, ErrorData> {
response: Result<CallToolResult, ErrorData>,
) -> Result<CallToolResult, ErrorData> {
match response {
Ok(contents) => {
Ok(mut result) => {
let mut processed_contents = Vec::new();
for content in contents {
for content in result.content {
match content.as_text() {
Some(text_content) => {
// Check if text exceeds threshold
@@ -51,7 +51,8 @@ pub fn process_tool_response(
}
}
Ok(processed_contents)
result.content = processed_contents;
Ok(result)
}
Err(e) => Err(e),
}
@@ -89,14 +90,19 @@ mod tests {
let small_text = "This is a small text response";
let content = Content::text(small_text.to_string());
let response = Ok(vec![content]);
let response = Ok(CallToolResult {
content: vec![content],
structured_content: None,
is_error: Some(false),
meta: None,
});
// Process the response
let processed = process_tool_response(response).unwrap();
// Verify the response is unchanged
assert_eq!(processed.len(), 1);
if let Some(text_content) = processed[0].as_text() {
assert_eq!(processed.content.len(), 1);
if let Some(text_content) = processed.content[0].as_text() {
assert_eq!(text_content.text, small_text);
} else {
panic!("Expected text content");
@@ -109,14 +115,19 @@ mod tests {
let large_text = "a".repeat(LARGE_TEXT_THRESHOLD + 1000);
let content = Content::text(large_text.clone());
let response = Ok(vec![content]);
let response = Ok(CallToolResult {
content: vec![content],
structured_content: None,
is_error: Some(false),
meta: None,
});
// Process the response
let processed = process_tool_response(response).unwrap();
// Verify the response contains a message about the file
assert_eq!(processed.len(), 1);
if let Some(text_content) = processed[0].as_text() {
assert_eq!(processed.content.len(), 1);
if let Some(text_content) = processed.content[0].as_text() {
assert!(text_content
.text
.contains("The response returned from the tool call was larger"));
@@ -146,14 +157,19 @@ mod tests {
// Create an image content
let image_content = Content::image("base64data".to_string(), "image/png".to_string());
let response = Ok(vec![image_content]);
let response = Ok(CallToolResult {
content: vec![image_content],
structured_content: None,
is_error: Some(false),
meta: None,
});
// Process the response
let processed = process_tool_response(response).unwrap();
// Verify the response is unchanged
assert_eq!(processed.len(), 1);
if let Some(img) = processed[0].as_image() {
assert_eq!(processed.content.len(), 1);
if let Some(img) = processed.content[0].as_image() {
assert_eq!(img.data, "base64data");
assert_eq!(img.mime_type, "image/png");
} else {
@@ -168,23 +184,28 @@ mod tests {
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]);
let response = Ok(CallToolResult {
content: vec![small_text, large_text, image],
structured_content: None,
is_error: Some(false),
meta: None,
});
// Process the response
let processed = process_tool_response(response).unwrap();
// Verify each item is handled correctly
assert_eq!(processed.len(), 3);
assert_eq!(processed.content.len(), 3);
// First item should be unchanged small text
if let Some(text_content) = processed[0].as_text() {
if let Some(text_content) = processed.content[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 Some(text_content) = processed[1].as_text() {
if let Some(text_content) = processed.content[1].as_text() {
assert!(text_content
.text
.contains("The response returned from the tool call was larger"));
@@ -201,7 +222,7 @@ mod tests {
}
// Third item should be unchanged image
if let Some(img) = processed[2].as_image() {
if let Some(img) = processed.content[2].as_image() {
assert_eq!(img.data, "image_data");
assert_eq!(img.mime_type, "image/jpeg");
} else {
@@ -217,7 +238,7 @@ mod tests {
message: Cow::from("Test error"),
data: None,
};
let response: Result<Vec<Content>, ErrorData> = Err(error);
let response: Result<CallToolResult, ErrorData> = Err(error);
// Process the response
let processed = process_tool_response(response);
crates/goose/src/agents/moim.rs
+18 -2
View File
@@ -106,7 +106,15 @@ mod tests {
arguments: None,
}),
),
Message::user().with_tool_response("search_1", Ok(vec![])),
Message::user().with_tool_response(
"search_1",
Ok(rmcp::model::CallToolResult {
content: vec![],
structured_content: None,
is_error: Some(false),
meta: None,
}),
),
Message::assistant()
.with_text("I need to search more")
.with_tool_request(
@@ -116,7 +124,15 @@ mod tests {
arguments: None,
}),
),
Message::user().with_tool_response("search_2", Ok(vec![])),
Message::user().with_tool_response(
"search_2",
Ok(rmcp::model::CallToolResult {
content: vec![],
structured_content: None,
is_error: Some(false),
meta: None,
}),
),
]);
let result = inject_moim(conv, &em).await;
crates/goose/src/agents/recipe_tools/dynamic_task_tools.rs
+6 -1
View File
@@ -370,5 +370,10 @@ pub async fn create_dynamic_task(
};
tasks_manager.save_tasks(tasks).await;
ToolCallResult::from(Ok(vec![Content::text(tasks_json)]))
ToolCallResult::from(Ok(rmcp::model::CallToolResult {
content: vec![Content::text(tasks_json)],
structured_content: None,
is_error: Some(false),
meta: None,
}))
}
crates/goose/src/agents/sub_recipe_manager.rs
+6 -1
View File
@@ -63,7 +63,12 @@ impl SubRecipeManager {
.call_sub_recipe_tool(tool_name, params, tasks_manager, parent_working_dir)
.await;
match result {
Ok(call_result) => ToolCallResult::from(Ok(call_result)),
Ok(content) => ToolCallResult::from(Ok(rmcp::model::CallToolResult {
content,
structured_content: None,
is_error: Some(false),
meta: None,
})),
Err(e) => ToolCallResult::from(Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(e.to_string()),
crates/goose/src/agents/subagent_execution_tool/subagent_execute_task_tool.rs
+6 -1
View File
@@ -82,7 +82,12 @@ pub async fn run_tasks(
{
Ok(result) => {
let output = serde_json::to_string(&result).unwrap();
Ok(vec![Content::text(output)])
Ok(rmcp::model::CallToolResult {
content: vec![Content::text(output)],
structured_content: None,
is_error: Some(false),
meta: None,
})
}
Err(e) => Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
crates/goose/src/agents/subagent_handler.rs
+3 -2
View File
@@ -61,8 +61,9 @@ pub async fn run_complete_subagent_task(
tool_response,
) => {
// Extract text from tool response
if let Ok(contents) = &tool_response.tool_result {
let texts: Vec<String> = contents
if let Ok(result) = &tool_response.tool_result {
let texts: Vec<String> = result
.content
.iter()
.filter_map(|content| {
if let rmcp::model::RawContent::Text(raw_text_content) =
crates/goose/src/agents/tool_execution.rs
+9 -4
View File
@@ -16,12 +16,12 @@ use rmcp::model::{Content, ServerNotification};
// ToolCallResult combines the result of a tool call with an optional notification stream that
// can be used to receive notifications from the tool.
pub struct ToolCallResult {
pub result: Box<dyn Future<Output = ToolResult<Vec<Content>>> + Send + Unpin>,
pub result: Box<dyn Future<Output = ToolResult<rmcp::model::CallToolResult>> + Send + Unpin>,
pub notification_stream: Option<Box<dyn Stream<Item = ServerNotification> + Send + Unpin>>,
}
impl From<ToolResult<Vec<Content>>> for ToolCallResult {
fn from(result: ToolResult<Vec<Content>>) -> Self {
impl From<ToolResult<rmcp::model::CallToolResult>> for ToolCallResult {
fn from(result: ToolResult<rmcp::model::CallToolResult>) -> Self {
Self {
result: Box::new(futures::future::ready(result)),
notification_stream: None,
@@ -122,7 +122,12 @@ impl Agent {
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response(
request.id.clone(),
Ok(vec![Content::text(DECLINED_RESPONSE)]),
Ok(rmcp::model::CallToolResult {
content: vec![Content::text(DECLINED_RESPONSE)],
structured_content: None,
is_error: Some(true),
meta: None,
}),
);
}
}
crates/goose/src/agents/tool_route_manager.rs
+6 -1
View File
@@ -56,7 +56,12 @@ impl ToolRouteManager {
let selector = self.router_tool_selector.lock().await.clone();
match selector.as_ref() {
Some(selector) => match selector.select_tools(arguments).await {
Ok(tools) => Ok(ToolCallResult::from(Ok(tools))),
Ok(content) => Ok(ToolCallResult::from(Ok(rmcp::model::CallToolResult {
content,
structured_content: None,
is_error: Some(false),
meta: None,
}))),
Err(e) => Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to select tools: {}", e),
crates/goose/src/agents/types.rs
+2 -2
View File
@@ -1,13 +1,13 @@
use crate::mcp_utils::ToolResult;
use crate::providers::base::Provider;
use rmcp::model::{Content, Tool};
use rmcp::model::{CallToolResult, Tool};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use utoipa::ToSchema;
/// Type alias for the tool result channel receiver
pub type ToolResultReceiver = Arc<Mutex<mpsc::Receiver<(String, ToolResult<Vec<Content>>)>>>;
pub type ToolResultReceiver = Arc<Mutex<mpsc::Receiver<(String, ToolResult<CallToolResult>)>>>;
// We use double Arc here to allow easy provider swaps while sharing concurrent access
pub type SharedProvider = Arc<Mutex<Option<Arc<dyn Provider>>>>;