huskies: merge 571_story_expose_agent_remaining_turns_and_budget_via_mcp_tool

2026-04-15 18:26:28 +00:00
parent 149a383447
commit ce37281333
2 changed files with 217 additions and 1 deletions
@@ -230,6 +230,92 @@ pub(super) fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String>
    .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<String, String> {
+    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<String, String> {
    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"));
+    }
 }