acp: ToolCallLocations and working cancellation (#5588)

This commit is contained in:
Alex Hancock
2025-11-05 23:11:35 -05:00
committed by GitHub
parent 147d884251
commit b2aa26a0f1
+218 -95
View File
@@ -7,10 +7,11 @@ use goose::agents::{Agent, SessionConfig};
use goose::config::{get_all_extensions, Config};
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::mcp_utils::ToolResult;
use goose::providers::create;
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use rmcp::model::{RawContent, ResourceContents};
use rmcp::model::{Content, RawContent, ResourceContents};
use std::collections::{HashMap, HashSet};
use std::fs;
use std::sync::Arc;
@@ -24,6 +25,7 @@ use url::Url;
struct GooseAcpSession {
messages: Conversation,
tool_call_ids: HashMap<String, String>, // Maps internal tool IDs to ACP tool call IDs
tool_requests: HashMap<String, goose::conversation::message::ToolRequest>, // Store tool requests by ID for location extraction
cancel_token: Option<CancellationToken>, // Active cancellation token for prompt processing
}
@@ -33,6 +35,116 @@ struct GooseAcpAgent {
agent: Agent, // Shared agent instance
}
/// Create a ToolCallLocation with common defaults
fn create_tool_location(path: &str, line: Option<u32>) -> acp::ToolCallLocation {
acp::ToolCallLocation {
path: path.into(),
line,
meta: None,
}
}
/// Extract file locations from tool request and response
fn extract_tool_locations(
tool_request: &goose::conversation::message::ToolRequest,
tool_response: &goose::conversation::message::ToolResponse,
) -> Vec<acp::ToolCallLocation> {
let mut locations = Vec::new();
// Get the tool call details
if let Ok(tool_call) = &tool_request.tool_call {
// Only process text_editor tool
if tool_call.name != "developer__text_editor" {
return locations;
}
// Extract the path from arguments
let path_str = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("path"))
.and_then(|p| p.as_str());
if let Some(path_str) = path_str {
// Get the command type
let command = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("command"))
.and_then(|c| c.as_str());
// Extract line numbers from the response content
if let Ok(content_items) = &tool_response.tool_result {
for content in content_items {
if let RawContent::Text(text_content) = &content.raw {
let text = &text_content.text;
// Parse line numbers based on command type and response format
match command {
Some("view") => {
// For view command, look for "lines X-Y" pattern in header
let line = extract_view_line_range(text)
.map(|range| range.0 as u32)
.or(Some(1));
locations.push(create_tool_location(path_str, line));
}
Some("str_replace") | Some("insert") => {
// For edits, extract the first line number from the snippet
let line = extract_first_line_number(text)
.map(|l| l as u32)
.or(Some(1));
locations.push(create_tool_location(path_str, line));
}
Some("write") => {
// For write, just point to the beginning of the file
locations.push(create_tool_location(path_str, Some(1)));
}
_ => {
// For other commands or unknown, default to line 1
locations.push(create_tool_location(path_str, Some(1)));
}
}
break; // Only process first text content
}
}
}
// If we didn't find any locations yet, add a default one
if locations.is_empty() {
locations.push(create_tool_location(path_str, Some(1)));
}
}
}
locations
}
/// Extract line range from view command output (e.g., "### path/to/file.rs (lines 10-20)")
fn extract_view_line_range(text: &str) -> Option<(usize, usize)> {
// Look for pattern like "(lines X-Y)" or "(lines X-end)"
let re = regex::Regex::new(r"\(lines (\d+)-(\d+|end)\)").ok()?;
if let Some(caps) = re.captures(text) {
let start = caps.get(1)?.as_str().parse::<usize>().ok()?;
let end = if caps.get(2)?.as_str() == "end" {
start // Use start as a reasonable default
} else {
caps.get(2)?.as_str().parse::<usize>().ok()?
};
return Some((start, end));
}
None
}
/// Extract the first line number from code snippet (e.g., "123: some code")
fn extract_first_line_number(text: &str) -> Option<usize> {
// Look for pattern like "123: " at the start of a line within a code block
let re = regex::Regex::new(r"```[^\n]*\n(\d+):").ok()?;
if let Some(caps) = re.captures(text) {
return caps.get(1)?.as_str().parse::<usize>().ok();
}
None
}
fn read_resource_link(link: acp::ResourceLink) -> Option<String> {
let url = Url::parse(&link.uri).ok()?;
if url.scheme() == "file" {
@@ -280,36 +392,19 @@ impl GooseAcpAgent {
.tool_call_ids
.insert(tool_request.id.clone(), acp_tool_id.clone());
// Extract tool name and parameters from the ToolCall if successful
let (tool_name, locations) = match &tool_request.tool_call {
Ok(tool_call) => {
// Extract file locations from certain tools for client tracking
let mut locs = Vec::new();
if tool_call.name == "developer__text_editor" {
// Try to extract the path from the arguments
if let Some(path_str) = tool_call
.arguments
.as_ref()
.and_then(|args_map| args_map.get("path"))
.and_then(|p| p.as_str())
{
let path = std::path::PathBuf::from(path_str);
if path.exists() && path.is_file() {
locs.push(acp::ToolCallLocation {
path: path_str.into(),
line: Some(1),
meta: None,
});
}
}
}
// Store the tool request for later use in response handling
session
.tool_requests
.insert(tool_request.id.clone(), tool_request.clone());
(tool_call.name.to_string(), locs)
}
Err(_) => ("error".to_string(), vec![]),
// Extract tool name from the ToolCall if successful
let tool_name = match &tool_request.tool_call {
Ok(tool_call) => tool_call.name.to_string(),
Err(_) => "error".to_string(),
};
// Send tool call notification
// Send tool call notification with empty locations initially
// We'll update with real locations when we get the response
let (tx, rx) = oneshot::channel();
self.session_update_tx
.send((
@@ -321,7 +416,7 @@ impl GooseAcpAgent {
kind: acp::ToolKind::default(),
status: acp::ToolCallStatus::Pending,
content: Vec::new(),
locations,
locations: Vec::new(), // Will be populated in handle_tool_response
raw_input: None,
raw_output: None,
meta: None,
@@ -351,74 +446,17 @@ impl GooseAcpAgent {
acp::ToolCallStatus::Failed
};
let content: Vec<ToolCallContent> = match &tool_response.tool_result {
Ok(content_items) => content_items
.iter()
.filter_map(|content| match &content.raw {
RawContent::Text(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Text(TextContent {
annotations: None,
text: val.text.clone(),
meta: None,
}),
}),
RawContent::Image(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Image(ImageContent {
annotations: None,
data: val.data.clone(),
mime_type: val.mime_type.clone(),
uri: None,
meta: None,
}),
}),
RawContent::Resource(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Resource(EmbeddedResource {
annotations: None,
resource: match &val.resource {
ResourceContents::TextResourceContents {
mime_type,
text,
uri,
..
} => acp::EmbeddedResourceResource::TextResourceContents(
acp::TextResourceContents {
mime_type: mime_type.clone(),
text: text.clone(),
uri: uri.clone(),
meta: None,
},
),
ResourceContents::BlobResourceContents {
mime_type,
blob,
uri,
..
} => acp::EmbeddedResourceResource::BlobResourceContents(
acp::BlobResourceContents {
mime_type: mime_type.clone(),
blob: blob.clone(),
uri: uri.clone(),
meta: None,
},
),
},
meta: None,
}),
}),
RawContent::Audio(_) => {
// Audio content is not supported in ACP ContentBlock, skip it
None
}
RawContent::ResourceLink(_) => {
// ResourceLink content is not supported in ACP ContentBlock, skip it
None
}
})
.collect(),
Err(_) => Vec::new(),
let content = build_tool_call_content(&tool_response.tool_result);
// Extract locations from the tool request and response
let locations = if let Some(tool_request) = session.tool_requests.get(&tool_response.id)
{
extract_tool_locations(tool_request, tool_response)
} else {
Vec::new()
};
// Send status update (completed or failed)
// Send status update (completed or failed) with locations
let (tx, rx) = oneshot::channel();
self.session_update_tx
.send((
@@ -429,6 +467,11 @@ impl GooseAcpAgent {
fields: acp::ToolCallUpdateFields {
status: Some(status),
content: Some(content),
locations: if locations.is_empty() {
None
} else {
Some(locations)
},
..Default::default()
},
meta: None,
@@ -445,6 +488,76 @@ impl GooseAcpAgent {
}
}
/// Build tool call content from tool result
fn build_tool_call_content(tool_result: &ToolResult<Vec<Content>>) -> Vec<ToolCallContent> {
match tool_result {
Ok(content_items) => content_items
.iter()
.filter_map(|content| match &content.raw {
RawContent::Text(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Text(TextContent {
annotations: None,
text: val.text.clone(),
meta: None,
}),
}),
RawContent::Image(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Image(ImageContent {
annotations: None,
data: val.data.clone(),
mime_type: val.mime_type.clone(),
uri: None,
meta: None,
}),
}),
RawContent::Resource(val) => Some(ToolCallContent::Content {
content: acp::ContentBlock::Resource(EmbeddedResource {
annotations: None,
resource: match &val.resource {
ResourceContents::TextResourceContents {
mime_type,
text,
uri,
..
} => acp::EmbeddedResourceResource::TextResourceContents(
acp::TextResourceContents {
mime_type: mime_type.clone(),
text: text.clone(),
uri: uri.clone(),
meta: None,
},
),
ResourceContents::BlobResourceContents {
mime_type,
blob,
uri,
..
} => acp::EmbeddedResourceResource::BlobResourceContents(
acp::BlobResourceContents {
mime_type: mime_type.clone(),
blob: blob.clone(),
uri: uri.clone(),
meta: None,
},
),
},
meta: None,
}),
}),
RawContent::Audio(_) => {
// Audio content is not supported in ACP ContentBlock, skip it
None
}
RawContent::ResourceLink(_) => {
// ResourceLink content is not supported in ACP ContentBlock, skip it
None
}
})
.collect(),
Err(_) => Vec::new(),
}
}
#[async_trait::async_trait(?Send)]
impl acp::Agent for GooseAcpAgent {
async fn initialize(
@@ -498,6 +611,7 @@ impl acp::Agent for GooseAcpAgent {
let session = GooseAcpSession {
messages: Conversation::new_unvalidated(Vec::new()),
tool_call_ids: HashMap::new(),
tool_requests: HashMap::new(),
cancel_token: None,
};
@@ -543,6 +657,15 @@ impl acp::Agent for GooseAcpAgent {
// Create and store cancellation token for this prompt
let cancel_token = CancellationToken::new();
// Store the cancel token in the session BEFORE starting the prompt
{
let mut sessions = self.sessions.lock().await;
let session = sessions
.get_mut(&session_id)
.ok_or_else(acp::Error::invalid_params)?;
session.cancel_token = Some(cancel_token.clone());
}
let user_message = self.convert_acp_prompt_to_message(args.prompt);
let session = SessionManager::create_session(