Accept story 41: Agent Completion Notification via MCP

Add wait_for_agent MCP tool that blocks until an agent reaches a terminal
state (completed, failed, stopped). Returns final status with session_id,
worktree_path, and git commits made by the agent.

- Subscribe-before-check pattern avoids race conditions
- Handles lagged receivers, channel closure, and configurable timeout
- Default timeout 5 minutes, includes git log of agent commits in response
- 11 new tests covering all paths

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Dave
2026-02-20 13:16:04 +00:00
parent fa5d013fe2
commit c6a04f5e53
3 changed files with 387 additions and 16 deletions

View File

@@ -499,6 +499,28 @@ fn handle_tools_list(id: Option<Value>) -> JsonRpcResponse {
},
"required": ["story_id", "agent_name"]
}
},
{
"name": "wait_for_agent",
"description": "Block until the agent reaches a terminal state (completed, failed, stopped). Returns final status and summary including session_id, worktree_path, and any commits made. Use this instead of polling get_agent_output when you want to fire-and-forget and be notified on completion.",
"inputSchema": {
"type": "object",
"properties": {
"story_id": {
"type": "string",
"description": "Story identifier"
},
"agent_name": {
"type": "string",
"description": "Agent name to wait for"
},
"timeout_ms": {
"type": "integer",
"description": "Maximum time to wait in milliseconds (default: 300000 = 5 minutes)"
}
},
"required": ["story_id", "agent_name"]
}
}
]
}),
@@ -533,6 +555,7 @@ async fn handle_tools_call(
"get_agent_config" => tool_get_agent_config(ctx),
"reload_agent_config" => tool_get_agent_config(ctx),
"get_agent_output" => tool_get_agent_output_poll(&args, ctx).await,
"wait_for_agent" => tool_wait_for_agent(&args, ctx).await,
_ => Err(format!("Unknown tool: {tool_name}")),
};
@@ -789,6 +812,71 @@ fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String> {
.map_err(|e| format!("Serialization error: {e}"))
}
async fn tool_wait_for_agent(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")?;
let timeout_ms = args
.get("timeout_ms")
.and_then(|v| v.as_u64())
.unwrap_or(300_000); // default: 5 minutes
let info = ctx
.agents
.wait_for_agent(story_id, agent_name, timeout_ms)
.await?;
let commits = match (&info.worktree_path, &info.base_branch) {
(Some(wt_path), Some(base)) => get_worktree_commits(wt_path, base).await,
_ => None,
};
serde_json::to_string_pretty(&json!({
"story_id": info.story_id,
"agent_name": info.agent_name,
"status": info.status.to_string(),
"session_id": info.session_id,
"worktree_path": info.worktree_path,
"base_branch": info.base_branch,
"commits": commits,
}))
.map_err(|e| format!("Serialization error: {e}"))
}
/// Run `git log <base>..HEAD --oneline` in the worktree and return the commit
/// summaries, or `None` if git is unavailable or there are no new commits.
async fn get_worktree_commits(worktree_path: &str, base_branch: &str) -> Option<Vec<String>> {
let wt = worktree_path.to_string();
let base = base_branch.to_string();
tokio::task::spawn_blocking(move || {
let output = std::process::Command::new("git")
.args(["log", &format!("{base}..HEAD"), "--oneline"])
.current_dir(&wt)
.output()
.ok()?;
if output.status.success() {
let lines: Vec<String> = String::from_utf8(output.stdout)
.ok()?
.lines()
.filter(|l| !l.is_empty())
.map(|l| l.to_string())
.collect();
Some(lines)
} else {
None
}
})
.await
.ok()
.flatten()
}
// ── Helpers ───────────────────────────────────────────────────────
fn parse_test_cases(value: Option<&Value>) -> Result<Vec<TestCaseResult>, String> {
@@ -904,7 +992,8 @@ mod tests {
assert!(names.contains(&"get_agent_config"));
assert!(names.contains(&"reload_agent_config"));
assert!(names.contains(&"get_agent_output"));
assert_eq!(tools.len(), 12);
assert!(names.contains(&"wait_for_agent"));
assert_eq!(tools.len(), 13);
}
#[test]
@@ -1084,4 +1173,69 @@ mod tests {
"application/json"
);
}
#[test]
fn wait_for_agent_tool_in_list() {
let resp = handle_tools_list(Some(json!(1)));
let tools = resp.result.unwrap()["tools"].as_array().unwrap().clone();
let wait_tool = tools.iter().find(|t| t["name"] == "wait_for_agent");
assert!(wait_tool.is_some(), "wait_for_agent missing from tools list");
let t = wait_tool.unwrap();
assert!(t["description"].as_str().unwrap().contains("block") || t["description"].as_str().unwrap().contains("Block"));
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"));
}
#[tokio::test]
async fn wait_for_agent_tool_missing_story_id() {
let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path());
let result = tool_wait_for_agent(&json!({"agent_name": "bot"}), &ctx).await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("story_id"));
}
#[tokio::test]
async fn wait_for_agent_tool_missing_agent_name() {
let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path());
let result = tool_wait_for_agent(&json!({"story_id": "1_test"}), &ctx).await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("agent_name"));
}
#[tokio::test]
async fn wait_for_agent_tool_nonexistent_agent_returns_error() {
let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path());
let result =
tool_wait_for_agent(&json!({"story_id": "99_nope", "agent_name": "bot", "timeout_ms": 50}), &ctx)
.await;
// No agent registered — should error
assert!(result.is_err());
}
#[tokio::test]
async fn wait_for_agent_tool_returns_completed_agent() {
use crate::agents::AgentStatus;
let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path());
ctx.agents
.inject_test_agent("41_story", "worker", AgentStatus::Completed);
let result = tool_wait_for_agent(
&json!({"story_id": "41_story", "agent_name": "worker"}),
&ctx,
)
.await
.unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["status"], "completed");
assert_eq!(parsed["story_id"], "41_story");
assert_eq!(parsed["agent_name"], "worker");
// commits key present (may be null since no real worktree)
assert!(parsed.get("commits").is_some());
}
}