fix(mcp): wrap websocket activity broadcasts in JSON-RPC notifications
Fixes a critical bug where the Antigravity MCP client would crash with 'invalid message version tag' during long-running tool executions. The server was broadcasting raw JSON activity objects ({'type': 'activity', 'data': ...}) over the WebSocket connection without wrapping them in the required JSON-RPC 2.0 Notification envelope, violating the protocol expectation on the proxy stub.
This commit is contained in:
Riz Ashraf committed
1 parent
208c5d448f
commit
da1b7cdc9d
6 files changed
+90
-86
No files matched your search
@@ -27,8 +27,7 @@ impl McpTool for AddTaskHandler {
|
||||
.as_secs();
|
||||
let task_id = uuid::Uuid::new_v4().to_string();
|
||||
|
||||
let parent_id = req.parent_id.clone();
|
||||
let deps = req.dependencies.clone().unwrap_or_default();
|
||||
let deps = req.dependencies.unwrap_or_default();
|
||||
|
||||
let task = Task {
|
||||
id: task_id.clone(),
|
||||
@@ -38,12 +37,12 @@ impl McpTool for AddTaskHandler {
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
git_branch: req.git_branch,
|
||||
parent_id,
|
||||
parent_id: req.parent_id,
|
||||
dependencies: deps,
|
||||
acceptance_criteria: vec![],
|
||||
};
|
||||
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||
let _ = idx.index_task(&task).await;
|
||||
let _ = idx.index_task(&task);
|
||||
state.tasks.modify(|tasks| {
|
||||
tasks.push(task);
|
||||
});
|
||||
@@ -69,33 +68,40 @@ impl McpTool for DeleteTaskHandler {
|
||||
let mut actually_deleted = Vec::new();
|
||||
state.tasks.modify(|tasks| {
|
||||
let initial_len = tasks.len();
|
||||
// Collect IDs of tasks to delete (this task + all its recursive children)
|
||||
let mut to_delete = std::collections::HashSet::new();
|
||||
to_delete.insert(req.id.clone());
|
||||
|
||||
let mut children_map: std::collections::HashMap<String, Vec<String>> =
|
||||
std::collections::HashMap::new();
|
||||
for t in tasks.iter() {
|
||||
|
||||
// Build index-based children map
|
||||
let mut children_map: std::collections::HashMap<usize, Vec<usize>> = std::collections::HashMap::new();
|
||||
let mut id_to_index = std::collections::HashMap::new();
|
||||
for (idx, t) in tasks.iter().enumerate() {
|
||||
id_to_index.insert(t.id.as_str(), idx);
|
||||
}
|
||||
|
||||
for (idx, t) in tasks.iter().enumerate() {
|
||||
if let Some(pid) = &t.parent_id {
|
||||
children_map
|
||||
.entry(pid.clone())
|
||||
.or_default()
|
||||
.push(t.id.clone());
|
||||
if let Some(&parent_idx) = id_to_index.get(pid.as_str()) {
|
||||
children_map.entry(parent_idx).or_default().push(idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back(req.id.clone());
|
||||
let mut to_delete_idx = std::collections::HashSet::new();
|
||||
if let Some(&start_idx) = id_to_index.get(req.id.as_str()) {
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back(start_idx);
|
||||
|
||||
while let Some(curr) = queue.pop_front() {
|
||||
if to_delete.insert(curr.clone())
|
||||
&& let Some(children) = children_map.get(&curr)
|
||||
{
|
||||
queue.extend(children.iter().cloned());
|
||||
while let Some(curr) = queue.pop_front() {
|
||||
if to_delete_idx.insert(curr) {
|
||||
if let Some(children) = children_map.get(&curr) {
|
||||
queue.extend(children.iter().copied());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
actually_deleted = to_delete.into_iter().collect();
|
||||
for &idx in &to_delete_idx {
|
||||
actually_deleted.push(tasks[idx].id.clone());
|
||||
}
|
||||
|
||||
tasks.retain(|t| !actually_deleted.contains(&t.id));
|
||||
deleted_count = initial_len - tasks.len();
|
||||
});
|
||||
@@ -103,7 +109,7 @@ impl McpTool for DeleteTaskHandler {
|
||||
if deleted_count > 0 {
|
||||
let idx = state.search_index.read().unwrap_or_else(|e| e.into_inner()).clone();
|
||||
for id in actually_deleted {
|
||||
let _ = idx.delete_document(&id).await;
|
||||
let _ = idx.delete_document(&id);
|
||||
}
|
||||
Ok(vec![
|
||||
format!("Deleted task and its children ({} total).", deleted_count).to_string(),
|
||||
@@ -139,21 +145,18 @@ impl McpTool for UpdateTaskStatusHandler {
|
||||
|
||||
state.tasks.modify(|tasks| {
|
||||
// Find target task
|
||||
let mut target_id = String::new();
|
||||
if let Some(t) = tasks.iter().find(|t| t.id == req.id || t.title == req.id) {
|
||||
target_id = t.id.clone();
|
||||
}
|
||||
|
||||
if target_id.is_empty() {
|
||||
return;
|
||||
}
|
||||
let target_idx = tasks.iter().position(|t| t.id == req.id || t.title == req.id);
|
||||
let target_idx = match target_idx {
|
||||
Some(idx) => idx,
|
||||
None => return,
|
||||
};
|
||||
|
||||
found = true;
|
||||
let target_id = tasks[target_idx].id.clone();
|
||||
|
||||
if target_status == "done" || target_status == "completed" {
|
||||
// 1. Check Acceptance Criteria
|
||||
if let Some(t) = tasks.iter().find(|t| t.id == target_id)
|
||||
&& t.acceptance_criteria.iter().any(|c| !c.is_met)
|
||||
{
|
||||
if tasks[target_idx].acceptance_criteria.iter().any(|c| !c.is_met) {
|
||||
blocked = true;
|
||||
blocker_details = "Unmet acceptance criteria exist.".to_string();
|
||||
}
|
||||
@@ -161,53 +164,41 @@ impl McpTool for UpdateTaskStatusHandler {
|
||||
// 2. Check dependencies
|
||||
if !blocked {
|
||||
let mut uncompleted_deps = Vec::new();
|
||||
if let Some(t) = tasks.iter().find(|t| t.id == target_id) {
|
||||
for dep_id in &t.dependencies {
|
||||
if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id)
|
||||
&& dep_task.status != "completed"
|
||||
&& dep_task.status != "done"
|
||||
{
|
||||
uncompleted_deps.push(dep_task.title.clone());
|
||||
for dep_id in &tasks[target_idx].dependencies {
|
||||
if let Some(dep_task) = tasks.iter().find(|dt| dt.id == *dep_id) {
|
||||
if dep_task.status != "completed" && dep_task.status != "done" {
|
||||
uncompleted_deps.push(dep_task.title.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
if !uncompleted_deps.is_empty() {
|
||||
blocked = true;
|
||||
blocker_details =
|
||||
format!("Blocked by dependencies: {}", uncompleted_deps.join(", "));
|
||||
blocker_details = format!("Blocked by dependencies: {}", uncompleted_deps.join(", "));
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Check child tasks
|
||||
if !blocked {
|
||||
let mut uncompleted_children = Vec::new();
|
||||
for child in tasks
|
||||
.iter()
|
||||
.filter(|t| t.parent_id.as_ref() == Some(&target_id))
|
||||
{
|
||||
for child in tasks.iter().filter(|t| t.parent_id.as_ref() == Some(&target_id)) {
|
||||
if child.status != "completed" && child.status != "done" {
|
||||
uncompleted_children.push(child.title.clone());
|
||||
uncompleted_children.push(child.title.as_str());
|
||||
}
|
||||
}
|
||||
if !uncompleted_children.is_empty() {
|
||||
blocked = true;
|
||||
blocker_details = format!(
|
||||
"Blocked by child tasks: {}",
|
||||
uncompleted_children.join(", ")
|
||||
);
|
||||
blocker_details = format!("Blocked by child tasks: {}", uncompleted_children.join(", "));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !blocked {
|
||||
// Apply update
|
||||
if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) {
|
||||
t.status = target_status.clone();
|
||||
t.updated_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
}
|
||||
tasks[target_idx].status = target_status.clone();
|
||||
tasks[target_idx].updated_at = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
// Cascade cancellation to children
|
||||
if target_status == "cancelled" || target_status == "abandoned" {
|
||||
|
||||
Reference in new issue
Block a user