diff --git a/server/src/http/mcp/agent_tools.rs b/server/src/http/mcp/agent_tools.rs index 8b0c2979..22ca0806 100644 --- a/server/src/http/mcp/agent_tools.rs +++ b/server/src/http/mcp/agent_tools.rs @@ -230,6 +230,92 @@ pub(super) fn tool_get_agent_config(ctx: &AppContext) -> Result .map_err(|e| format!("Serialization error: {e}")) } +/// Get remaining turns and budget for a running agent. +/// +/// Returns turns used, max turns, remaining turns, budget used, max budget, +/// and remaining budget for the named agent. Fails if the agent is not +/// currently running or pending. +pub(super) fn tool_get_agent_remaining_turns_and_budget( + args: &Value, + ctx: &AppContext, +) -> Result { + let story_id = args + .get("story_id") + .and_then(|v| v.as_str()) + .ok_or("Missing required argument: story_id")?; + let agent_name = args + .get("agent_name") + .and_then(|v| v.as_str()) + .ok_or("Missing required argument: agent_name")?; + + // Verify the agent exists and is running/pending. + let agents = ctx.agents.list_agents()?; + let agent_info = agents + .iter() + .find(|a| a.story_id == story_id && a.agent_name == agent_name) + .ok_or_else(|| format!("No agent '{agent_name}' found for story '{story_id}'"))?; + + if !matches!( + agent_info.status, + crate::agents::AgentStatus::Running | crate::agents::AgentStatus::Pending + ) { + return Err(format!( + "Agent '{agent_name}' for story '{story_id}' is not running (status: {})", + agent_info.status + )); + } + + let project_root = ctx.agents.get_project_root(&ctx.state)?; + let config = ProjectConfig::load(&project_root)?; + + // Find the agent config (max_turns, max_budget_usd). + let agent_config = config.agent.iter().find(|a| a.name == agent_name); + + let max_turns = agent_config.and_then(|a| a.max_turns); + let max_budget_usd = agent_config.and_then(|a| a.max_budget_usd); + + // Count turns by reading log files and counting assistant events. + let log_files = + crate::agent_log::list_story_log_files(&project_root, story_id, Some(agent_name)); + let mut turns_used: u64 = 0; + for path in &log_files { + if let Ok(entries) = crate::agent_log::read_log(path) { + for entry in &entries { + if entry.event.get("type").and_then(|v| v.as_str()) == Some("agent_json") + && let Some(data) = entry.event.get("data") + && data.get("type").and_then(|v| v.as_str()) == Some("assistant") + { + turns_used += 1; + } + } + } + } + + // Compute budget used from completed-session token usage records. + let all_records = crate::agents::token_usage::read_all(&project_root).unwrap_or_default(); + let budget_used_usd: f64 = all_records + .iter() + .filter(|r| r.story_id == story_id && r.agent_name == agent_name) + .map(|r| r.usage.total_cost_usd) + .sum(); + + let remaining_turns = max_turns.map(|max| (max as i64) - (turns_used as i64)); + let remaining_budget_usd = max_budget_usd.map(|max| max - budget_used_usd); + + serde_json::to_string_pretty(&json!({ + "story_id": story_id, + "agent_name": agent_name, + "status": agent_info.status.to_string(), + "turns_used": turns_used, + "max_turns": max_turns, + "remaining_turns": remaining_turns, + "budget_used_usd": budget_used_usd, + "max_budget_usd": max_budget_usd, + "remaining_budget_usd": remaining_budget_usd, + })) + .map_err(|e| format!("Serialization error: {e}")) +} + pub(super) async fn tool_wait_for_agent(args: &Value, ctx: &AppContext) -> Result { let story_id = args .get("story_id") @@ -840,4 +926,112 @@ stage = "coder" let pct = read_coverage_percent_from_json(tmp.path()); assert!(pct.is_none()); } + + // ── get_agent_remaining_turns_and_budget tests ────────────────────────── + + #[test] + fn tool_get_agent_remaining_turns_and_budget_missing_story_id() { + let tmp = tempfile::tempdir().unwrap(); + let ctx = test_ctx(tmp.path()); + let result = + tool_get_agent_remaining_turns_and_budget(&json!({"agent_name": "coder-1"}), &ctx); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("story_id")); + } + + #[test] + fn tool_get_agent_remaining_turns_and_budget_missing_agent_name() { + let tmp = tempfile::tempdir().unwrap(); + let ctx = test_ctx(tmp.path()); + let result = + tool_get_agent_remaining_turns_and_budget(&json!({"story_id": "1_test"}), &ctx); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("agent_name")); + } + + #[test] + fn tool_get_agent_remaining_turns_and_budget_no_agent_returns_error() { + let tmp = tempfile::tempdir().unwrap(); + let ctx = test_ctx(tmp.path()); + let result = tool_get_agent_remaining_turns_and_budget( + &json!({"story_id": "99_nope", "agent_name": "coder-1"}), + &ctx, + ); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + err.contains("No agent"), + "expected 'No agent' error, got: {err}" + ); + } + + #[test] + fn tool_get_agent_remaining_turns_and_budget_completed_agent_returns_error() { + use crate::agents::AgentStatus; + let tmp = tempfile::tempdir().unwrap(); + let ctx = test_ctx(tmp.path()); + ctx.agents + .inject_test_agent("42_story", "coder-1", AgentStatus::Completed); + + let result = tool_get_agent_remaining_turns_and_budget( + &json!({"story_id": "42_story", "agent_name": "coder-1"}), + &ctx, + ); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + err.contains("not running"), + "expected 'not running' error, got: {err}" + ); + } + + #[test] + fn tool_get_agent_remaining_turns_and_budget_running_agent_returns_data() { + use crate::agents::AgentStatus; + use crate::store::StoreOps; + + let tmp = tempfile::tempdir().unwrap(); + let ctx = test_ctx(tmp.path()); + ctx.store + .set("project_root", json!(tmp.path().to_string_lossy().as_ref())); + ctx.agents + .inject_test_agent("42_story", "coder-1", AgentStatus::Running); + + let result = tool_get_agent_remaining_turns_and_budget( + &json!({"story_id": "42_story", "agent_name": "coder-1"}), + &ctx, + ) + .unwrap(); + let parsed: Value = serde_json::from_str(&result).unwrap(); + + assert_eq!(parsed["story_id"], "42_story"); + assert_eq!(parsed["agent_name"], "coder-1"); + assert_eq!(parsed["status"], "running"); + assert!(parsed.get("turns_used").is_some()); + assert!(parsed.get("budget_used_usd").is_some()); + // max_turns and max_budget_usd may be null if not configured + assert!(parsed.get("max_turns").is_some()); + assert!(parsed.get("remaining_turns").is_some()); + assert!(parsed.get("max_budget_usd").is_some()); + assert!(parsed.get("remaining_budget_usd").is_some()); + } + + #[test] + fn tool_get_agent_remaining_turns_and_budget_in_tools_list() { + use super::super::handle_tools_list; + let resp = handle_tools_list(Some(json!(1))); + let tools = resp.result.unwrap()["tools"].as_array().unwrap().clone(); + let tool = tools + .iter() + .find(|t| t["name"] == "get_agent_remaining_turns_and_budget"); + assert!( + tool.is_some(), + "get_agent_remaining_turns_and_budget missing from tools list" + ); + let t = tool.unwrap(); + let required = t["inputSchema"]["required"].as_array().unwrap(); + let req_names: Vec<&str> = required.iter().map(|v| v.as_str().unwrap()).collect(); + assert!(req_names.contains(&"story_id")); + assert!(req_names.contains(&"agent_name")); + } } diff --git a/server/src/http/mcp/mod.rs b/server/src/http/mcp/mod.rs index bcb087a1..ab1cff5c 100644 --- a/server/src/http/mcp/mod.rs +++ b/server/src/http/mcp/mod.rs @@ -431,6 +431,24 @@ fn handle_tools_list(id: Option) -> JsonRpcResponse { "required": ["story_id", "agent_name"] } }, + { + "name": "get_agent_remaining_turns_and_budget", + "description": "Get remaining turns and budget for a running agent. Returns turns used, max turns, remaining turns, budget used (from completed sessions), max budget, and remaining budget. Only works for agents in running or pending state.", + "inputSchema": { + "type": "object", + "properties": { + "story_id": { + "type": "string", + "description": "Story identifier (e.g. '42_story_my_feature')" + }, + "agent_name": { + "type": "string", + "description": "Agent name (e.g. 'coder-1', 'mergemaster', 'qa')" + } + }, + "required": ["story_id", "agent_name"] + } + }, { "name": "create_worktree", "description": "Create a git worktree for a story under .huskies/worktrees/{story_id} with deterministic naming. Writes .mcp.json and runs component setup. Returns the worktree path.", @@ -1272,6 +1290,9 @@ async fn handle_tools_call(id: Option, params: &Value, ctx: &AppContext) "reload_agent_config" => agent_tools::tool_get_agent_config(ctx), "get_agent_output" => agent_tools::tool_get_agent_output(&args, ctx).await, "wait_for_agent" => agent_tools::tool_wait_for_agent(&args, ctx).await, + "get_agent_remaining_turns_and_budget" => { + agent_tools::tool_get_agent_remaining_turns_and_budget(&args, ctx) + } // Worktree tools "create_worktree" => agent_tools::tool_create_worktree(&args, ctx).await, "list_worktrees" => agent_tools::tool_list_worktrees(ctx), @@ -1423,6 +1444,7 @@ mod tests { assert!(names.contains(&"reload_agent_config")); assert!(names.contains(&"get_agent_output")); assert!(names.contains(&"wait_for_agent")); + assert!(names.contains(&"get_agent_remaining_turns_and_budget")); assert!(names.contains(&"create_worktree")); assert!(names.contains(&"list_worktrees")); assert!(names.contains(&"remove_worktree")); @@ -1469,7 +1491,7 @@ mod tests { assert!(names.contains(&"dump_crdt")); assert!(names.contains(&"get_version")); assert!(names.contains(&"remove_criterion")); - assert_eq!(tools.len(), 65); + assert_eq!(tools.len(), 66); } #[test]