acp: ToolCallLocations and working cancellation (#5588)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user