From 1e4d6a4161061ca93d83409832af0192ec1c16a8 Mon Sep 17 00:00:00 2001 From: Affaan Mustafa Date: Thu, 9 Apr 2026 22:43:16 -0700 Subject: [PATCH] feat: add ecc2 agent profiles --- ecc2/src/config/mod.rs | 173 +++++++++++++++ ecc2/src/main.rs | 70 ++++-- ecc2/src/session/manager.rs | 431 ++++++++++++++++++++++++++++++++++-- ecc2/src/session/mod.rs | 2 + ecc2/src/session/store.rs | 166 +++++++++++++- ecc2/src/tui/dashboard.rs | 73 +++++- 6 files changed, 873 insertions(+), 42 deletions(-) diff --git a/ecc2/src/config/mod.rs b/ecc2/src/config/mod.rs index ffe9cdbc..938903a1 100644 --- a/ecc2/src/config/mod.rs +++ b/ecc2/src/config/mod.rs @@ -1,6 +1,7 @@ use anyhow::{Context, Result}; use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::path::PathBuf; use crate::notifications::{ @@ -48,6 +49,35 @@ pub struct ConflictResolutionConfig { pub notify_lead: bool, } +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct AgentProfileConfig { + pub inherits: Option, + pub agent: Option, + pub model: Option, + pub allowed_tools: Vec, + pub disallowed_tools: Vec, + pub permission_mode: Option, + pub add_dirs: Vec, + pub max_budget_usd: Option, + pub token_budget: Option, + pub append_system_prompt: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct ResolvedAgentProfile { + pub profile_name: String, + pub agent: Option, + pub model: Option, + pub allowed_tools: Vec, + pub disallowed_tools: Vec, + pub permission_mode: Option, + pub add_dirs: Vec, + pub max_budget_usd: Option, + pub token_budget: Option, + pub append_system_prompt: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default)] pub struct Config { @@ -61,6 +91,8 @@ pub struct Config { pub heartbeat_interval_secs: u64, pub auto_terminate_stale_sessions: bool, pub default_agent: String, + pub default_agent_profile: Option, + pub agent_profiles: BTreeMap, pub auto_dispatch_unread_handoffs: bool, pub auto_dispatch_limit_per_session: usize, pub auto_create_worktrees: bool, @@ -122,6 +154,8 @@ impl Default for Config { heartbeat_interval_secs: 30, auto_terminate_stale_sessions: false, default_agent: "claude".to_string(), + default_agent_profile: None, + agent_profiles: BTreeMap::new(), auto_dispatch_unread_handoffs: false, auto_dispatch_limit_per_session: 5, auto_create_worktrees: true, @@ -180,6 +214,41 @@ impl Config { self.budget_alert_thresholds.sanitized() } + pub fn resolve_agent_profile(&self, name: &str) -> Result { + let mut chain = Vec::new(); + self.resolve_agent_profile_inner(name, &mut chain) + } + + fn resolve_agent_profile_inner( + &self, + name: &str, + chain: &mut Vec, + ) -> Result { + if chain.iter().any(|existing| existing == name) { + chain.push(name.to_string()); + anyhow::bail!( + "agent profile inheritance cycle: {}", + chain.join(" -> ") + ); + } + + let profile = self + .agent_profiles + .get(name) + .ok_or_else(|| anyhow::anyhow!("Unknown agent profile: {name}"))?; + + chain.push(name.to_string()); + let mut resolved = if let Some(parent) = profile.inherits.as_deref() { + self.resolve_agent_profile_inner(parent, chain)? + } else { + ResolvedAgentProfile::default() + }; + chain.pop(); + + resolved.apply(name, profile); + Ok(resolved) + } + pub fn load() -> Result { let global_paths = Self::global_config_paths(); let project_paths = std::env::current_dir() @@ -437,6 +506,50 @@ impl Default for ConflictResolutionConfig { } } +impl ResolvedAgentProfile { + fn apply(&mut self, profile_name: &str, config: &AgentProfileConfig) { + self.profile_name = profile_name.to_string(); + if let Some(agent) = config.agent.as_ref() { + self.agent = Some(agent.clone()); + } + if let Some(model) = config.model.as_ref() { + self.model = Some(model.clone()); + } + merge_unique(&mut self.allowed_tools, &config.allowed_tools); + merge_unique(&mut self.disallowed_tools, &config.disallowed_tools); + if let Some(permission_mode) = config.permission_mode.as_ref() { + self.permission_mode = Some(permission_mode.clone()); + } + merge_unique(&mut self.add_dirs, &config.add_dirs); + if let Some(max_budget_usd) = config.max_budget_usd { + self.max_budget_usd = Some(max_budget_usd); + } + if let Some(token_budget) = config.token_budget { + self.token_budget = Some(token_budget); + } + self.append_system_prompt = match ( + self.append_system_prompt.take(), + config.append_system_prompt.as_ref(), + ) { + (Some(parent), Some(child)) => Some(format!("{parent}\n\n{child}")), + (Some(parent), None) => Some(parent), + (None, Some(child)) => Some(child.clone()), + (None, None) => None, + }; + } +} + +fn merge_unique(base: &mut Vec, additions: &[T]) +where + T: Clone + PartialEq, +{ + for value in additions { + if !base.contains(value) { + base.push(value.clone()); + } + } +} + impl BudgetAlertThresholds { pub fn sanitized(self) -> Self { let values = [self.advisory, self.warning, self.critical]; @@ -461,6 +574,7 @@ mod tests { PaneLayout, }; use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + use std::path::PathBuf; use uuid::Uuid; #[test] @@ -806,6 +920,65 @@ notify_lead = false ); } + #[test] + fn agent_profiles_resolve_inheritance_and_defaults() { + let config: Config = toml::from_str( + r#" +default_agent_profile = "reviewer" + +[agent_profiles.base] +model = "sonnet" +allowed_tools = ["Read"] +permission_mode = "plan" +add_dirs = ["docs"] +append_system_prompt = "Be careful." + +[agent_profiles.reviewer] +inherits = "base" +allowed_tools = ["Edit"] +disallowed_tools = ["Bash"] +token_budget = 1200 +append_system_prompt = "Review thoroughly." +"#, + ) + .unwrap(); + + let profile = config.resolve_agent_profile("reviewer").unwrap(); + assert_eq!(config.default_agent_profile.as_deref(), Some("reviewer")); + assert_eq!(profile.profile_name, "reviewer"); + assert_eq!(profile.model.as_deref(), Some("sonnet")); + assert_eq!(profile.allowed_tools, vec!["Read", "Edit"]); + assert_eq!(profile.disallowed_tools, vec!["Bash"]); + assert_eq!(profile.permission_mode.as_deref(), Some("plan")); + assert_eq!(profile.add_dirs, vec![PathBuf::from("docs")]); + assert_eq!(profile.token_budget, Some(1200)); + assert_eq!( + profile.append_system_prompt.as_deref(), + Some("Be careful.\n\nReview thoroughly.") + ); + } + + #[test] + fn agent_profile_resolution_rejects_inheritance_cycles() { + let config: Config = toml::from_str( + r#" +[agent_profiles.a] +inherits = "b" + +[agent_profiles.b] +inherits = "a" +"#, + ) + .unwrap(); + + let error = config + .resolve_agent_profile("a") + .expect_err("profile inheritance cycles must fail"); + assert!(error + .to_string() + .contains("agent profile inheritance cycle")); + } + #[test] fn completion_summary_notifications_deserialize_from_toml() { let config: Config = toml::from_str( diff --git a/ecc2/src/main.rs b/ecc2/src/main.rs index 756dadcb..ee1fcb50 100644 --- a/ecc2/src/main.rs +++ b/ecc2/src/main.rs @@ -53,6 +53,9 @@ enum Commands { /// Agent type (claude, codex, custom) #[arg(short, long, default_value = "claude")] agent: String, + /// Agent profile defined in ecc2.toml + #[arg(long)] + profile: Option, #[command(flatten)] worktree: WorktreePolicyArgs, /// Source session to delegate from @@ -69,6 +72,9 @@ enum Commands { /// Agent type (claude, codex, custom) #[arg(short, long, default_value = "claude")] agent: String, + /// Agent profile defined in ecc2.toml + #[arg(long)] + profile: Option, #[command(flatten)] worktree: WorktreePolicyArgs, }, @@ -82,6 +88,9 @@ enum Commands { /// Agent type (claude, codex, custom) #[arg(short, long, default_value = "claude")] agent: String, + /// Agent profile defined in ecc2.toml + #[arg(long)] + profile: Option, #[command(flatten)] worktree: WorktreePolicyArgs, }, @@ -381,6 +390,7 @@ async fn main() -> Result<()> { Some(Commands::Start { task, agent, + profile, worktree, from_session, }) => { @@ -394,18 +404,34 @@ async fn main() -> Result<()> { } else { None }; - let session_id = session::manager::create_session_with_grouping( - &db, - &cfg, - &task, - &agent, - use_worktree, - session::SessionGrouping { - project: source.as_ref().map(|session| session.project.clone()), - task_group: source.as_ref().map(|session| session.task_group.clone()), - }, - ) - .await?; + let grouping = session::SessionGrouping { + project: source.as_ref().map(|session| session.project.clone()), + task_group: source.as_ref().map(|session| session.task_group.clone()), + }; + let session_id = if let Some(source) = source.as_ref() { + session::manager::create_session_from_source_with_profile_and_grouping( + &db, + &cfg, + &task, + &agent, + use_worktree, + profile.as_deref(), + &source.id, + grouping, + ) + .await? + } else { + session::manager::create_session_with_profile_and_grouping( + &db, + &cfg, + &task, + &agent, + use_worktree, + profile.as_deref(), + grouping, + ) + .await? + }; if let Some(source) = source { let from_id = source.id; send_handoff_message(&db, &from_id, &session_id)?; @@ -416,6 +442,7 @@ async fn main() -> Result<()> { from_session, task, agent, + profile, worktree, }) => { let use_worktree = worktree.resolve(&cfg); @@ -431,12 +458,14 @@ async fn main() -> Result<()> { ) }); - let session_id = session::manager::create_session_with_grouping( + let session_id = session::manager::create_session_from_source_with_profile_and_grouping( &db, &cfg, &task, &agent, use_worktree, + profile.as_deref(), + &source.id, session::SessionGrouping { project: Some(source.project.clone()), task_group: Some(source.task_group.clone()), @@ -454,13 +483,22 @@ async fn main() -> Result<()> { from_session, task, agent, + profile, worktree, }) => { let use_worktree = worktree.resolve(&cfg); let lead_id = resolve_session_id(&db, &from_session)?; - let outcome = - session::manager::assign_session(&db, &cfg, &lead_id, &task, &agent, use_worktree) - .await?; + let outcome = session::manager::assign_session_with_profile_and_grouping( + &db, + &cfg, + &lead_id, + &task, + &agent, + use_worktree, + profile.as_deref(), + session::SessionGrouping::default(), + ) + .await?; if session::manager::assignment_action_routes_work(outcome.action) { println!( "Assignment routed: {} -> {} ({})", diff --git a/ecc2/src/session/manager.rs b/ecc2/src/session/manager.rs index 8311a4f2..28669e9d 100644 --- a/ecc2/src/session/manager.rs +++ b/ecc2/src/session/manager.rs @@ -11,7 +11,7 @@ use super::runtime::capture_command_output; use super::store::StateStore; use super::{ default_project_label, default_task_group_label, normalize_group_label, Session, - SessionGrouping, SessionMetrics, SessionState, + SessionAgentProfile, SessionGrouping, SessionMetrics, SessionState, }; use crate::comms::{self, MessageType}; use crate::config::Config; @@ -25,12 +25,13 @@ pub async fn create_session( agent_type: &str, use_worktree: bool, ) -> Result { - create_session_with_grouping( + create_session_with_profile_and_grouping( db, cfg, task, agent_type, use_worktree, + None, SessionGrouping::default(), ) .await @@ -43,6 +44,27 @@ pub async fn create_session_with_grouping( agent_type: &str, use_worktree: bool, grouping: SessionGrouping, +) -> Result { + create_session_with_profile_and_grouping( + db, + cfg, + task, + agent_type, + use_worktree, + None, + grouping, + ) + .await +} + +pub async fn create_session_with_profile_and_grouping( + db: &StateStore, + cfg: &Config, + task: &str, + agent_type: &str, + use_worktree: bool, + profile_name: Option<&str>, + grouping: SessionGrouping, ) -> Result { let repo_root = std::env::current_dir().context("Failed to resolve current working directory")?; @@ -53,6 +75,34 @@ pub async fn create_session_with_grouping( agent_type, use_worktree, &repo_root, + profile_name, + None, + grouping, + ) + .await +} + +pub async fn create_session_from_source_with_profile_and_grouping( + db: &StateStore, + cfg: &Config, + task: &str, + agent_type: &str, + use_worktree: bool, + profile_name: Option<&str>, + source_session_id: &str, + grouping: SessionGrouping, +) -> Result { + let repo_root = + std::env::current_dir().context("Failed to resolve current working directory")?; + queue_session_in_dir( + db, + cfg, + task, + agent_type, + use_worktree, + &repo_root, + profile_name, + Some(source_session_id), grouping, ) .await @@ -66,6 +116,7 @@ pub fn get_status(db: &StateStore, id: &str) -> Result { let session = resolve_session(db, id)?; let session_id = session.id.clone(); Ok(SessionStatus { + profile: db.get_session_profile(&session_id)?, session, parent_session: db.latest_task_handoff_source(&session_id)?, delegated_children: db.delegated_children(&session_id, 5)?, @@ -159,13 +210,14 @@ pub async fn assign_session( agent_type: &str, use_worktree: bool, ) -> Result { - assign_session_with_grouping( + assign_session_with_profile_and_grouping( db, cfg, lead_id, task, agent_type, use_worktree, + None, SessionGrouping::default(), ) .await @@ -179,6 +231,29 @@ pub async fn assign_session_with_grouping( agent_type: &str, use_worktree: bool, grouping: SessionGrouping, +) -> Result { + assign_session_with_profile_and_grouping( + db, + cfg, + lead_id, + task, + agent_type, + use_worktree, + None, + grouping, + ) + .await +} + +pub async fn assign_session_with_profile_and_grouping( + db: &StateStore, + cfg: &Config, + lead_id: &str, + task: &str, + agent_type: &str, + use_worktree: bool, + profile_name: Option<&str>, + grouping: SessionGrouping, ) -> Result { let repo_root = std::env::current_dir().context("Failed to resolve current working directory")?; @@ -191,6 +266,7 @@ pub async fn assign_session_with_grouping( use_worktree, &repo_root, &std::env::current_exe().context("Failed to resolve ECC executable path")?, + profile_name, grouping, ) .await @@ -228,6 +304,7 @@ pub async fn drain_inbox( use_worktree, &repo_root, &runner_program, + None, SessionGrouping::default(), ) .await?; @@ -434,6 +511,7 @@ pub async fn rebalance_team_backlog( use_worktree, &repo_root, &runner_program, + None, SessionGrouping::default(), ) .await?; @@ -464,12 +542,15 @@ pub async fn stop_session(db: &StateStore, id: &str) -> Result<()> { pub struct BudgetEnforcementOutcome { pub token_budget_exceeded: bool, pub cost_budget_exceeded: bool, + pub profile_token_budget_exceeded: bool, pub paused_sessions: Vec, } impl BudgetEnforcementOutcome { pub fn hard_limit_exceeded(&self) -> bool { - self.token_budget_exceeded || self.cost_budget_exceeded + self.token_budget_exceeded + || self.cost_budget_exceeded + || self.profile_token_budget_exceeded } } @@ -490,18 +571,51 @@ pub fn enforce_budget_hard_limits( let mut outcome = BudgetEnforcementOutcome { token_budget_exceeded: cfg.token_budget > 0 && total_tokens >= cfg.token_budget, cost_budget_exceeded: cfg.cost_budget_usd > 0.0 && total_cost >= cfg.cost_budget_usd, + profile_token_budget_exceeded: false, paused_sessions: Vec::new(), }; + let mut sessions_to_pause = HashSet::new(); + + if outcome.token_budget_exceeded || outcome.cost_budget_exceeded { + for session in sessions.iter().filter(|session| { + matches!( + session.state, + SessionState::Pending | SessionState::Running | SessionState::Idle + ) + }) { + sessions_to_pause.insert(session.id.clone()); + } + } + + for session in sessions.iter().filter(|session| { + matches!( + session.state, + SessionState::Pending | SessionState::Running | SessionState::Idle + ) + }) { + let Some(profile) = db.get_session_profile(&session.id)? else { + continue; + }; + let Some(token_budget) = profile.token_budget else { + continue; + }; + if token_budget > 0 && session.metrics.tokens_used >= token_budget { + outcome.profile_token_budget_exceeded = true; + sessions_to_pause.insert(session.id.clone()); + } + } + if !outcome.hard_limit_exceeded() { return Ok(outcome); } for session in sessions.into_iter().filter(|session| { - matches!( - session.state, - SessionState::Pending | SessionState::Running | SessionState::Idle - ) + sessions_to_pause.contains(&session.id) + && matches!( + session.state, + SessionState::Pending | SessionState::Running | SessionState::Idle + ) }) { stop_session_recorded(db, &session, false)?; outcome.paused_sessions.push(session.id); @@ -820,6 +934,7 @@ async fn assign_session_in_dir_with_runner_program( use_worktree: bool, repo_root: &Path, runner_program: &Path, + profile_name: Option<&str>, grouping: SessionGrouping, ) -> Result { let lead = resolve_session(db, lead_id)?; @@ -868,6 +983,8 @@ async fn assign_session_in_dir_with_runner_program( use_worktree, repo_root, runner_program, + profile_name, + Some(&lead.id), inherited_grouping.clone(), ) .await?; @@ -943,6 +1060,8 @@ async fn assign_session_in_dir_with_runner_program( use_worktree, repo_root, runner_program, + profile_name, + Some(&lead.id), inherited_grouping, ) .await?; @@ -1623,7 +1742,8 @@ pub async fn run_session( } let agent_program = agent_program(agent_type)?; - let command = build_agent_command(&agent_program, task, session_id, working_dir); + let profile = db.get_session_profile(session_id)?; + let command = build_agent_command(&agent_program, task, session_id, working_dir, profile.as_ref()); capture_command_output( cfg.db_path.clone(), session_id.to_string(), @@ -1750,6 +1870,8 @@ async fn queue_session_in_dir( agent_type: &str, use_worktree: bool, repo_root: &Path, + profile_name: Option<&str>, + inherited_profile_session_id: Option<&str>, grouping: SessionGrouping, ) -> Result { queue_session_in_dir_with_runner_program( @@ -1760,6 +1882,8 @@ async fn queue_session_in_dir( use_worktree, repo_root, &std::env::current_exe().context("Failed to resolve ECC executable path")?, + profile_name, + inherited_profile_session_id, grouping, ) .await @@ -1773,11 +1897,29 @@ async fn queue_session_in_dir_with_runner_program( use_worktree: bool, repo_root: &Path, runner_program: &Path, + profile_name: Option<&str>, + inherited_profile_session_id: Option<&str>, grouping: SessionGrouping, ) -> Result { - let session = - build_session_record(db, task, agent_type, use_worktree, cfg, repo_root, grouping)?; + let profile = + resolve_launch_profile(db, cfg, profile_name, inherited_profile_session_id)?; + let effective_agent_type = profile + .as_ref() + .and_then(|profile| profile.agent.as_deref()) + .unwrap_or(agent_type); + let session = build_session_record( + db, + task, + effective_agent_type, + use_worktree, + cfg, + repo_root, + grouping, + )?; db.insert_session(&session)?; + if let Some(profile) = profile.as_ref() { + db.upsert_session_profile(&session.id, profile)?; + } if use_worktree && session.worktree.is_none() { db.enqueue_pending_worktree(&session.id, repo_root)?; @@ -1793,7 +1935,7 @@ async fn queue_session_in_dir_with_runner_program( match spawn_session_runner_for_program( task, &session.id, - agent_type, + &session.agent_type, working_dir, runner_program, ) @@ -1911,6 +2053,27 @@ async fn create_session_in_dir( } } +fn resolve_launch_profile( + db: &StateStore, + cfg: &Config, + explicit_profile_name: Option<&str>, + inherited_profile_session_id: Option<&str>, +) -> Result> { + let inherited_profile_name = match inherited_profile_session_id { + Some(session_id) => db.get_session_profile(session_id)?.map(|profile| profile.profile_name), + None => None, + }; + let profile_name = explicit_profile_name + .map(ToOwned::to_owned) + .or(inherited_profile_name) + .or_else(|| cfg.default_agent_profile.clone()); + + profile_name + .as_deref() + .map(|name| cfg.resolve_agent_profile(name)) + .transpose() +} + fn attached_worktree_count(db: &StateStore) -> Result { Ok(db .list_sessions()? @@ -2075,16 +2238,44 @@ fn build_agent_command( task: &str, session_id: &str, working_dir: &Path, + profile: Option<&SessionAgentProfile>, ) -> Command { let mut command = Command::new(agent_program); + command.env("ECC_SESSION_ID", session_id); command - .env("ECC_SESSION_ID", session_id) .arg("--print") .arg("--name") - .arg(format!("ecc-{session_id}")) - .arg(task) - .current_dir(working_dir) - .stdin(Stdio::null()); + .arg(format!("ecc-{session_id}")); + if let Some(profile) = profile { + if let Some(model) = profile.model.as_ref() { + command.arg("--model").arg(model); + } + if !profile.allowed_tools.is_empty() { + command + .arg("--allowed-tools") + .arg(profile.allowed_tools.join(",")); + } + if !profile.disallowed_tools.is_empty() { + command + .arg("--disallowed-tools") + .arg(profile.disallowed_tools.join(",")); + } + if let Some(permission_mode) = profile.permission_mode.as_ref() { + command.arg("--permission-mode").arg(permission_mode); + } + for dir in &profile.add_dirs { + command.arg("--add-dir").arg(dir); + } + if let Some(max_budget_usd) = profile.max_budget_usd { + command + .arg("--max-budget-usd") + .arg(max_budget_usd.to_string()); + } + if let Some(prompt) = profile.append_system_prompt.as_ref() { + command.arg("--append-system-prompt").arg(prompt); + } + } + command.arg(task).current_dir(working_dir).stdin(Stdio::null()); command } @@ -2094,7 +2285,7 @@ async fn spawn_claude_code( session_id: &str, working_dir: &Path, ) -> Result { - let mut command = build_agent_command(agent_program, task, session_id, working_dir); + let mut command = build_agent_command(agent_program, task, session_id, working_dir, None); let child = command .stdout(Stdio::null()) .stderr(Stdio::null()) @@ -2194,6 +2385,7 @@ async fn kill_process(pid: u32) -> Result<()> { } pub struct SessionStatus { + profile: Option, session: Session, parent_session: Option, delegated_children: Vec, @@ -2363,6 +2555,21 @@ impl fmt::Display for SessionStatus { writeln!(f, "Task: {}", s.task)?; writeln!(f, "Agent: {}", s.agent_type)?; writeln!(f, "State: {}", s.state)?; + if let Some(profile) = self.profile.as_ref() { + writeln!(f, "Profile: {}", profile.profile_name)?; + if let Some(model) = profile.model.as_ref() { + writeln!(f, "Model: {}", model)?; + } + if let Some(permission_mode) = profile.permission_mode.as_ref() { + writeln!(f, "Perms: {}", permission_mode)?; + } + if let Some(token_budget) = profile.token_budget { + writeln!(f, "Profile tokens: {}", token_budget)?; + } + if let Some(max_budget_usd) = profile.max_budget_usd { + writeln!(f, "Profile cost: ${max_budget_usd:.4}")?; + } + } if let Some(parent) = self.parent_session.as_ref() { writeln!(f, "Parent: {}", parent)?; } @@ -2590,7 +2797,7 @@ fn session_state_label(state: &SessionState) -> &'static str { mod tests { use super::*; use crate::config::{Config, PaneLayout, Theme}; - use crate::session::{Session, SessionMetrics, SessionState}; + use crate::session::{Session, SessionAgentProfile, SessionMetrics, SessionState}; use anyhow::{Context, Result}; use chrono::{Duration, Utc}; use std::fs; @@ -2635,6 +2842,8 @@ mod tests { heartbeat_interval_secs: 5, auto_terminate_stale_sessions: false, default_agent: "claude".to_string(), + default_agent_profile: None, + agent_profiles: Default::default(), auto_dispatch_unread_handoffs: false, auto_dispatch_limit_per_session: 5, auto_create_worktrees: true, @@ -2674,6 +2883,61 @@ mod tests { } } + #[test] + fn build_agent_command_applies_profile_runner_flags() { + let profile = SessionAgentProfile { + profile_name: "reviewer".to_string(), + agent: None, + model: Some("sonnet".to_string()), + allowed_tools: vec!["Read".to_string(), "Edit".to_string()], + disallowed_tools: vec!["Bash".to_string()], + permission_mode: Some("plan".to_string()), + add_dirs: vec![PathBuf::from("docs"), PathBuf::from("specs")], + max_budget_usd: Some(1.25), + token_budget: Some(750), + append_system_prompt: Some("Review thoroughly.".to_string()), + }; + + let command = build_agent_command( + Path::new("claude"), + "review this change", + "sess-1234", + Path::new("/tmp/repo"), + Some(&profile), + ); + let args = command + .as_std() + .get_args() + .map(|value| value.to_string_lossy().to_string()) + .collect::>(); + + assert_eq!( + args, + vec![ + "--print", + "--name", + "ecc-sess-1234", + "--model", + "sonnet", + "--allowed-tools", + "Read,Edit", + "--disallowed-tools", + "Bash", + "--permission-mode", + "plan", + "--add-dir", + "docs", + "--add-dir", + "specs", + "--max-budget-usd", + "1.25", + "--append-system-prompt", + "Review thoroughly.", + "review this change", + ] + ); + } + #[test] fn enforce_session_heartbeats_marks_overdue_running_sessions_stale() -> Result<()> { let tempdir = TestDir::new("manager-heartbeat-stale")?; @@ -3099,6 +3363,62 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "current_thread")] + async fn create_session_uses_default_agent_profile_and_persists_launch_settings() -> Result<()> { + let tempdir = TestDir::new("manager-default-agent-profile")?; + let repo_root = tempdir.path().join("repo"); + init_git_repo(&repo_root)?; + + let mut cfg = build_config(tempdir.path()); + cfg.default_agent_profile = Some("reviewer".to_string()); + cfg.agent_profiles.insert( + "reviewer".to_string(), + crate::config::AgentProfileConfig { + model: Some("sonnet".to_string()), + allowed_tools: vec!["Read".to_string(), "Edit".to_string()], + disallowed_tools: vec!["Bash".to_string()], + permission_mode: Some("plan".to_string()), + add_dirs: vec![PathBuf::from("docs")], + token_budget: Some(800), + append_system_prompt: Some("Review thoroughly.".to_string()), + ..Default::default() + }, + ); + let db = StateStore::open(&cfg.db_path)?; + let (fake_runner, _) = write_fake_claude(tempdir.path())?; + + let session_id = queue_session_in_dir_with_runner_program( + &db, + &cfg, + "review work", + "claude", + false, + &repo_root, + &fake_runner, + None, + None, + SessionGrouping::default(), + ) + .await?; + + let profile = db + .get_session_profile(&session_id)? + .context("session profile should be persisted")?; + assert_eq!(profile.profile_name, "reviewer"); + assert_eq!(profile.model.as_deref(), Some("sonnet")); + assert_eq!(profile.allowed_tools, vec!["Read", "Edit"]); + assert_eq!(profile.disallowed_tools, vec!["Bash"]); + assert_eq!(profile.permission_mode.as_deref(), Some("plan")); + assert_eq!(profile.add_dirs, vec![PathBuf::from("docs")]); + assert_eq!(profile.token_budget, Some(800)); + assert_eq!( + profile.append_system_prompt.as_deref(), + Some("Review thoroughly.") + ); + + Ok(()) + } + #[test] fn enforce_budget_hard_limits_stops_active_sessions_without_cleaning_worktrees() -> Result<()> { let tempdir = TestDir::new("manager-budget-pause")?; @@ -3214,6 +3534,73 @@ mod tests { Ok(()) } + #[test] + fn enforce_budget_hard_limits_pauses_sessions_over_profile_token_budget() -> Result<()> { + let tempdir = TestDir::new("manager-profile-token-budget")?; + let cfg = build_config(tempdir.path()); + let db = StateStore::open(&cfg.db_path)?; + let now = Utc::now(); + + db.insert_session(&Session { + id: "profile-over-budget".to_string(), + task: "review work".to_string(), + project: "workspace".to_string(), + task_group: "general".to_string(), + agent_type: "claude".to_string(), + working_dir: tempdir.path().to_path_buf(), + state: SessionState::Running, + pid: Some(999_998), + worktree: None, + created_at: now - Duration::minutes(1), + updated_at: now, + last_heartbeat_at: now, + metrics: SessionMetrics::default(), + })?; + db.upsert_session_profile( + "profile-over-budget", + &SessionAgentProfile { + profile_name: "reviewer".to_string(), + agent: None, + model: Some("sonnet".to_string()), + allowed_tools: vec!["Read".to_string()], + disallowed_tools: Vec::new(), + permission_mode: Some("plan".to_string()), + add_dirs: Vec::new(), + max_budget_usd: None, + token_budget: Some(75), + append_system_prompt: None, + }, + )?; + db.update_metrics( + "profile-over-budget", + &SessionMetrics { + input_tokens: 60, + output_tokens: 30, + tokens_used: 90, + tool_calls: 0, + files_changed: 0, + duration_secs: 60, + cost_usd: 0.0, + }, + )?; + + let outcome = enforce_budget_hard_limits(&db, &cfg)?; + assert!(!outcome.token_budget_exceeded); + assert!(!outcome.cost_budget_exceeded); + assert!(outcome.profile_token_budget_exceeded); + assert_eq!( + outcome.paused_sessions, + vec!["profile-over-budget".to_string()] + ); + + let session = db + .get_session("profile-over-budget")? + .context("session should still exist")?; + assert_eq!(session.state, SessionState::Stopped); + + Ok(()) + } + #[tokio::test(flavor = "current_thread")] async fn resume_session_requeues_failed_session() -> Result<()> { let tempdir = TestDir::new("manager-resume-session")?; @@ -4108,6 +4495,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; @@ -4181,6 +4569,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; @@ -4266,6 +4655,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; @@ -4338,6 +4728,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; @@ -4394,6 +4785,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; @@ -4467,6 +4859,7 @@ mod tests { true, &repo_root, &fake_runner, + None, SessionGrouping::default(), ) .await?; diff --git a/ecc2/src/session/mod.rs b/ecc2/src/session/mod.rs index b66f6ee0..301f3384 100644 --- a/ecc2/src/session/mod.rs +++ b/ecc2/src/session/mod.rs @@ -10,6 +10,8 @@ use std::fmt; use std::path::Path; use std::path::PathBuf; +pub type SessionAgentProfile = crate::config::ResolvedAgentProfile; + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Session { pub id: String, diff --git a/ecc2/src/session/store.rs b/ecc2/src/session/store.rs index b7029b57..d0b82686 100644 --- a/ecc2/src/session/store.rs +++ b/ecc2/src/session/store.rs @@ -14,8 +14,8 @@ use crate::observability::{ToolCallEvent, ToolLogEntry, ToolLogPage}; use super::output::{OutputLine, OutputStream, OUTPUT_BUFFER_LIMIT}; use super::{ default_project_label, default_task_group_label, normalize_group_label, DecisionLogEntry, - FileActivityAction, FileActivityEntry, Session, SessionMessage, SessionMetrics, SessionState, - WorktreeInfo, + FileActivityAction, FileActivityEntry, Session, SessionAgentProfile, SessionMessage, + SessionMetrics, SessionState, WorktreeInfo, }; pub struct StateStore { @@ -194,6 +194,19 @@ impl StateStore { file_events_json TEXT NOT NULL DEFAULT '[]' ); + CREATE TABLE IF NOT EXISTS session_profiles ( + session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE, + profile_name TEXT NOT NULL, + model TEXT, + allowed_tools_json TEXT NOT NULL DEFAULT '[]', + disallowed_tools_json TEXT NOT NULL DEFAULT '[]', + permission_mode TEXT, + add_dirs_json TEXT NOT NULL DEFAULT '[]', + max_budget_usd REAL, + token_budget INTEGER, + append_system_prompt TEXT + ); + CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, from_session TEXT NOT NULL, @@ -569,6 +582,98 @@ impl StateStore { Ok(()) } + pub fn upsert_session_profile( + &self, + session_id: &str, + profile: &SessionAgentProfile, + ) -> Result<()> { + let allowed_tools_json = serde_json::to_string(&profile.allowed_tools) + .context("serialize allowed agent profile tools")?; + let disallowed_tools_json = serde_json::to_string(&profile.disallowed_tools) + .context("serialize disallowed agent profile tools")?; + let add_dirs_json = serde_json::to_string(&profile.add_dirs) + .context("serialize agent profile add_dirs")?; + + self.conn.execute( + "INSERT INTO session_profiles ( + session_id, + profile_name, + model, + allowed_tools_json, + disallowed_tools_json, + permission_mode, + add_dirs_json, + max_budget_usd, + token_budget, + append_system_prompt + ) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10) + ON CONFLICT(session_id) DO UPDATE SET + profile_name = excluded.profile_name, + model = excluded.model, + allowed_tools_json = excluded.allowed_tools_json, + disallowed_tools_json = excluded.disallowed_tools_json, + permission_mode = excluded.permission_mode, + add_dirs_json = excluded.add_dirs_json, + max_budget_usd = excluded.max_budget_usd, + token_budget = excluded.token_budget, + append_system_prompt = excluded.append_system_prompt", + rusqlite::params![ + session_id, + profile.profile_name, + profile.model, + allowed_tools_json, + disallowed_tools_json, + profile.permission_mode, + add_dirs_json, + profile.max_budget_usd, + profile.token_budget, + profile.append_system_prompt, + ], + )?; + Ok(()) + } + + pub fn get_session_profile(&self, session_id: &str) -> Result> { + self.conn + .query_row( + "SELECT + profile_name, + model, + allowed_tools_json, + disallowed_tools_json, + permission_mode, + add_dirs_json, + max_budget_usd, + token_budget, + append_system_prompt + FROM session_profiles + WHERE session_id = ?1", + [session_id], + |row| { + let allowed_tools_json: String = row.get(2)?; + let disallowed_tools_json: String = row.get(3)?; + let add_dirs_json: String = row.get(5)?; + Ok(SessionAgentProfile { + profile_name: row.get(0)?, + model: row.get(1)?, + allowed_tools: serde_json::from_str(&allowed_tools_json) + .unwrap_or_default(), + disallowed_tools: serde_json::from_str(&disallowed_tools_json) + .unwrap_or_default(), + permission_mode: row.get(4)?, + add_dirs: serde_json::from_str(&add_dirs_json).unwrap_or_default(), + max_budget_usd: row.get(6)?, + token_budget: row.get(7)?, + append_system_prompt: row.get(8)?, + agent: None, + }) + }, + ) + .optional() + .map_err(Into::into) + } + pub fn update_state_and_pid( &self, session_id: &str, @@ -2532,6 +2637,63 @@ mod tests { Ok(()) } + #[test] + fn session_profile_round_trips_with_launch_settings() -> Result<()> { + let tempdir = TestDir::new("store-session-profile")?; + let db = StateStore::open(&tempdir.path().join("state.db"))?; + let now = Utc::now(); + + db.insert_session(&Session { + id: "session-1".to_string(), + task: "review work".to_string(), + project: "workspace".to_string(), + task_group: "general".to_string(), + agent_type: "claude".to_string(), + working_dir: PathBuf::from("/tmp"), + state: SessionState::Pending, + pid: None, + worktree: None, + created_at: now, + updated_at: now, + last_heartbeat_at: now, + metrics: SessionMetrics::default(), + })?; + + db.upsert_session_profile( + "session-1", + &crate::session::SessionAgentProfile { + agent: None, + profile_name: "reviewer".to_string(), + model: Some("sonnet".to_string()), + allowed_tools: vec!["Read".to_string(), "Edit".to_string()], + disallowed_tools: vec!["Bash".to_string()], + permission_mode: Some("plan".to_string()), + add_dirs: vec![PathBuf::from("docs"), PathBuf::from("specs")], + max_budget_usd: Some(1.5), + token_budget: Some(1200), + append_system_prompt: Some("Review thoroughly.".to_string()), + }, + )?; + + let profile = db + .get_session_profile("session-1")? + .expect("profile should be stored"); + assert_eq!(profile.profile_name, "reviewer"); + assert_eq!(profile.model.as_deref(), Some("sonnet")); + assert_eq!(profile.allowed_tools, vec!["Read", "Edit"]); + assert_eq!(profile.disallowed_tools, vec!["Bash"]); + assert_eq!(profile.permission_mode.as_deref(), Some("plan")); + assert_eq!(profile.add_dirs, vec![PathBuf::from("docs"), PathBuf::from("specs")]); + assert_eq!(profile.max_budget_usd, Some(1.5)); + assert_eq!(profile.token_budget, Some(1200)); + assert_eq!( + profile.append_system_prompt.as_deref(), + Some("Review thoroughly.") + ); + + Ok(()) + } + #[test] fn sync_cost_tracker_metrics_aggregates_usage_into_sessions() -> Result<()> { let tempdir = TestDir::new("store-cost-metrics")?; diff --git a/ecc2/src/tui/dashboard.rs b/ecc2/src/tui/dashboard.rs index 68603fb9..a471d637 100644 --- a/ecc2/src/tui/dashboard.rs +++ b/ecc2/src/tui/dashboard.rs @@ -5392,6 +5392,11 @@ impl Dashboard { fn selected_session_metrics_text(&self) -> String { if let Some(session) = self.sessions.get(self.selected_session) { let metrics = &session.metrics; + let selected_profile = self + .db + .get_session_profile(&session.id) + .ok() + .flatten(); let group_peers = self .sessions .iter() @@ -5413,6 +5418,57 @@ impl Dashboard { ), ]; + if let Some(profile) = selected_profile.as_ref() { + let model = profile.model.as_deref().unwrap_or("default"); + let permission_mode = profile.permission_mode.as_deref().unwrap_or("default"); + lines.push(format!( + "Profile {} | Model {} | Permissions {}", + profile.profile_name, model, permission_mode + )); + let mut profile_details = Vec::new(); + if let Some(token_budget) = profile.token_budget { + profile_details.push(format!( + "Profile tokens {}", + format_token_count(token_budget) + )); + } + if let Some(max_budget_usd) = profile.max_budget_usd { + profile_details.push(format!( + "Profile cost {}", + format_currency(max_budget_usd) + )); + } + if !profile.allowed_tools.is_empty() { + profile_details.push(format!( + "Allow {}", + truncate_for_dashboard(&profile.allowed_tools.join(", "), 36) + )); + } + if !profile.disallowed_tools.is_empty() { + profile_details.push(format!( + "Deny {}", + truncate_for_dashboard(&profile.disallowed_tools.join(", "), 36) + )); + } + if !profile.add_dirs.is_empty() { + profile_details.push(format!( + "Dirs {}", + truncate_for_dashboard( + &profile + .add_dirs + .iter() + .map(|path| path.display().to_string()) + .collect::>() + .join(", "), + 36 + ) + )); + } + if !profile_details.is_empty() { + lines.push(profile_details.join(" | ")); + } + } + if let Some(parent) = self.selected_parent_session.as_ref() { lines.push(format!("Delegated from {}", format_session_id(parent))); } @@ -7878,11 +7934,16 @@ fn heartbeat_enforcement_note(outcome: &manager::HeartbeatEnforcementOutcome) -> } fn budget_auto_pause_note(outcome: &manager::BudgetEnforcementOutcome) -> String { - let cause = match (outcome.token_budget_exceeded, outcome.cost_budget_exceeded) { - (true, true) => "token and cost budgets exceeded", - (true, false) => "token budget exceeded", - (false, true) => "cost budget exceeded", - (false, false) => "budget exceeded", + let cause = match ( + outcome.token_budget_exceeded, + outcome.cost_budget_exceeded, + outcome.profile_token_budget_exceeded, + ) { + (true, true, _) => "token and cost budgets exceeded", + (true, false, _) => "token budget exceeded", + (false, true, _) => "cost budget exceeded", + (false, false, true) => "profile token budget exceeded", + (false, false, false) => "budget exceeded", }; format!( @@ -13011,6 +13072,8 @@ diff --git a/src/lib.rs b/src/lib.rs heartbeat_interval_secs: 5, auto_terminate_stale_sessions: false, default_agent: "claude".to_string(), + default_agent_profile: None, + agent_profiles: Default::default(), auto_dispatch_unread_handoffs: false, auto_dispatch_limit_per_session: 5, auto_create_worktrees: true,