Tool reply meta (#6074)

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
+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);
}