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

chore: small refactor on agent.rs (#3703)

Browse Source
This commit is contained in:
Lifei Zhou
2025-07-30 09:22:19 +10:00
committed by GitHub
parent 0fb58c4095
commit 1544e1fb27
2 changed files with 168 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+166 -113
View File
@@ -1,4 +1,4 @@
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@@ -35,8 +35,8 @@ use crate::agents::tool_vectordb::generate_table_id;
use crate::agents::types::SessionConfig;
use crate::agents::types::{FrontendTool, ToolResultReceiver};
use crate::config::{Config, ExtensionConfigManager, PermissionManager};
use crate::message::{push_message, Message};
use crate::permission::permission_judge::check_tool_permissions;
use crate::message::{push_message, Message, ToolRequest};
use crate::permission::permission_judge::{check_tool_permissions, PermissionCheckResult};
use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
@@ -61,6 +61,26 @@ use crate::conversation_fixer::{debug_conversation_fix, ConversationFixer};
const DEFAULT_MAX_TURNS: u32 = 1000;
/// Context needed for the reply function
pub struct ReplyContext {
pub messages: Vec<Message>,
pub tools: Vec<Tool>,
pub toolshim_tools: Vec<Tool>,
pub system_prompt: String,
pub goose_mode: String,
pub initial_messages: Vec<Message>,
pub config: &'static Config,
}
/// Result of processing tool requests
pub struct ToolProcessingResult {
pub frontend_requests: Vec<ToolRequest>,
pub remaining_requests: Vec<ToolRequest>,
pub filtered_response: Message,
pub readonly_tools: HashSet<String>,
pub regular_tools: HashSet<String>,
}
/// The main goose Agent
pub struct Agent {
pub(super) provider: Mutex<Option<Arc<dyn Provider>>>,
@@ -162,17 +182,6 @@ impl Agent {
*tool_monitor = Some(ToolMonitor::new(max_repetitions));
}
pub async fn get_tool_stats(&self) -> Option<HashMap<String, u32>> {
let tool_monitor = self.tool_monitor.lock().await;
tool_monitor.as_ref().map(|monitor| monitor.get_stats())
}
pub async fn reset_tool_monitor(&self) {
if let Some(monitor) = self.tool_monitor.lock().await.as_mut() {
monitor.reset();
}
}
/// Reset the retry attempts counter to 0
pub async fn reset_retry_attempts(&self) {
self.retry_manager.reset_attempts().await;
@@ -208,6 +217,119 @@ impl Agent {
}
}
async fn prepare_reply_context(
&self,
unfixed_messages: &[Message],
session: &Option<SessionConfig>,
) -> Result<ReplyContext> {
let (messages, issues) = ConversationFixer::fix_conversation(Vec::from(unfixed_messages));
if !issues.is_empty() {
tracing::warn!(
"Conversation issue fixed: {}",
debug_conversation_fix(&messages, unfixed_messages, &issues)
);
}
let initial_messages = messages.clone();
let config = Config::global();
let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?;
let goose_mode = Self::determine_goose_mode(session.as_ref(), config);
Ok(ReplyContext {
messages,
tools,
toolshim_tools,
system_prompt,
goose_mode,
initial_messages,
config,
})
}
/// Process tool requests by categorizing them and recording them in the router selector
async fn process_tool_requests(
&self,
response: &Message,
tools: &[rmcp::model::Tool],
) -> ToolProcessingResult {
let (readonly_tools, regular_tools) = Self::categorize_tools_by_annotation(tools);
// Categorize tool requests
let (frontend_requests, remaining_requests, filtered_response) =
self.categorize_tool_requests(response).await;
// Record tool calls in the router selector
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
for request in &frontend_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await {
error!("Failed to record frontend tool call: {}", e);
}
}
}
for request in &remaining_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await {
error!("Failed to record tool call: {}", e);
}
}
}
}
ToolProcessingResult {
frontend_requests,
remaining_requests,
filtered_response,
readonly_tools,
regular_tools,
}
}
async fn handle_approved_and_denied_tools(
&self,
permission_check_result: &PermissionCheckResult,
message_tool_response: Arc<Mutex<Message>>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
) -> Result<Vec<(String, ToolStream)>> {
let mut tool_futures: Vec<(String, ToolStream)> = Vec::new();
// Handle pre-approved and read-only tools
for request in &permission_check_result.approved {
if let Ok(tool_call) = request.tool_call.clone() {
let (req_id, tool_result) = self
.dispatch_tool_call(tool_call, request.id.clone(), cancel_token.clone())
.await;
tool_futures.push((
req_id,
match tool_result {
Ok(result) => tool_stream(
result
.notification_stream
.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => {
tool_stream(Box::new(stream::empty()), futures::future::ready(Err(e)))
}
},
));
}
}
// Handle denied tools
for request in &permission_check_result.denied {
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(
request.id.clone(),
Ok(vec![rmcp::model::Content::text(DECLINED_RESPONSE)]),
);
}
Ok(tool_futures)
}
/// Set the scheduler service for this agent
pub async fn set_scheduler(&self, scheduler: Arc<dyn SchedulerTrait>) {
let mut scheduler_service = self.scheduler_service.lock().await;
@@ -232,24 +354,6 @@ impl Agent {
self.frontend_tools.lock().await.get(name).cloned()
}
/// Get all tools from all clients with proper prefixing
pub async fn get_prefixed_tools(&self) -> ExtensionResult<Vec<Tool>> {
let mut tools = self
.extension_manager
.read()
.await
.get_prefixed_tools(None)
.await?;
// Add frontend tools directly - they don't need prefixing since they're already uniquely named
let frontend_tools = self.frontend_tools.lock().await;
for frontend_tool in frontend_tools.values() {
tools.push(frontend_tool.tool.clone());
}
Ok(tools)
}
pub async fn add_final_output_tool(&self, response: Response) {
let mut final_output_tool = self.final_output_tool.lock().await;
let created_final_output_tool = FinalOutputTool::new(response);
@@ -728,22 +832,21 @@ impl Agent {
session: Option<SessionConfig>,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let (mut messages, issues) =
ConversationFixer::fix_conversation(Vec::from(unfixed_messages));
if !issues.is_empty() {
tracing::warn!(
"Conversation issue fixed: {}",
debug_conversation_fix(&messages, unfixed_messages, &issues)
);
}
let initial_messages = messages.clone();
let context = self
.prepare_reply_context(unfixed_messages, &session)
.await?;
let ReplyContext {
mut messages,
mut tools,
mut toolshim_tools,
mut system_prompt,
goose_mode,
initial_messages,
config,
} = context;
let reply_span = tracing::Span::current();
self.reset_retry_attempts().await;
let config = Config::global();
let (mut tools, mut toolshim_tools, mut system_prompt) =
self.prepare_tools_and_prompt().await?;
let goose_mode = Self::determine_goose_mode(session.as_ref(), config);
if let Some(content) = messages
.last()
@@ -792,8 +895,7 @@ impl Agent {
&messages,
&tools,
&toolshim_tools,
)
.await?;
).await?;
let mut added_message = false;
let mut messages_to_add = Vec::new();
@@ -836,33 +938,14 @@ impl Agent {
}
if let Some(response) = response {
let (tools_with_readonly_annotation, tools_without_annotation) =
Self::categorize_tools_by_annotation(&tools);
// Categorize tool requests
let (frontend_requests, remaining_requests, filtered_response) =
self.categorize_tool_requests(&response).await;
// Record tool calls in the router selector
let selector = self.router_tool_selector.lock().await.clone();
if let Some(selector) = selector {
for request in &frontend_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await
{
error!("Failed to record frontend tool call: {}", e);
}
}
}
for request in &remaining_requests {
if let Ok(tool_call) = &request.tool_call {
if let Err(e) = selector.record_tool_call(&tool_call.name).await
{
error!("Failed to record tool call: {}", e);
}
}
}
}
let tool_result = self.process_tool_requests(&response, &tools).await;
let ToolProcessingResult {
frontend_requests,
remaining_requests,
filtered_response,
readonly_tools,
regular_tools,
} = tool_result;
yield AgentEvent::Message(filtered_response.clone());
tokio::task::yield_now().await;
@@ -901,47 +984,17 @@ impl Agent {
check_tool_permissions(
&remaining_requests,
&mode,
tools_with_readonly_annotation.clone(),
tools_without_annotation.clone(),
readonly_tools.clone(),
regular_tools.clone(),
&mut permission_manager,
self.provider().await?,
)
.await;
).await;
let mut tool_futures: Vec<(String, ToolStream)> = Vec::new();
// Handle pre-approved and read-only tools
for request in &permission_check_result.approved {
if let Ok(tool_call) = request.tool_call.clone() {
let (req_id, tool_result) = self
.dispatch_tool_call(tool_call, request.id.clone(), cancel_token.clone())
.await;
tool_futures.push((
req_id,
match tool_result {
Ok(result) => tool_stream(
result
.notification_stream
.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => tool_stream(
Box::new(stream::empty()),
futures::future::ready(Err(e)),
),
},
));
}
}
for request in &permission_check_result.denied {
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(
request.id.clone(),
Ok(vec![Content::text(DECLINED_RESPONSE)]),
);
}
let mut tool_futures = self.handle_approved_and_denied_tools(
&permission_check_result,
message_tool_response.clone(),
cancel_token.clone()
).await?;
let tool_futures_arc = Arc::new(Mutex::new(tool_futures));
crates/goose/src/agents/subagent.rs
+2 -48
View File
@@ -91,11 +91,6 @@ impl SubAgent {
Ok(subagent)
}
/// Get the current status of the subagent
pub async fn get_status(&self) -> SubAgentStatus {
self.status.read().await.clone()
}
/// Update the status of the subagent
async fn set_status(&self, status: SubAgentStatus) {
// Update the status first, then release the lock
@@ -105,26 +100,6 @@ impl SubAgent {
} // Write lock is released here!
}
/// Get current progress information
pub async fn get_progress(&self) -> SubAgentProgress {
let status = self.get_status().await;
let turn_count = *self.turn_count.lock().await;
SubAgentProgress {
subagent_id: self.id.clone(),
status: status.clone(),
message: match &status {
SubAgentStatus::Ready => "Ready to process messages".to_string(),
SubAgentStatus::Processing => "Processing request...".to_string(),
SubAgentStatus::Completed(msg) => msg.clone(),
SubAgentStatus::Terminated => "Subagent terminated".to_string(),
},
turn: turn_count,
max_turns: self.config.max_turns,
timestamp: Utc::now(),
}
}
/// Process a message and generate a response using the subagent's provider
#[instrument(skip(self, message))]
pub async fn reply_subagent(
@@ -283,37 +258,16 @@ impl SubAgent {
}
/// Add a message to the conversation (for tracking agent responses)
pub async fn add_message(&self, message: Message) {
async fn add_message(&self, message: Message) {
let mut conversation = self.conversation.lock().await;
conversation.push(message);
}
/// Get the full conversation history
pub async fn get_conversation(&self) -> Vec<Message> {
async fn get_conversation(&self) -> Vec<Message> {
self.conversation.lock().await.clone()
}
/// Check if the subagent has completed its task
pub async fn is_completed(&self) -> bool {
matches!(
self.get_status().await,
SubAgentStatus::Completed(_) | SubAgentStatus::Terminated
)
}
/// Terminate the subagent
pub async fn terminate(&self) -> Result<(), anyhow::Error> {
debug!("Terminating subagent {}", self.id);
self.set_status(SubAgentStatus::Terminated).await;
Ok(())
}
/// Filter out subagent spawning tools to prevent infinite recursion
fn _filter_subagent_tools(tools: Vec<Tool>) -> Vec<Tool> {
// TODO: add this in subagent loop
tools
}
/// Build the system prompt for the subagent using the template
async fn build_system_prompt(&self, available_tools: &[Tool]) -> Result<String, anyhow::Error> {
let mut context = HashMap::new();