-
Notifications
You must be signed in to change notification settings - Fork 300
feat: implement delegate tool for agent task management #2890
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
use std::collections::HashMap; | ||
|
||
use eyre::{ | ||
Result, | ||
eyre, | ||
}; | ||
|
||
use crate::cli::chat::tools::delegate::agent_paths::AgentPaths; | ||
use crate::cli::chat::tools::delegate::errors::AgentError; | ||
use crate::cli::chat::tools::delegate::file_ops::{ | ||
load_json, | ||
save_json, | ||
}; | ||
use crate::cli::chat::tools::delegate::types::{ | ||
AgentConfig, | ||
AgentExecution, | ||
}; | ||
use crate::cli::chat::tools::delegate::ui::{ | ||
display_agent_info, | ||
display_default_agent_warning, | ||
get_user_confirmation, | ||
}; | ||
use crate::os::Os; | ||
|
||
const DEFAULT_AGENT: &str = "default"; | ||
|
||
pub async fn validate_agent_availability(os: &Os, agent: &str) -> Result<()> { | ||
if let Some(existing) = load_agent_execution(os, agent).await? { | ||
if existing.is_active() { | ||
return Err(eyre!("{}", AgentError::already_running(agent))); | ||
} | ||
} | ||
|
||
if agent != DEFAULT_AGENT { | ||
let agents_config = load_available_agents(os).await?; | ||
if !agents_config.contains_key(agent) { | ||
let available: Vec<String> = agents_config.keys().cloned().collect(); | ||
return Err(eyre!("{}", AgentError::not_found(agent, &available))); | ||
} | ||
} | ||
|
||
Ok(()) | ||
} | ||
|
||
pub async fn request_user_approval(os: &Os, agent: &str, task: &str) -> Result<()> { | ||
if agent != DEFAULT_AGENT { | ||
let agents_config = load_available_agents(os).await?; | ||
if let Some(agent_config) = agents_config.get(agent) { | ||
display_agent_info(agent, task, agent_config)?; | ||
if !get_user_confirmation()? { | ||
return Err(eyre!("✗ Task delegation cancelled by user.")); | ||
} | ||
} | ||
} else { | ||
display_default_agent_warning()?; | ||
} | ||
Ok(()) | ||
} | ||
|
||
pub async fn load_agent_execution(os: &Os, agent: &str) -> Result<Option<AgentExecution>> { | ||
let file_path = AgentPaths::agent_file(os, agent).await?; | ||
load_json(os, &file_path).await | ||
} | ||
|
||
pub async fn save_agent_execution(os: &Os, execution: &AgentExecution) -> Result<()> { | ||
let file_path = AgentPaths::agent_file(os, &execution.agent).await?; | ||
save_json(os, &file_path, execution).await | ||
} | ||
|
||
pub async fn load_available_agents(os: &Os) -> Result<HashMap<String, AgentConfig>> { | ||
let agents_dir = AgentPaths::cli_agents_dir(os).await?; | ||
let mut agents = HashMap::new(); | ||
|
||
if agents_dir.exists() { | ||
if let Ok(entries) = std::fs::read_dir(&agents_dir) { | ||
for entry in entries.flatten() { | ||
if let Some(file_name) = entry.file_name().to_str() { | ||
if file_name.ends_with(".json") { | ||
let agent_name = file_name.trim_end_matches(".json"); | ||
if let Ok(content) = std::fs::read_to_string(entry.path()) { | ||
if let Ok(config) = serde_json::from_str::<AgentConfig>(&content) { | ||
agents.insert(agent_name.to_string(), config); | ||
} | ||
} | ||
} | ||
} | ||
} | ||
} | ||
} | ||
|
||
Ok(agents) | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,30 @@ | ||
use std::path::PathBuf; | ||
|
||
use eyre::Result; | ||
|
||
use crate::os::Os; | ||
|
||
pub struct AgentPaths; | ||
|
||
impl AgentPaths { | ||
pub async fn subagents_dir(os: &Os) -> Result<PathBuf> { | ||
let home_dir = os.env.home().unwrap_or_default(); | ||
let agents_dir = home_dir.join(".aws").join("amazonq").join(".subagents"); | ||
|
||
if !agents_dir.exists() { | ||
std::fs::create_dir_all(&agents_dir)?; | ||
} | ||
|
||
Ok(agents_dir) | ||
} | ||
|
||
pub async fn cli_agents_dir(os: &Os) -> Result<PathBuf> { | ||
let home_dir = os.env.home().unwrap_or_default(); | ||
Ok(home_dir.join(".aws").join("amazonq").join("cli-agents")) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We should have this already defined in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ok, I can check that out. |
||
} | ||
|
||
pub async fn agent_file(os: &Os, agent: &str) -> Result<PathBuf> { | ||
let agents_dir = Self::subagents_dir(os).await?; | ||
Ok(agents_dir.join(format!("{}.json", agent))) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
pub struct AgentError; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This isn't being used as an error? We should implement There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Need to discuss this. |
||
|
||
impl AgentError { | ||
pub fn not_found(agent: &str, available: &[String]) -> String { | ||
if available.is_empty() { | ||
format!( | ||
"✗ I can't find agent '{}'. No agents are configured. You need to set up agents first.", | ||
agent | ||
) | ||
} else { | ||
format!( | ||
"✗ I can't find agent '{}'. Available agents: {}\n\nPlease use one of the available agents or set up the '{}' agent first.", | ||
agent, | ||
available.join(", "), | ||
agent | ||
) | ||
} | ||
} | ||
|
||
pub fn already_running(agent: &str) -> String { | ||
format!("Agent '{}' is already running a task", agent) | ||
} | ||
|
||
pub fn no_execution_found(agent: &str) -> String { | ||
format!("No execution found for agent '{}'", agent) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,23 @@ | ||
use std::path::Path; | ||
|
||
use eyre::Result; | ||
use serde::Serialize; | ||
use serde::de::DeserializeOwned; | ||
|
||
use crate::os::Os; | ||
|
||
pub async fn load_json<T: DeserializeOwned>(os: &Os, path: &Path) -> Result<Option<T>> { | ||
if !path.exists() { | ||
return Ok(None); | ||
} | ||
|
||
let content = os.fs.read_to_string(path).await?; | ||
let data: T = serde_json::from_str(&content)?; | ||
Ok(Some(data)) | ||
} | ||
|
||
pub async fn save_json<T: Serialize>(os: &Os, path: &Path, data: &T) -> Result<()> { | ||
let content = serde_json::to_string_pretty(data)?; | ||
os.fs.write(path, content).await?; | ||
Ok(()) | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,138 @@ | ||
mod agent_manager; | ||
mod agent_paths; | ||
mod errors; | ||
mod file_ops; | ||
mod process; | ||
mod status; | ||
mod types; | ||
mod ui; | ||
|
||
// Re-export types for external use | ||
use std::io::Write; | ||
|
||
use agent_manager::{ | ||
request_user_approval, | ||
validate_agent_availability, | ||
}; | ||
use eyre::{ | ||
Result, | ||
eyre, | ||
}; | ||
use process::{ | ||
format_launch_success, | ||
spawn_agent_process, | ||
start_monitoring, | ||
}; | ||
use serde::{ | ||
Deserialize, | ||
Serialize, | ||
}; | ||
use status::{ | ||
status_agent, | ||
status_all_agents, | ||
}; | ||
#[allow(unused_imports)] | ||
pub use types::{ | ||
AgentConfig, | ||
AgentExecution, | ||
AgentStatus, | ||
}; | ||
|
||
use crate::cli::chat::tools::{ | ||
InvokeOutput, | ||
OutputKind, | ||
}; | ||
use crate::database::settings::Setting; | ||
use crate::os::Os; | ||
|
||
const OPERATION_LAUNCH: &str = "launch"; | ||
const OPERATION_STATUS: &str = "status"; | ||
const DEFAULT_AGENT: &str = "default"; | ||
const ALL_AGENTS: &str = "all"; | ||
|
||
#[derive(Debug, Clone, Serialize, Deserialize)] | ||
pub struct Delegate { | ||
/// Operation to perform: launch or status | ||
pub operation: String, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You can define operation as an enum, and use I think something like:
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sure |
||
/// Agent name to use (optional - uses "default_agent" if not specified) | ||
#[serde(default)] | ||
pub agent: Option<String>, | ||
/// Task description (required for launch operation) | ||
#[serde(default)] | ||
pub task: Option<String>, | ||
} | ||
|
||
impl Delegate { | ||
pub async fn invoke(&self, os: &Os, _stdout: &mut impl Write) -> Result<InvokeOutput> { | ||
if !is_enabled(os) { | ||
return Ok(InvokeOutput { | ||
output: OutputKind::Text( | ||
"Delegate tool is experimental and not enabled. Use /experiment to enable it.".to_string(), | ||
), | ||
}); | ||
} | ||
|
||
let agent_name = self.get_agent_name(); | ||
let result = match self.operation.as_str() { | ||
OPERATION_LAUNCH => { | ||
let task = self | ||
.task | ||
.as_ref() | ||
.ok_or_else(|| eyre!("Task description required for launch operation"))?; | ||
launch_agent(os, agent_name, task).await? | ||
}, | ||
OPERATION_STATUS => { | ||
if agent_name == ALL_AGENTS { | ||
status_all_agents(os).await? | ||
} else { | ||
status_agent(os, agent_name).await? | ||
} | ||
}, | ||
_ => { | ||
return Err(eyre!( | ||
"Invalid operation. Use: {} or {}", | ||
OPERATION_LAUNCH, | ||
OPERATION_STATUS | ||
)); | ||
}, | ||
}; | ||
|
||
Ok(InvokeOutput { | ||
output: OutputKind::Text(result), | ||
}) | ||
} | ||
|
||
pub fn queue_description(&self, output: &mut impl Write) -> Result<()> { | ||
let agent_name = self.get_agent_name(); | ||
match self.operation.as_str() { | ||
OPERATION_LAUNCH => writeln!(output, "Launching agent '{}'", agent_name)?, | ||
OPERATION_STATUS => writeln!(output, "Checking status of agent '{}'", agent_name)?, | ||
_ => writeln!( | ||
output, | ||
"Delegate operation '{}' on agent '{}'", | ||
self.operation, agent_name | ||
)?, | ||
} | ||
Ok(()) | ||
} | ||
|
||
fn get_agent_name(&self) -> &str { | ||
match self.operation.as_str() { | ||
OPERATION_LAUNCH => self.agent.as_deref().unwrap_or(DEFAULT_AGENT), | ||
OPERATION_STATUS => self.agent.as_deref().unwrap_or(ALL_AGENTS), | ||
_ => self.agent.as_deref().unwrap_or(DEFAULT_AGENT), | ||
} | ||
} | ||
} | ||
|
||
async fn launch_agent(os: &Os, agent: &str, task: &str) -> Result<String> { | ||
validate_agent_availability(os, agent).await?; | ||
request_user_approval(os, agent, task).await?; | ||
let execution = spawn_agent_process(os, agent, task).await?; | ||
start_monitoring(execution, os.clone()).await; | ||
Ok(format_launch_success(agent, task)) | ||
} | ||
|
||
fn is_enabled(os: &Os) -> bool { | ||
os.database.settings.get_bool(Setting::EnabledDelegate).unwrap_or(false) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we add this to
src/util/directories.rs
instead for consistency with the rest of the code base? Common paradigm to keep constants like this insidedirectories.rs
currently.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can move if that is the convention. Since this is definitely experimental don't want to have others accidentally relying on this code.