Skip to content

chore: small refactor on agent.rs #3703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lifeizhou-ap merged 5 commits into from
Jul 29, 2025
Merged

chore: small refactor on agent.rs #3703

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
85fc5f1
removed unused functions
lifeizhou-ap Jul 29, 2025
a846b1f
removed unused functions in agent
lifeizhou-ap Jul 29, 2025
374f879
refactored agent (extract some functions)
lifeizhou-ap Jul 29, 2025
e5d4173
cleanup
lifeizhou-ap Jul 29, 2025
fe93ca5
rename with a better function name
lifeizhou-ap Jul 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
279 changes: 166 additions & 113 deletions crates/goose/src/agents/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -34,8 +34,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;
Expand All @@ -60,6 +60,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>>>,
Expand Down Expand Up @@ -161,17 +181,6 @@ impl Agent {
*tool_monitor = Some(ToolMonitor::new(max_repetitions));
}

pub async fn get_tool_stats(&self) -> Option<HashMap<String, u32>> {

@lifeizhou-ap lifeizhou-ap Jul 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these functions are not used

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;
Expand Down Expand Up @@ -207,6 +216,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(

@DOsinga DOsinga Jul 29, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we're doing two things here, maybe would be better to just do record_tool_calls and move that to router_tool_selector to record more than one call and then we can just call that directly from agent with frontend_requests.iter().chain(&remaining_requests)

@lifeizhou-ap lifeizhou-ap Jul 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review! Yes, the function has a couple of responsibilities. I feel the tool handling in reply could be cleaner with more refactoring and will have a followup pr to refactor more in this area.

&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;
Expand All @@ -231,24 +353,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);
Expand Down Expand Up @@ -714,22 +818,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()
Expand Down Expand Up @@ -778,8 +881,7 @@ impl Agent {
&messages,
&tools,
&toolshim_tools,
)
.await?;
).await?;

let mut added_message = false;
let mut messages_to_add = Vec::new();
Expand Down Expand Up @@ -822,33 +924,14 @@ impl Agent {
}

if let Some(response) = response {
let (tools_with_readonly_annotation, tools_without_annotation) =

@lifeizhou-ap lifeizhou-ap Jul 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extracted to process_tool_requests function

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;
Expand Down Expand Up @@ -885,47 +968,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;

let mut tool_futures: Vec<(String, ToolStream)> = Vec::new();

@lifeizhou-ap lifeizhou-ap Jul 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extracted to execute_approved_tools function


// 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)),
),
},
));
}
}
).await;

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

Expand Down
Loading
Loading