diff --git a/server/src/api/setup.rs b/server/src/api/setup.rs index f95451e..c24b9c3 100644 --- a/server/src/api/setup.rs +++ b/server/src/api/setup.rs @@ -34,7 +34,7 @@ pub fn create_router(app_state: Arc) -> Router { let entity_count = graph.entities.len(); let relation_count = graph.relations.len(); let tasks = state_clone.project.tasks.cache.read().unwrap(); - let active_tasks = tasks.iter().filter(|t| t.status != "completed" && t.status != "done").count(); + let active_tasks = tasks.iter().filter(|t| t.is_active()).count(); let adrs = state_clone.code.adrs.cache.read().unwrap(); let adr_count = adrs.len(); let tech_debts = state_clone.code.tech_debts.cache.read().unwrap(); @@ -464,10 +464,7 @@ mod tests { async fn test_ping_endpoint() { let (app, _, _dir) = setup_app().await; - let request = Request::builder() - .uri("/ping") - .body(Body::empty()) - .unwrap(); + let request = Request::builder().uri("/ping").body(Body::empty()).unwrap(); let response = app.oneshot(request).await.unwrap(); assert_eq!(response.status(), StatusCode::OK); @@ -611,12 +608,15 @@ mod tests { .method("POST") .uri("/gate/set") .header("content-type", "application/json") - .body(Body::from(serde_json::json!({ - "action": "deploy", - "target": "prod", - "authorize": true, - "reason": "Tests passed" - }).to_string())) + .body(Body::from( + serde_json::json!({ + "action": "deploy", + "target": "prod", + "authorize": true, + "reason": "Tests passed" + }) + .to_string(), + )) .unwrap(); let response = app.oneshot(set_req).await.unwrap(); @@ -655,10 +655,7 @@ mod tests { for ep in endpoints { let app_inst = create_router(app_state.clone()); - let req = Request::builder() - .uri(ep) - .body(Body::empty()) - .unwrap(); + let req = Request::builder().uri(ep).body(Body::empty()).unwrap(); let resp = app_inst.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK, "Failed endpoint: {}", ep); } @@ -693,22 +690,39 @@ mod tests { let req_html = Request::builder().uri("/").body(Body::empty()).unwrap(); let resp_html = app_html.oneshot(req_html).await.unwrap(); assert_eq!(resp_html.status(), StatusCode::OK); - let content_type = resp_html.headers().get(axum::http::header::CONTENT_TYPE).unwrap().to_str().unwrap(); + let content_type = resp_html + .headers() + .get(axum::http::header::CONTENT_TYPE) + .unwrap() + .to_str() + .unwrap(); assert!(content_type.contains("text/html")); - let body_bytes = axum::body::to_bytes(resp_html.into_body(), usize::MAX).await.unwrap(); + let body_bytes = axum::body::to_bytes(resp_html.into_body(), usize::MAX) + .await + .unwrap(); let html_str = String::from_utf8(body_bytes.to_vec()).unwrap(); assert!(html_str.contains("")); // Test GET /dashboard.js let app_js = create_router(app_state.clone()); - let req_js = Request::builder().uri("/dashboard.js").body(Body::empty()).unwrap(); + let req_js = Request::builder() + .uri("/dashboard.js") + .body(Body::empty()) + .unwrap(); let resp_js = app_js.oneshot(req_js).await.unwrap(); assert_eq!(resp_js.status(), StatusCode::OK); - let content_type_js = resp_js.headers().get(axum::http::header::CONTENT_TYPE).unwrap().to_str().unwrap(); + let content_type_js = resp_js + .headers() + .get(axum::http::header::CONTENT_TYPE) + .unwrap() + .to_str() + .unwrap(); assert!(content_type_js.contains("application/javascript")); - let js_bytes = axum::body::to_bytes(resp_js.into_body(), usize::MAX).await.unwrap(); + let js_bytes = axum::body::to_bytes(resp_js.into_body(), usize::MAX) + .await + .unwrap(); let js_str = String::from_utf8(js_bytes.to_vec()).unwrap(); assert!(js_str.contains("escapeHtml")); assert!(js_str.contains("setupWS")); @@ -722,15 +736,24 @@ mod tests { // 1. Initial GET /api/activity should return empty array [] let app_act = create_router(app_state.clone()); - let req_act = Request::builder().uri("/api/activity").body(Body::empty()).unwrap(); + let req_act = Request::builder() + .uri("/api/activity") + .body(Body::empty()) + .unwrap(); let resp_act = app_act.oneshot(req_act).await.unwrap(); assert_eq!(resp_act.status(), StatusCode::OK); - let body_bytes = axum::body::to_bytes(resp_act.into_body(), usize::MAX).await.unwrap(); + let body_bytes = axum::body::to_bytes(resp_act.into_body(), usize::MAX) + .await + .unwrap(); let json_str = String::from_utf8(body_bytes.to_vec()).unwrap(); assert_eq!(json_str.trim(), "[]"); // 2. Record a direct activity - app_state.handler.state.record_activity("code_change", "Refactored live activity UI", Some("Added empty state placeholder")); + app_state.handler.state.record_activity( + "code_change", + "Refactored live activity UI", + Some("Added empty state placeholder"), + ); // 3. Post terminal telemetry activity let app_term = create_router(app_state.clone()); @@ -738,13 +761,16 @@ mod tests { .method("POST") .uri("/terminal/telemetry") .header("content-type", "application/json") - .body(Body::from(serde_json::json!({ - "command": "cargo test --workspace", - "exit_code": 0, - "cwd": "C:\\workspace\\mcp-memory", - "os": "windows", - "timestamp": 1728130000000u64 - }).to_string())) + .body(Body::from( + serde_json::json!({ + "command": "cargo test --workspace", + "exit_code": 0, + "cwd": "C:\\workspace\\mcp-memory", + "os": "windows", + "timestamp": 1728130000000u64 + }) + .to_string(), + )) .unwrap(); let term_resp = app_term.oneshot(term_req).await.unwrap(); assert_eq!(term_resp.status(), StatusCode::OK); @@ -755,23 +781,31 @@ mod tests { .method("POST") .uri("/nvim/telemetry") .header("content-type", "application/json") - .body(Body::from(serde_json::json!({ - "session_id": "test_session_1", - "event": "BufWritePost", - "file": "server/src/dashboard.ts" - }).to_string())) + .body(Body::from( + serde_json::json!({ + "session_id": "test_session_1", + "event": "BufWritePost", + "file": "server/src/dashboard.ts" + }) + .to_string(), + )) .unwrap(); let nvim_resp = app_nvim.oneshot(nvim_req).await.unwrap(); assert_eq!(nvim_resp.status(), StatusCode::OK); // 5. GET /api/activity should return all recorded items let app_act2 = create_router(app_state.clone()); - let req_act2 = Request::builder().uri("/api/activity").body(Body::empty()).unwrap(); + let req_act2 = Request::builder() + .uri("/api/activity") + .body(Body::empty()) + .unwrap(); let resp_act2 = app_act2.oneshot(req_act2).await.unwrap(); assert_eq!(resp_act2.status(), StatusCode::OK); - let body_bytes2 = axum::body::to_bytes(resp_act2.into_body(), usize::MAX).await.unwrap(); + let body_bytes2 = axum::body::to_bytes(resp_act2.into_body(), usize::MAX) + .await + .unwrap(); let json_str2 = String::from_utf8(body_bytes2.to_vec()).unwrap(); - + assert!(json_str2.contains("CODE_CHANGE")); assert!(json_str2.contains("Refactored live activity UI")); assert!(json_str2.contains("TERMINAL")); @@ -781,19 +815,18 @@ mod tests { // 6. Verify /dashboard.js contains empty state text and stream logic let app_js = create_router(app_state.clone()); - let req_js = Request::builder().uri("/dashboard.js").body(Body::empty()).unwrap(); + let req_js = Request::builder() + .uri("/dashboard.js") + .body(Body::empty()) + .unwrap(); let resp_js = app_js.oneshot(req_js).await.unwrap(); assert_eq!(resp_js.status(), StatusCode::OK); - let js_bytes = axum::body::to_bytes(resp_js.into_body(), usize::MAX).await.unwrap(); + let js_bytes = axum::body::to_bytes(resp_js.into_body(), usize::MAX) + .await + .unwrap(); let js_str = String::from_utf8(js_bytes.to_vec()).unwrap(); assert!(js_str.contains("No recent activity recorded yet.")); assert!(js_str.contains("/api/activity/stream")); assert!(js_str.contains("parseActivityPayload")); } } - - - - - - diff --git a/server/src/db.rs b/server/src/db.rs index 96b0732..c87a2d2 100644 --- a/server/src/db.rs +++ b/server/src/db.rs @@ -14,36 +14,55 @@ pub fn init_redb(base: &Path) -> Arc { } else { let redb_path = base.join("mcp_store.redb"); if redb_path.exists() { - match redb::Database::open(&redb_path) { - Ok(db) => Arc::new(db), - Err(open_err) => { - let err_msg = format!( - "Failed to open existing redb database at {:?}: {}. Attempting to recreate database.", - redb_path, open_err - ); - tracing::warn!("{}", err_msg); - match redb::Database::create(&redb_path) { - Ok(db) => Arc::new(db), - Err(create_err) => { - if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { - let temp_path = std::env::temp_dir().join(format!( - "mcp_store_fallback_{}.redb", - uuid::Uuid::new_v4() - )); - tracing::warn!( - "CRITICAL PERSISTENCE ALERT: MCP_ALLOW_TMP_FALLBACK=1 set. Using temporary redb database {:?}. Changes will be discarded upon application exit.", - temp_path - ); - Arc::new( - redb::Database::create(&temp_path) - .expect("Failed to create fallback redb database"), - ) - } else { - panic!( - "CRITICAL PERSISTENCE FAILURE: Unable to open or create primary database at {:?}: (open: {}, create: {}). To prevent silent data loss on restart, the application cannot start without accessible persistence.", - redb_path, open_err, create_err - ); - } + let mut db_opt = None; + let mut last_open_err = String::new(); + for attempt in 1..=3 { + match redb::Database::open(&redb_path) { + Ok(db) => { + db_opt = Some(Arc::new(db)); + break; + } + Err(open_err) => { + last_open_err = open_err.to_string(); + if attempt < 3 { + tracing::warn!( + "Transient lock contention opening redb at {:?} (attempt {}/3: {}). Retrying...", + redb_path, + attempt, + open_err + ); + std::thread::sleep(std::time::Duration::from_millis(150)); + } + } + } + } + if let Some(db) = db_opt { + db + } else { + let err_msg = format!( + "Failed to open existing redb database at {:?}: {}. Attempting to recreate database.", + redb_path, last_open_err + ); + tracing::warn!("{}", err_msg); + match redb::Database::create(&redb_path) { + Ok(db) => Arc::new(db), + Err(create_err) => { + if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { + let temp_path = std::env::temp_dir() + .join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4())); + tracing::warn!( + "CRITICAL PERSISTENCE ALERT: MCP_ALLOW_TMP_FALLBACK=1 set. Using temporary redb database {:?}. Changes will be discarded upon application exit.", + temp_path + ); + Arc::new( + redb::Database::create(&temp_path) + .expect("Failed to create fallback redb database"), + ) + } else { + panic!( + "CRITICAL PERSISTENCE FAILURE: Unable to open or create primary database at {:?}: (open: {}, create: {}). To prevent silent data loss on restart, the application cannot start without accessible persistence.", + redb_path, last_open_err, create_err + ); } } } diff --git a/server/src/handlers/meta.rs b/server/src/handlers/meta.rs index 8f6ff5a..03c81aa 100644 --- a/server/src/handlers/meta.rs +++ b/server/src/handlers/meta.rs @@ -49,7 +49,9 @@ impl McpTool for LogDecisionHandler { status: "accepted".to_string(), supersedes: req.supersedes, timestamp: crate::handlers::utils::now_secs(), - namespace: req.namespace.unwrap_or_else(|| crate::models::default_namespace()), + namespace: req + .namespace + .unwrap_or_else(|| crate::models::default_namespace()), repo_name: req.repo_name, alternatives_considered: req.alternatives_considered.unwrap_or_default(), affected_components: req.affected_components.unwrap_or_default(), @@ -65,7 +67,11 @@ impl McpTool for LogDecisionHandler { drop(idx.index_adr(adr)); } - state.record_activity("decision", &format!("Logged {}: {}", final_id, req.title), Some(&req.decision)); + state.record_activity( + "decision", + &format!("Logged {}: {}", final_id, req.title), + Some(&req.decision), + ); Ok(format!("Logged decision {}: {}", final_id, req.title)) } } @@ -107,12 +113,17 @@ impl McpTool for QueryDecisionsHandler { if include_body { Ok::(serde_json::to_string(&filtered)?) } else { - let compact: Vec<_> = filtered.iter().map(|a| serde_json::json!({ - "id": a.id, - "title": a.title, - "status": a.status, - "timestamp": a.timestamp, - })).collect(); + let compact: Vec<_> = filtered + .iter() + .map(|a| { + serde_json::json!({ + "id": a.id, + "title": a.title, + "status": a.status, + "timestamp": a.timestamp, + }) + }) + .collect(); Ok::(serde_json::to_string(&compact)?) } })?; @@ -212,8 +223,15 @@ impl McpTool for LogErrorFixHandler { fixes.remove(0); } }); - state.record_activity("error_fix", &format!("Fixed error: {}", req.signature), Some(&solution)); - Ok(format!("Logged error fix for {}: {}", req.signature, solution)) + state.record_activity( + "error_fix", + &format!("Fixed error: {}", req.signature), + Some(&solution), + ); + Ok(format!( + "Logged error fix for {}: {}", + req.signature, solution + )) } } @@ -267,22 +285,26 @@ impl McpTool for SearchErrorFixesHandler { scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| { - if include_body { - serde_json::json!({ - "signature": f.signature, - "solution": f.solution, - "git_commit": f.git_commit, - "git_branch": f.git_branch, - "match_score": score - }) - } else { - serde_json::json!({ - "signature": f.signature, - "match_score": score - }) - } - }).collect(); + let suggestions: Vec<_> = scored + .into_iter() + .take(limit) + .map(|(f, score)| { + if include_body { + serde_json::json!({ + "signature": f.signature, + "solution": f.solution, + "git_commit": f.git_commit, + "git_branch": f.git_branch, + "match_score": score + }) + } else { + serde_json::json!({ + "signature": f.signature, + "match_score": score + }) + } + }) + .collect(); Ok::(serde_json::to_string_pretty(&suggestions)?) })?; @@ -357,9 +379,18 @@ impl McpTool for LogCodeChangeHandler { } } } - let change_kind = match req.change_kind.as_deref().map(|s| s.to_lowercase()).as_deref() { - Some("added") | Some("create") | Some("created") | Some("new") => crate::models::ChangeKind::Added, - Some("deleted") | Some("remove") | Some("removed") => crate::models::ChangeKind::Deleted, + let change_kind = match req + .change_kind + .as_deref() + .map(|s| s.to_lowercase()) + .as_deref() + { + Some("added") | Some("create") | Some("created") | Some("new") => { + crate::models::ChangeKind::Added + } + Some("deleted") | Some("remove") | Some("removed") => { + crate::models::ChangeKind::Deleted + } Some("renamed") | Some("move") | Some("moved") => crate::models::ChangeKind::Renamed, _ => crate::models::ChangeKind::Modified, }; @@ -380,13 +411,21 @@ impl McpTool for LogCodeChangeHandler { } else if let Some(ref rev) = effective_rev { if rev.starts_with('r') && rev[1..].chars().all(|c| c.is_ascii_digit()) { Some("svn".to_string()) - } else if req.repo_url.as_deref().map(|u| u.contains("/svn/")).unwrap_or(false) { + } else if req + .repo_url + .as_deref() + .map(|u| u.contains("/svn/")) + .unwrap_or(false) + { Some("svn".to_string()) } else { Some("git".to_string()) } } else if let Some(ref br) = effective_branch { - if br.eq_ignore_ascii_case("trunk") || br.starts_with("branches/") || br.starts_with("tags/") { + if br.eq_ignore_ascii_case("trunk") + || br.starts_with("branches/") + || br.starts_with("tags/") + { Some("svn".to_string()) } else { Some("git".to_string()) @@ -419,8 +458,15 @@ impl McpTool for LogCodeChangeHandler { ledger.remove(0); } }); - state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description)); - Ok(format!("Logged code change for {}: {}", req.file_path, description)) + state.record_activity( + "code_change", + &format!("Modified {}", req.file_path), + Some(&description), + ); + Ok(format!( + "Logged code change for {}: {}", + req.file_path, description + )) } } @@ -440,12 +486,13 @@ impl McpTool for QueryRecentChangesHandler { } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: QueryRecentChangesTool = serde_json::from_value(args).unwrap_or(QueryRecentChangesTool { - namespace: None, - repo_name: None, - vcs_type: None, - limit: None, - }); + let req: QueryRecentChangesTool = + serde_json::from_value(args).unwrap_or(QueryRecentChangesTool { + namespace: None, + repo_name: None, + vcs_type: None, + limit: None, + }); let limit = req.limit.unwrap_or(50); let data = state.code.ledger.read_with(|l| { let filtered: Vec<_> = l @@ -458,7 +505,11 @@ impl McpTool for QueryRecentChangesHandler { } } if let Some(repo) = &req.repo_name { - if c.repo_name.as_ref().map(|rn| !rn.eq_ignore_ascii_case(repo)).unwrap_or(true) { + if c.repo_name + .as_ref() + .map(|rn| !rn.eq_ignore_ascii_case(repo)) + .unwrap_or(true) + { return false; } } @@ -573,7 +624,10 @@ impl McpTool for LogTechDebtHandler { line_range: req.line_range, workaround: req.workaround, effort_estimate: req.effort_estimate, - }) + }); + if debts.len() > 300 { + debts.remove(0); + } }); Ok("Tech debt logged".to_string()) } @@ -687,7 +741,9 @@ impl McpTool for DecisionsHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: DecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; - let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace()); + let ns = req + .namespace + .unwrap_or_else(|| crate::models::default_namespace()); match req.action { DecisionAction::Log => { let title = req.title.ok_or_else(|| { @@ -697,30 +753,39 @@ impl McpTool for DecisionsHandler { let context = req.context.unwrap_or_default(); let decision = req.decision.unwrap_or_default(); let consequences = req.consequences.unwrap_or_default(); - LogDecisionHandler.execute( - serde_json::json!({ - "title": title, - "status": status, - "context": context, - "decision": decision, - "consequences": consequences, - "namespace": ns, - "repo_name": req.repo_name, - "author": req.author, - "affected_components": req.affected_components, - "alternatives_considered": req.alternatives_considered, - }), - state, - ).await + LogDecisionHandler + .execute( + serde_json::json!({ + "title": title, + "status": status, + "context": context, + "decision": decision, + "consequences": consequences, + "namespace": ns, + "repo_name": req.repo_name, + "author": req.author, + "affected_components": req.affected_components, + "alternatives_considered": req.alternatives_considered, + }), + state, + ) + .await } DecisionAction::Query => { - QueryDecisionsHandler.execute(serde_json::json!({"query": req.query, "namespace": ns}), state).await + QueryDecisionsHandler + .execute( + serde_json::json!({"query": req.query, "namespace": ns}), + state, + ) + .await } DecisionAction::Delete => { let id = req.id.ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'id' for action 'delete'. Next step: Provide ADR 'id' string in request and retry.".to_string()) })?; - DeleteDecisionHandler.execute(serde_json::json!({"id": id}), state).await + DeleteDecisionHandler + .execute(serde_json::json!({"id": id}), state) + .await } } } @@ -743,40 +808,51 @@ impl McpTool for TechDebtHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: TechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?; - let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace()); + let ns = req + .namespace + .unwrap_or_else(|| crate::models::default_namespace()); match req.action { TechDebtAction::Log => { let desc = req.description.ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'description' for action 'log'. Next step: Provide tech debt 'description' in request and retry.".to_string()) })?; let ideal = req.ideal_solution.unwrap_or_default(); - LogTechDebtHandler.execute( - serde_json::json!({ - "description": desc, - "ideal_solution": ideal, - "git_commit": req.git_commit, - "git_branch": req.git_branch, - "symbol_references": req.symbol_references, - "line_range": req.line_range, - "namespace": ns, - "repo_name": req.repo_name, - "severity": req.severity, - "file_path": req.file_path, - "workaround": req.workaround, - "effort_estimate": req.effort_estimate, - }), - state, - ).await + LogTechDebtHandler + .execute( + serde_json::json!({ + "description": desc, + "ideal_solution": ideal, + "git_commit": req.git_commit, + "git_branch": req.git_branch, + "symbol_references": req.symbol_references, + "line_range": req.line_range, + "namespace": ns, + "repo_name": req.repo_name, + "severity": req.severity, + "file_path": req.file_path, + "workaround": req.workaround, + "effort_estimate": req.effort_estimate, + }), + state, + ) + .await } TechDebtAction::Resolve => { let id = req.id.ok_or_else(|| { crate::error::AppError::Internal("Missing required parameter 'id' for action 'resolve'. Next step: Provide tech debt 'id' string in request and retry.".to_string()) })?; - ResolveTechDebtHandler.execute(serde_json::json!({"id": id}), state).await + ResolveTechDebtHandler + .execute(serde_json::json!({"id": id}), state) + .await } TechDebtAction::List => { let inc = req.include_resolved.unwrap_or(false); - ListTechDebtHandler.execute(serde_json::json!({"include_resolved": inc, "namespace": ns}), state).await + ListTechDebtHandler + .execute( + serde_json::json!({"include_resolved": inc, "namespace": ns}), + state, + ) + .await } } } @@ -861,8 +937,7 @@ impl McpTool for OmniSearchHandler { let mut ranked_items: Vec<_> = rrf_scores.into_values().collect(); ranked_items.sort_by(|a, b| b.0.total_cmp(&a.0)); - let matches: Vec = - ranked_items.into_iter().map(|(_, item)| item).collect(); + let matches: Vec = ranked_items.into_iter().map(|(_, item)| item).collect(); let kg_json = state.read_graph(|full| { let mut kg_results = serde_json::Map::new(); @@ -872,14 +947,16 @@ impl McpTool for OmniSearchHandler { let mut adj_map: std::collections::HashMap<&str, Vec<(&str, &str, &str)>> = std::collections::HashMap::new(); for rel in &full.relations { - adj_map - .entry(rel.from.as_str()) - .or_default() - .push((rel.to.as_str(), rel.relation_type.as_str(), "outgoing")); - adj_map - .entry(rel.to.as_str()) - .or_default() - .push((rel.from.as_str(), rel.relation_type.as_str(), "incoming")); + adj_map.entry(rel.from.as_str()).or_default().push(( + rel.to.as_str(), + rel.relation_type.as_str(), + "outgoing", + )); + adj_map.entry(rel.to.as_str()).or_default().push(( + rel.from.as_str(), + rel.relation_type.as_str(), + "incoming", + )); } for res in &matches { @@ -912,7 +989,10 @@ impl McpTool for OmniSearchHandler { }; if let Some(obj) = entity_val.as_object_mut() { - obj.insert("subgraph_relations".to_string(), serde_json::Value::Array(connected_rels)); + obj.insert( + "subgraph_relations".to_string(), + serde_json::Value::Array(connected_rels), + ); } kg_results.insert(res.id.clone(), entity_val); } @@ -1130,17 +1210,18 @@ impl McpTool for OmniSearchHandler { let mut pruned = false; for key in prune_keys { while out_str.len() > max_chars { - let popped = if let Some(arr) = report.get_mut(key).and_then(|v| v.as_array_mut()) { - if arr.len() > 1 { - arr.pop(); - pruned = true; - true + let popped = + if let Some(arr) = report.get_mut(key).and_then(|v| v.as_array_mut()) { + if arr.len() > 1 { + arr.pop(); + pruned = true; + true + } else { + false + } } else { false - } - } else { - false - }; + }; if popped { out_str = report.to_string(); } else { @@ -1180,10 +1261,12 @@ impl McpTool for GetProjectHealthHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: GetProjectHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?; - let active_tasks = state - .project - .tasks - .read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count()); + let active_tasks = state.project.tasks.read_with(|tasks| { + tasks + .iter() + .filter(|t| t.namespace == req.namespace && t.is_active()) + .count() + }); let unresolved_debt = state.code.tech_debts.read_with(|debts| { debts .iter() @@ -1268,15 +1351,19 @@ impl McpTool for SuggestErrorFixHandler { scored.sort_by(|a, b| b.1.total_cmp(&a.1)); - let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| { - serde_json::json!({ - "signature": f.signature, - "solution": f.solution, - "git_commit": f.git_commit, - "git_branch": f.git_branch, - "match_score": score + let suggestions: Vec<_> = scored + .into_iter() + .take(limit) + .map(|(f, score)| { + serde_json::json!({ + "signature": f.signature, + "solution": f.solution, + "git_commit": f.git_commit, + "git_branch": f.git_branch, + "match_score": score + }) }) - }).collect(); + .collect(); Ok::(serde_json::to_string_pretty(&suggestions)?) })?; @@ -1305,26 +1392,43 @@ impl McpTool for ManageCheckpointHandler { match req.action { CheckpointAction::Create => { let name = req.name_or_id.ok_or_else(|| { - crate::error::AppError::Internal("name_or_id is required for 'create' action".to_string()) + crate::error::AppError::Internal( + "name_or_id is required for 'create' action".to_string(), + ) })?; let target_dir = state.base_dir.join("checkpoints").join(&name); if let Err(e) = std::fs::create_dir_all(&target_dir) { - return Err(crate::error::AppError::Internal(format!("Failed to create checkpoint dir: {}", e))); + return Err(crate::error::AppError::Internal(format!( + "Failed to create checkpoint dir: {}", + e + ))); } let graph_json = state.read_graph(|g| serde_json::to_string(g).unwrap_or_default()); let _ = std::fs::write(target_dir.join("graph.json"), graph_json); - let tasks_json = state.project.tasks.read_with(|t| serde_json::to_string(t).unwrap_or_default()); + let tasks_json = state + .project + .tasks + .read_with(|t| serde_json::to_string(t).unwrap_or_default()); let _ = std::fs::write(target_dir.join("tasks.json"), tasks_json); - let debts_json = state.code.tech_debts.read_with(|d| serde_json::to_string(d).unwrap_or_default()); + let debts_json = state + .code + .tech_debts + .read_with(|d| serde_json::to_string(d).unwrap_or_default()); let _ = std::fs::write(target_dir.join("tech_debts.json"), debts_json); if let Some(desc) = &req.description { - let snapshot_id = format!("SNAP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase()); - let ns = req.namespace.clone().unwrap_or_else(|| "global".to_string()); + let snapshot_id = format!( + "SNAP-{}", + uuid::Uuid::new_v4().to_string()[..8].to_uppercase() + ); + let ns = req + .namespace + .clone() + .unwrap_or_else(|| "global".to_string()); let snapshot = crate::models::StateSnapshot { id: snapshot_id, timestamp: crate::handlers::utils::now_secs(), @@ -1339,16 +1443,27 @@ impl McpTool for ManageCheckpointHandler { } CheckpointAction::Restore => { let name = req.name_or_id.ok_or_else(|| { - crate::error::AppError::Internal("name_or_id is required for 'restore' action".to_string()) + crate::error::AppError::Internal( + "name_or_id is required for 'restore' action".to_string(), + ) })?; let target_dir = state.base_dir.join("checkpoints").join(&name); if !target_dir.exists() { - let found = state.project.snapshots.read_with(|snaps| snaps.iter().any(|s| s.id == name)); + let found = state + .project + .snapshots + .read_with(|snaps| snaps.iter().any(|s| s.id == name)); if found { - return Ok(format!("Successfully restored memory state from snapshot {}", name)); + return Ok(format!( + "Successfully restored memory state from snapshot {}", + name + )); } - return Err(crate::error::AppError::Internal(format!("Checkpoint or snapshot '{}' does not exist.", name))); + return Err(crate::error::AppError::Internal(format!( + "Checkpoint or snapshot '{}' does not exist.", + name + ))); } if let Ok(graph_content) = std::fs::read_to_string(target_dir.join("graph.json")) { @@ -1363,7 +1478,9 @@ impl McpTool for ManageCheckpointHandler { } } - if let Ok(debts_content) = std::fs::read_to_string(target_dir.join("tech_debts.json")) { + if let Ok(debts_content) = + std::fs::read_to_string(target_dir.join("tech_debts.json")) + { if let Ok(debts) = serde_json::from_str(&debts_content) { state.code.tech_debts.modify(|d| *d = debts); } @@ -1391,7 +1508,9 @@ impl McpTool for ManageCheckpointHandler { } CheckpointAction::Delete => { let name = req.name_or_id.ok_or_else(|| { - crate::error::AppError::Internal("name_or_id is required for 'delete' action".to_string()) + crate::error::AppError::Internal( + "name_or_id is required for 'delete' action".to_string(), + ) })?; let target_dir = state.base_dir.join("checkpoints").join(&name); if target_dir.exists() { @@ -1400,7 +1519,10 @@ impl McpTool for ManageCheckpointHandler { state.project.snapshots.modify(|snaps| { snaps.retain(|s| s.id != name); }); - Ok(format!("Checkpoint or snapshot '{}' deleted successfully.", name)) + Ok(format!( + "Checkpoint or snapshot '{}' deleted successfully.", + name + )) } } } @@ -1426,17 +1548,26 @@ impl McpTool for CheckpointStateHandler { let target_dir = state.base_dir.join("checkpoints").join(&req.name); if let Err(e) = std::fs::create_dir_all(&target_dir) { - return Err(crate::error::AppError::Internal(format!("Failed to create checkpoint dir: {}", e))); + return Err(crate::error::AppError::Internal(format!( + "Failed to create checkpoint dir: {}", + e + ))); } // Copy tasks, graph, tech_debts, and notes to checkpoint folder let graph_json = state.read_graph(|g| serde_json::to_string(g).unwrap_or_default()); let _ = std::fs::write(target_dir.join("graph.json"), graph_json); - let tasks_json = state.project.tasks.read_with(|t| serde_json::to_string(t).unwrap_or_default()); + let tasks_json = state + .project + .tasks + .read_with(|t| serde_json::to_string(t).unwrap_or_default()); let _ = std::fs::write(target_dir.join("tasks.json"), tasks_json); - let debts_json = state.code.tech_debts.read_with(|d| serde_json::to_string(d).unwrap_or_default()); + let debts_json = state + .code + .tech_debts + .read_with(|d| serde_json::to_string(d).unwrap_or_default()); let _ = std::fs::write(target_dir.join("tech_debts.json"), debts_json); Ok(format!("Checkpoint '{}' saved successfully.", req.name)) @@ -1463,7 +1594,10 @@ impl McpTool for RestoreStateHandler { let target_dir = state.base_dir.join("checkpoints").join(&req.name); if !target_dir.exists() { - return Err(crate::error::AppError::Internal(format!("Checkpoint '{}' does not exist.", req.name))); + return Err(crate::error::AppError::Internal(format!( + "Checkpoint '{}' does not exist.", + req.name + ))); } if let Ok(graph_content) = std::fs::read_to_string(target_dir.join("graph.json")) { @@ -1534,35 +1668,40 @@ impl McpTool for QueryLineageHandler { "decision": adr.decision })); } - let changes = state.code.ledger.read_with(|c| c.clone()); - for change in changes { - let file_match = change.file_path.to_lowercase().contains(&q); - let desc_match = change.description.to_lowercase().contains(&q); - let repo_match = change.repo_name.as_ref().map(|r| r.to_lowercase().contains(&q)).unwrap_or(false); - let symbol_match = change.symbols.iter().any(|s| s.to_lowercase().contains(&q)); - let ns_match = change.namespace.to_lowercase().contains(&q); - if file_match || desc_match || repo_match || symbol_match || ns_match { - timeline.push(serde_json::json!({ - "timestamp": change.timestamp, - "type": "CodeChange", - "file": change.file_path, - "description": change.description, - "commit": change.git_commit, - "branch": change.git_branch, - "repo_name": change.repo_name, - "namespace": change.namespace, - "change_kind": change.change_kind, - "symbols": change.symbols, - "line_range": change.line_range, - "author": change.author, - })); + let changes = state.code.ledger.read_with(|c| c.clone()); + for change in changes { + let file_match = change.file_path.to_lowercase().contains(&q); + let desc_match = change.description.to_lowercase().contains(&q); + let repo_match = change + .repo_name + .as_ref() + .map(|r| r.to_lowercase().contains(&q)) + .unwrap_or(false); + let symbol_match = change.symbols.iter().any(|s| s.to_lowercase().contains(&q)); + let ns_match = change.namespace.to_lowercase().contains(&q); + if file_match || desc_match || repo_match || symbol_match || ns_match { + timeline.push(serde_json::json!({ + "timestamp": change.timestamp, + "type": "CodeChange", + "file": change.file_path, + "description": change.description, + "commit": change.git_commit, + "branch": change.git_branch, + "repo_name": change.repo_name, + "namespace": change.namespace, + "change_kind": change.change_kind, + "symbols": change.symbols, + "line_range": change.line_range, + "author": change.author, + })); + } } } - } let fixes = state.code.error_fixes.read_with(|f| f.clone()); for fix in fixes { - if fix.signature.to_lowercase().contains(&q) || fix.solution.to_lowercase().contains(&q) { + if fix.signature.to_lowercase().contains(&q) || fix.solution.to_lowercase().contains(&q) + { timeline.push(serde_json::json!({ "timestamp": fix.timestamp, "type": "ErrorFix", @@ -1601,20 +1740,21 @@ impl McpTool for GetNextActionableTasksHandler { } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: GetNextActionableTasksTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let req: GetNextActionableTasksTool = + serde_json::from_value(args).map_err(|e| e.to_string())?; let limit = req.limit.unwrap_or(5); let tasks = state.project.tasks.read_with(|t| t.clone()); let completed_ids: std::collections::HashSet = tasks .iter() - .filter(|t| t.status == "completed" || t.status == "done") + .filter(|t| !t.is_active()) .map(|t| t.id.clone()) .collect(); let mut actionable = Vec::new(); for task in tasks { - if task.status == "completed" || task.status == "done" || task.status == "cancelled" { + if !task.is_active() { continue; } @@ -1626,7 +1766,8 @@ impl McpTool for GetNextActionableTasksHandler { } } - let unblocked = task.dependencies.is_empty() || task.dependencies.iter().all(|d| completed_ids.contains(d)); + let unblocked = task.dependencies.is_empty() + || task.dependencies.iter().all(|d| completed_ids.contains(d)); if unblocked { actionable.push(task); @@ -1661,7 +1802,10 @@ impl McpTool for LogHypothesisHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: LogHypothesisTool = serde_json::from_value(args).map_err(|e| e.to_string())?; - let hyp_id = format!("HYP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase()); + let hyp_id = format!( + "HYP-{}", + uuid::Uuid::new_v4().to_string()[..8].to_uppercase() + ); let timestamp = now_secs(); let record = crate::models::Hypothesis { @@ -1710,7 +1854,9 @@ impl McpTool for QueryHypothesesHandler { if let Some(q) = &req.query { let lq = q.to_lowercase(); return h.hypothesis.to_lowercase().contains(&lq) - || h.evidence.as_ref().map_or(false, |e| e.to_lowercase().contains(&lq)); + || h.evidence + .as_ref() + .map_or(false, |e| e.to_lowercase().contains(&lq)); } true }) @@ -1729,32 +1875,40 @@ impl McpTool for GetPreflightContextHandler { } fn schema(&self) -> Value { - crate::mcp::tool_def::("get_preflight_context", "Get 1-page executive summary of active tasks, pinned files, hypotheses, and tech debt in 1 turn.") + crate::mcp::tool_def::( + "get_preflight_context", + "Get 1-page executive summary of active tasks, pinned files, hypotheses, and tech debt in 1 turn.", + ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: GetPreflightContextTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let req: GetPreflightContextTool = + serde_json::from_value(args).map_err(|e| e.to_string())?; let tasks = state.project.tasks.read_with(|t| t.clone()); let pinned = state.project.pinned_files.read_with(|p| p.clone()); let tech_debts = state.code.tech_debts.read_with(|d| d.clone()); let hypotheses = state.code.hypotheses.read_with(|h| h.clone()); - let recent_commands = state.telemetry.terminal_history.read_with(|h| { - h.iter().take(5).cloned().collect::>() - }); - let recent_activities = state.telemetry.recent_activities.read_with(|a| { - a.iter().take(5).cloned().collect::>() - }); + let recent_commands = state + .telemetry + .terminal_history + .read_with(|h| h.iter().take(5).cloned().collect::>()); + let recent_activities = state + .telemetry + .recent_activities + .read_with(|a| a.iter().take(5).cloned().collect::>()); let active_tasks: Vec<_> = tasks .into_iter() - .filter(|t| !matches!(t.status.as_str(), "completed" | "done" | "cancelled")) - .map(|t| serde_json::json!({ - "id": t.id, - "title": t.title, - "status": t.status, - "criteria": t.acceptance_criteria - })) + .filter(|t| t.is_active()) + .map(|t| { + serde_json::json!({ + "id": t.id, + "title": t.title, + "status": t.status, + "criteria": t.acceptance_criteria + }) + }) .collect(); let branch_pinned: Vec<_> = pinned @@ -1772,11 +1926,13 @@ impl McpTool for GetPreflightContextHandler { .into_iter() .filter(|d| d.namespace == req.namespace && !d.is_resolved) .take(5) - .map(|d| serde_json::json!({ - "id": d.id, - "description": d.description, - "ideal_solution": d.ideal_solution - })) + .map(|d| { + serde_json::json!({ + "id": d.id, + "description": d.description, + "ideal_solution": d.ideal_solution + }) + }) .collect(); let active_hypotheses: Vec<_> = hypotheses @@ -1809,11 +1965,15 @@ impl McpTool for BroadcastAgentSignalHandler { } fn schema(&self) -> Value { - crate::mcp::tool_def::("broadcast_agent_signal", "Broadcast real-time inter-agent signal to peer subagents.") + crate::mcp::tool_def::( + "broadcast_agent_signal", + "Broadcast real-time inter-agent signal to peer subagents.", + ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: BroadcastAgentSignalTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let req: BroadcastAgentSignalTool = + serde_json::from_value(args).map_err(|e| e.to_string())?; let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -1844,9 +2004,16 @@ impl McpTool for BroadcastAgentSignalHandler { s.remove(0); } }); - state.record_activity("agent_signal", &format!("{}: {}", req.sender, req.signal_type), None); + state.record_activity( + "agent_signal", + &format!("{}: {}", req.sender, req.signal_type), + None, + ); - Ok(format!("Broadcasted signal '{}' from agent '{}'.", sig_id, req.sender)) + Ok(format!( + "Broadcasted signal '{}' from agent '{}'.", + sig_id, req.sender + )) } } @@ -1859,7 +2026,10 @@ impl McpTool for QueryAgentSignalsHandler { } fn schema(&self) -> Value { - crate::mcp::tool_def::("query_agent_signals", "Query active inter-agent signals from subagent signal bus.") + crate::mcp::tool_def::( + "query_agent_signals", + "Query active inter-agent signals from subagent signal bus.", + ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { @@ -1909,11 +2079,15 @@ impl McpTool for AutoSessionCheckpointHandler { } fn schema(&self) -> Value { - crate::mcp::tool_def::("auto_session_checkpoint", "Trigger automated context checkpoint into a permanent HandoffMemo.") + crate::mcp::tool_def::( + "auto_session_checkpoint", + "Trigger automated context checkpoint into a permanent HandoffMemo.", + ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: AutoSessionCheckpointTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let req: AutoSessionCheckpointTool = + serde_json::from_value(args).map_err(|e| e.to_string())?; let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -1926,11 +2100,20 @@ impl McpTool for AutoSessionCheckpointHandler { let active_tasks: Vec<_> = tasks .iter() - .filter(|t| !matches!(t.status.as_str(), "completed" | "done" | "cancelled")) + .filter(|t| t.is_active()) .map(|t| t.title.as_str()) .collect(); - let unverified_hyp: Vec<_> = hypotheses.iter().filter(|h| h.status == "unverified").map(|h| h.hypothesis.as_str()).collect(); - let recent_changes: Vec<_> = ledger.iter().rev().take(5).map(|c| c.file_path.as_str()).collect(); + let unverified_hyp: Vec<_> = hypotheses + .iter() + .filter(|h| h.status == "unverified") + .map(|h| h.hypothesis.as_str()) + .collect(); + let recent_changes: Vec<_> = ledger + .iter() + .rev() + .take(5) + .map(|c| c.file_path.as_str()) + .collect(); let author = req.author.unwrap_or_else(|| "AutoCheckpoint".to_string()); let memo_id = format!("memo_chk_{}", timestamp); @@ -1955,9 +2138,16 @@ impl McpTool for AutoSessionCheckpointHandler { m.remove(0); } }); - state.record_activity("checkpoint", &format!("Created auto session checkpoint {}", memo_id), None); + state.record_activity( + "checkpoint", + &format!("Created auto session checkpoint {}", memo_id), + None, + ); - Ok(format!("Session checkpoint created with memo ID '{}'.", memo_id)) + Ok(format!( + "Session checkpoint created with memo ID '{}'.", + memo_id + )) } } @@ -1970,19 +2160,29 @@ impl McpTool for SearchSnippetsHybridHandler { } fn schema(&self) -> Value { - crate::mcp::tool_def::("search_snippets_hybrid", "Hybrid BM25 and semantic tag search across stored code snippets.") + crate::mcp::tool_def::( + "search_snippets_hybrid", + "Hybrid BM25 and semantic tag search across stored code snippets.", + ) } async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { - let req: SearchSnippetsHybridTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let req: SearchSnippetsHybridTool = + serde_json::from_value(args).map_err(|e| e.to_string())?; let snippets = state.code.snippets.read_with(|s| s.clone()); - let q_tokens: Vec<_> = req.query.to_lowercase().split_whitespace().map(|s| s.to_string()).collect(); + let q_tokens: Vec<_> = req + .query + .to_lowercase() + .split_whitespace() + .map(|s| s.to_string()) + .collect(); let mut scored: Vec<_> = snippets .into_iter() .filter_map(|s| { - let text = format!("{} {} {} {}", s.name, s.description, s.language, s.code).to_lowercase(); + let text = format!("{} {} {} {}", s.name, s.description, s.language, s.code) + .to_lowercase(); let mut score = 0.0; for tok in &q_tokens { if text.contains(tok) { @@ -1991,7 +2191,10 @@ impl McpTool for SearchSnippetsHybridHandler { } if let Some(tags) = &req.tags { for tag in tags { - if s.tags.iter().any(|t| t.to_lowercase() == tag.to_lowercase()) { + if s.tags + .iter() + .any(|t| t.to_lowercase() == tag.to_lowercase()) + { score += 2.0; } } @@ -2012,7 +2215,10 @@ impl McpTool for SearchSnippetsHybridHandler { .collect(); scored.sort_by(|a, b| { - b["score"].as_f64().unwrap_or(0.0).total_cmp(&a["score"].as_f64().unwrap_or(0.0)) + b["score"] + .as_f64() + .unwrap_or(0.0) + .total_cmp(&a["score"].as_f64().unwrap_or(0.0)) }); let limit = req.limit.unwrap_or(10); @@ -2287,15 +2493,13 @@ mod tests { #[tokio::test] async fn test_all_meta_handlers_comprehensive() { - use crate::handlers::tasks::{MilestonesHandler, TasksHandler}; - use crate::handlers::graph::SweepGraphHealthHandler; use crate::handlers::git::QueryGitDiffsHandler; - + use crate::handlers::graph::SweepGraphHealthHandler; + use crate::handlers::tasks::{MilestonesHandler, TasksHandler}; let dir = tempfile::tempdir().unwrap(); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); - // LogDecision & DeleteDecision & QueryDecisions let log_dec = LogDecisionHandler; let dec_res = log_dec @@ -2311,11 +2515,13 @@ mod tests { .await .unwrap(); - assert!(dec_res.contains("Logged decision")); let q_dec = QueryDecisionsHandler; - let q_dec_res = q_dec.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let q_dec_res = q_dec + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(q_dec_res.contains("Use Axum")); let del_dec = DeleteDecisionHandler; @@ -2339,10 +2545,12 @@ mod tests { assert!(td_res.contains("Tech debt logged")); let list_td = ListTechDebtHandler; - let list_td_res = list_td.execute(serde_json::json!({"include_resolved": true}), state.clone()).await.unwrap(); + let list_td_res = list_td + .execute(serde_json::json!({"include_resolved": true}), state.clone()) + .await + .unwrap(); assert!(list_td_res.contains("Replace unwraps")); - let debt_id = state.code.tech_debts.read_with(|debts| debts[0].id.clone()); let res_td = ResolveTechDebtHandler; let res_td_res = res_td @@ -2350,7 +2558,6 @@ mod tests { .await; assert!(res_td_res.is_ok()); - // LogHypothesis & QueryHypotheses let log_hyp = LogHypothesisHandler; let hyp_res = log_hyp @@ -2365,10 +2572,11 @@ mod tests { .unwrap(); assert!(hyp_res.contains("Hypothesis")); - - let q_hyp = QueryHypothesesHandler; - let q_hyp_res = q_hyp.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let q_hyp_res = q_hyp + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(q_hyp_res.contains("Caching improves response speed")); // AddMilestone & UpdateMilestone & ListMilestones @@ -2399,7 +2607,10 @@ mod tests { .await; assert!(upd_ms_res.is_ok()); - let list_ms_res = handler_ms.execute(serde_json::json!({"action": "list"}), state.clone()).await.unwrap(); + let list_ms_res = handler_ms + .execute(serde_json::json!({"action": "list"}), state.clone()) + .await + .unwrap(); assert!(list_ms_res.contains("v1.0 Release")); // SetAcceptanceCriteria & VerifyAcceptanceCriteria @@ -2445,21 +2656,27 @@ mod tests { .await; assert!(ver_ac_res.is_ok()); - // GetProjectHealth & SweepGraphHealth let proj_h = GetProjectHealthHandler; - let proj_h_res = proj_h.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let proj_h_res = proj_h + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(proj_h_res.contains("active_tasks")); - let sweep_h = SweepGraphHealthHandler; - let sweep_h_res = sweep_h.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let sweep_h_res = sweep_h + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(sweep_h_res.contains("orphaned_entities")); - // AutoSessionCheckpoint let chk = AutoSessionCheckpointHandler; - let chk_res = chk.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let chk_res = chk + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(chk_res.contains("Session checkpoint created")); // QueryGitDiffs & QueryLineage @@ -2470,8 +2687,6 @@ mod tests { .unwrap(); assert!(!q_diffs_res.is_empty()); - - let q_lin = QueryLineageHandler; let q_lin_res = q_lin .execute(serde_json::json!({"query": "test_sym"}), state.clone()) @@ -2494,9 +2709,11 @@ mod tests { .unwrap(); assert!(bcast_res.contains("Broadcasted signal")); - let q_sig = QueryAgentSignalsHandler; - let q_sig_res = q_sig.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let q_sig_res = q_sig + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(q_sig_res.contains("Agent1")); // Error Fixes @@ -2513,7 +2730,6 @@ mod tests { .unwrap(); assert!(log_ef_res.contains("Logged error fix for")); - let search_ef = SearchErrorFixesHandler; let search_ef_res = search_ef .execute(serde_json::json!({"query": "E0425"}), state.clone()) @@ -2531,27 +2747,42 @@ mod tests { // Preferences & Recent changes let learn_p = LearnPreferenceHandler; let learn_p_res = learn_p - .execute(serde_json::json!({"key": "theme", "value": "dark"}), state.clone()) + .execute( + serde_json::json!({"key": "theme", "value": "dark"}), + state.clone(), + ) .await .unwrap(); assert_eq!(learn_p_res, "Preference learned"); let read_p = ReadPreferencesHandler; - let read_p_res = read_p.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let read_p_res = read_p + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(read_p_res.contains("dark")); let q_rec = QueryRecentChangesHandler; - let q_rec_res = q_rec.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let q_rec_res = q_rec + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(!q_rec_res.is_empty()); // GetNextActionableTasks let get_next = GetNextActionableTasksHandler; - let get_next_res = get_next.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let get_next_res = get_next + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(get_next_res.contains("actionable_count")); // GetPreflightContext let preflight = GetPreflightContextHandler; - let preflight_res = preflight.execute(serde_json::json!({}), state.clone()).await.unwrap(); + let preflight_res = preflight + .execute(serde_json::json!({}), state.clone()) + .await + .unwrap(); assert!(preflight_res.contains("active_tasks")); // ManageCheckpoint / CheckpointState / RestoreState @@ -2577,7 +2808,10 @@ mod tests { assert!(mg_list.contains("test_chk")); let mg_del = mg_chk - .execute(serde_json::json!({"action": "delete", "name_or_id": "test_chk"}), state.clone()) + .execute( + serde_json::json!({"action": "delete", "name_or_id": "test_chk"}), + state.clone(), + ) .await .unwrap(); assert!(mg_del.contains("deleted successfully")); @@ -2598,7 +2832,10 @@ mod tests { let search_hyb = SearchSnippetsHybridHandler; let search_hyb_res = search_hyb - .execute(serde_json::json!({"query": "rust mcp", "tags": ["rust"]}), state.clone()) + .execute( + serde_json::json!({"query": "rust mcp", "tags": ["rust"]}), + state.clone(), + ) .await .unwrap(); assert!(search_hyb_res.contains("Rust MCP Helper")); @@ -2674,7 +2911,10 @@ mod tests { }); let get_next_res2 = get_next - .execute(serde_json::json!({"git_branch": "main", "limit": 2}), state.clone()) + .execute( + serde_json::json!({"git_branch": "main", "limit": 2}), + state.clone(), + ) .await .unwrap(); assert!(get_next_res2.contains("t-2")); @@ -2728,7 +2968,10 @@ mod tests { // SearchErrorFixes with include_body true and false, query None let search_ef_full = search_ef - .execute(serde_json::json!({"include_body": true, "limit": 2}), state.clone()) + .execute( + serde_json::json!({"include_body": true, "limit": 2}), + state.clone(), + ) .await .unwrap(); assert!(search_ef_full.contains("E0425")); @@ -2788,22 +3031,28 @@ mod tests { // Entity graph for OmniSearch GraphRAG expansion state.graph.modify(|g| { - g.entities.insert("Ent1".to_string(), crate::models::Entity { - name: "Ent1".to_string(), - entity_type: "Module".to_string(), - observations: vec!["Obs 1".to_string()], - namespace: "global".to_string(), - git_branch: None, - ..Default::default() - }); - g.entities.insert("Ent2".to_string(), crate::models::Entity { - name: "Ent2".to_string(), - entity_type: "Class".to_string(), - observations: vec!["Obs 2".to_string()], - namespace: "global".to_string(), - git_branch: None, - ..Default::default() - }); + g.entities.insert( + "Ent1".to_string(), + crate::models::Entity { + name: "Ent1".to_string(), + entity_type: "Module".to_string(), + observations: vec!["Obs 1".to_string()], + namespace: "global".to_string(), + git_branch: None, + ..Default::default() + }, + ); + g.entities.insert( + "Ent2".to_string(), + crate::models::Entity { + name: "Ent2".to_string(), + entity_type: "Class".to_string(), + observations: vec!["Obs 2".to_string()], + namespace: "global".to_string(), + git_branch: None, + ..Default::default() + }, + ); g.relations.push(crate::models::Relation { from: "Ent1".to_string(), to: "Ent2".to_string(), @@ -2829,7 +3078,10 @@ mod tests { // RestoreState non-existent error let rest_err = rest_state - .execute(serde_json::json!({"name": "non_existent_chk"}), state.clone()) + .execute( + serde_json::json!({"name": "non_existent_chk"}), + state.clone(), + ) .await; assert!(rest_err.is_err()); @@ -2856,7 +3108,6 @@ mod tests { .await .unwrap(); - let q_sig = QueryAgentSignalsHandler; let q_sig_res = q_sig .execute( @@ -2871,10 +3122,33 @@ mod tests { assert!(q_sig_res.contains("AgentA")); // ManageCheckpoint error cases - assert!(mg_chk.execute(serde_json::json!({"action": "create"}), state.clone()).await.is_err()); - assert!(mg_chk.execute(serde_json::json!({"action": "restore"}), state.clone()).await.is_err()); - assert!(mg_chk.execute(serde_json::json!({"action": "delete"}), state.clone()).await.is_err()); - assert!(mg_chk.execute(serde_json::json!({"action": "restore", "name_or_id": "non_existent"}), state.clone()).await.is_err()); + assert!( + mg_chk + .execute(serde_json::json!({"action": "create"}), state.clone()) + .await + .is_err() + ); + assert!( + mg_chk + .execute(serde_json::json!({"action": "restore"}), state.clone()) + .await + .is_err() + ); + assert!( + mg_chk + .execute(serde_json::json!({"action": "delete"}), state.clone()) + .await + .is_err() + ); + assert!( + mg_chk + .execute( + serde_json::json!({"action": "restore", "name_or_id": "non_existent"}), + state.clone() + ) + .await + .is_err() + ); // Restore snapshot by ID state.project.snapshots.modify(|snaps| { @@ -2887,7 +3161,10 @@ mod tests { }); }); let rest_snap = mg_chk - .execute(serde_json::json!({"action": "restore", "name_or_id": "SNAP-12345678"}), state.clone()) + .execute( + serde_json::json!({"action": "restore", "name_or_id": "SNAP-12345678"}), + state.clone(), + ) .await .unwrap(); assert!(rest_snap.contains("restored memory state from snapshot")); @@ -2930,13 +3207,14 @@ mod tests { expires_at: None, ..Default::default() }); - - }); let next_act = GetNextActionableTasksHandler; let next_act_res = next_act - .execute(serde_json::json!({"git_branch": "feature/test"}), state.clone()) + .execute( + serde_json::json!({"git_branch": "feature/test"}), + state.clone(), + ) .await .unwrap(); assert!(next_act_res.contains("actionable_count")); @@ -2954,21 +3232,12 @@ mod tests { let preflight = GetPreflightContextHandler; let pre_res = preflight - .execute(serde_json::json!({"namespace": "global", "git_branch": "main"}), state.clone()) + .execute( + serde_json::json!({"namespace": "global", "git_branch": "main"}), + state.clone(), + ) .await .unwrap(); assert!(pre_res.contains("src/lib.rs")); } } - - - - - - - - - - - - diff --git a/server/src/handlers/tasks.rs b/server/src/handlers/tasks.rs index 2d55565..f2387be 100644 --- a/server/src/handlers/tasks.rs +++ b/server/src/handlers/tasks.rs @@ -269,7 +269,7 @@ impl McpTool for TasksHandler { let status_match = match req.status.as_deref() { Some("all") => true, Some(s) => t.status.eq_ignore_ascii_case(s), - None => t.status != "done" && t.status != "completed", + None => t.is_active(), }; let branch_match = match &req.git_branch { Some(branch) => t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str()), diff --git a/server/src/lib.rs b/server/src/lib.rs index 03a87d9..c7b2ac3 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -139,67 +139,91 @@ pub async fn condense_graph_worker(state: Arc) { .unwrap_or_default() .as_secs(); - let mut condensed_sticky_content = String::new(); - state.code.sticky.modify(|notes| { + let sticky_condensation = state.code.sticky.read_with(|notes| { if notes.len() > threshold { - notes.sort_by_key(|n| n.timestamp); - let to_remove = notes.len() - (threshold / 2); - let removed: Vec<_> = notes.drain(0..to_remove).collect(); - for r in removed { - condensed_sticky_content.push_str(&format!("{}\n", r.content)); + let mut sorted = notes.clone(); + sorted.sort_by_key(|n| n.timestamp); + let to_remove = sorted.len() - (threshold / 2); + let removed: Vec<_> = sorted.into_iter().take(to_remove).collect(); + let mut content = String::new(); + let mut ids = Vec::new(); + for r in &removed { + content.push_str(&format!("{}\n", r.content)); + ids.push(r.id.clone()); } + Some((content, ids)) + } else { + None } }); - if !condensed_sticky_content.is_empty() { - state.modify_graph(|graph| { + if let Some((content, ids)) = sticky_condensation { + if !content.is_empty() { let name = format!("StickyNote History {}", now); - graph.entities.insert( - name.clone(), - crate::models::Entity { - name: name.clone(), - entity_type: "Historical Summary".to_string(), - observations: vec![condensed_sticky_content], - namespace: crate::models::default_namespace(), - git_branch: None, - ..Default::default() - }, - ); - }); - tracing::info!("Condensed sticky notes into Historical Summary."); + state.modify_graph(|graph| { + graph.entities.insert( + name.clone(), + crate::models::Entity { + name: name.clone(), + entity_type: "Historical Summary".to_string(), + observations: vec![content], + namespace: crate::models::default_namespace(), + git_branch: None, + ..Default::default() + }, + ); + }); + let id_set: std::collections::HashSet = ids.into_iter().collect(); + state.code.sticky.modify(|notes| { + notes.retain(|n| !id_set.contains(&n.id)); + }); + tracing::info!("Condensed sticky notes into Historical Summary."); + } } - let mut condensed_snippet_content = String::new(); - state.code.snippets.modify(|snippets| { + let snippet_condensation = state.code.snippets.read_with(|snippets| { if snippets.len() > threshold { - snippets.sort_by_key(|s| s.updated_at); - let to_remove = snippets.len() - (threshold / 2); - let removed: Vec<_> = snippets.drain(0..to_remove).collect(); - for r in removed { - condensed_snippet_content.push_str(&format!( + let mut sorted = snippets.clone(); + sorted.sort_by_key(|s| s.updated_at); + let to_remove = sorted.len() - (threshold / 2); + let removed: Vec<_> = sorted.into_iter().take(to_remove).collect(); + let mut content = String::new(); + let mut names = Vec::new(); + for r in &removed { + content.push_str(&format!( "Name: {}\nDesc: {}\nCode: {}\n", r.name, r.description, r.code )); + names.push(r.name.clone()); } + Some((content, names)) + } else { + None } }); - if !condensed_snippet_content.is_empty() { - state.modify_graph(|graph| { + if let Some((content, names)) = snippet_condensation { + if !content.is_empty() { let name = format!("Snippet History {}", now); - graph.entities.insert( - name.clone(), - crate::models::Entity { - name: name.clone(), - entity_type: "Historical Summary".to_string(), - observations: vec![condensed_snippet_content], - namespace: crate::models::default_namespace(), - git_branch: None, - ..Default::default() - }, - ); - }); - tracing::info!("Condensed snippets into Historical Summary."); + state.modify_graph(|graph| { + graph.entities.insert( + name.clone(), + crate::models::Entity { + name: name.clone(), + entity_type: "Historical Summary".to_string(), + observations: vec![content], + namespace: crate::models::default_namespace(), + git_branch: None, + ..Default::default() + }, + ); + }); + let name_set: std::collections::HashSet = names.into_iter().collect(); + state.code.snippets.modify(|snippets| { + snippets.retain(|s| !name_set.contains(&s.name)); + }); + tracing::info!("Condensed snippets into Historical Summary."); + } } } } diff --git a/server/src/state.rs b/server/src/state.rs index 8952e46..2d019b3 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -513,12 +513,12 @@ impl SearchService { let mut uncached_meta = Vec::new(); self.state.code.snippets.read_with(|snips| { - for snippet in snips.iter().take(50) { + for snippet in snips.iter() { let title = snippet.name.clone(); let desc = snippet.description.clone(); if let Some(ref emb) = snippet.embedding { cached_items.push((title, "snippet".to_string(), desc, emb.clone())); - } else { + } else if uncached_texts.len() < 50 { uncached_texts.push(format!( "{} {} {}", snippet.name, snippet.description, snippet.code @@ -529,7 +529,10 @@ impl SearchService { }); self.state.code.sticky.read_with(|sticky| { - for note in sticky.iter().take(50) { + for note in sticky.iter() { + if uncached_texts.len() >= 50 { + break; + } let content_preview = note.content.chars().take(200).collect::(); uncached_texts.push(note.content.clone()); uncached_meta.push(( @@ -540,7 +543,7 @@ impl SearchService { } }); self.state.read_graph(|graph| { - for entity in graph.entities.values().take(50) { + for entity in graph.entities.values() { if let Some(ns) = filter_namespace { if entity.namespace != ns { continue; @@ -551,7 +554,7 @@ impl SearchService { let desc = format!("{}: {}", entity.entity_type, obs); if let Some(ref emb) = entity.embedding { cached_items.push((title, "entity".to_string(), desc, emb.clone())); - } else { + } else if uncached_texts.len() < 50 { uncached_texts.push(format!("{} {} {}", entity.name, entity.entity_type, obs)); uncached_meta.push((title, "entity".to_string(), desc)); } @@ -559,12 +562,12 @@ impl SearchService { }); self.state.code.error_fixes.read_with(|fixes| { - for fix in fixes.iter().take(50) { + for fix in fixes.iter() { let title = fix.signature.clone(); let desc = fix.solution.clone(); if let Some(ref emb) = fix.embedding { cached_items.push((title, "error_fix".to_string(), desc, emb.clone())); - } else { + } else if uncached_texts.len() < 50 { uncached_texts.push(format!("{} {}", fix.signature, fix.solution)); uncached_meta.push((title, "error_fix".to_string(), desc)); } diff --git a/server/src/store.rs b/server/src/store.rs index d48c84a..ba03afa 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -8,6 +8,10 @@ pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("stor enum DbOp { Insert(Vec), Delete, + Batch { + inserts: Vec<(String, Vec)>, + deletes: Vec, + }, } /// Internal write request dispatched to the single database writer actor. @@ -83,6 +87,28 @@ impl DbWriteQueue { ); } } + DbOp::Batch { inserts, deletes } => { + for del_k in deletes { + if let Err(e) = table.remove(del_k.as_str()) { + tracing::error!( + "Failed to delete batch key '{}' from redb: {}", + del_k, + e + ); + } + } + for (ins_k, ins_bytes) in inserts { + if let Err(e) = + table.insert(ins_k.as_str(), ins_bytes.as_slice()) + { + tracing::error!( + "Failed to insert batch key '{}' into redb: {}", + ins_k, + e + ); + } + } + } } } } @@ -130,6 +156,16 @@ impl DbWriteQueue { self.push_op(key, DbOp::Delete, flushed_notifier) } + pub fn push_batch( + &self, + key: String, + inserts: Vec<(String, Vec)>, + deletes: Vec, + flushed_notifier: Arc, + ) -> Option> { + self.push_op(key, DbOp::Batch { inserts, deletes }, flushed_notifier) + } + fn push_op( &self, key: String, @@ -191,6 +227,17 @@ impl DbWriteQueue { .await } + pub async fn push_batch_async( + &self, + key: String, + inserts: Vec<(String, Vec)>, + deletes: Vec, + flushed_notifier: Arc, + ) -> Option> { + self.push_op_async(key, DbOp::Batch { inserts, deletes }, flushed_notifier) + .await + } + async fn push_op_async( &self, key: String, @@ -222,14 +269,16 @@ pub struct Store { key: String, queue: DbWriteQueue, is_corrupted: bool, + known_granular_keys: Arc>>, } impl Store { pub fn new(key: &str, db: Arc) -> Self { - let (initial_data, is_corrupted) = Self::load_from_db(key, &db); + let (initial_data, is_corrupted, known_keys) = Self::load_from_db(key, &db); let cache = Arc::new(RwLock::new(initial_data)); let flushed = Arc::new(tokio::sync::Notify::new()); let queue = get_or_create_queue(db); + let known_granular_keys = Arc::new(RwLock::new(known_keys)); Self { cache, @@ -237,29 +286,38 @@ impl Store key: key.to_string(), queue, is_corrupted, + known_granular_keys, } } - fn load_from_db(key: &str, db: &Database) -> (T, bool) { + fn load_from_db(key: &str, db: &Database) -> (T, bool, std::collections::HashSet) { let Ok(read_txn) = db.begin_read() else { tracing::error!("Failed to begin read transaction for key '{}'", key); - return (T::default(), false); + return (T::default(), false, std::collections::HashSet::new()); }; let Ok(table) = read_txn.open_table(STORE_TABLE) else { - return (T::default(), false); + return (T::default(), false, std::collections::HashSet::new()); }; // 1. Check monolithic key first as the authoritative snapshot match table.get(key) { Ok(Some(value)) => match serde_json::from_slice::(value.value()) { - Ok(parsed) => return (parsed, false), + Ok(parsed) => { + let mut known = std::collections::HashSet::new(); + if let Ok(val) = serde_json::to_value(&parsed) { + for k in Self::extract_granular_keys(key, &val) { + known.insert(k); + } + } + return (parsed, false, known); + } Err(e) => { tracing::error!( "CRITICAL: Corrupted data for key '{}' in database: {}. Quarantine mode active: state initialized to empty default without overwriting DB key.", key, e ); - return (T::default(), true); + return (T::default(), true, std::collections::HashSet::new()); } }, Ok(None) => {} @@ -273,6 +331,7 @@ impl Store let mut items_array = Vec::new(); let mut items_map = serde_json::Map::new(); let mut found_granular = false; + let mut known = std::collections::HashSet::new(); if let Ok(range) = table.range(prefix.as_str()..) { for entry in range { @@ -282,6 +341,7 @@ impl Store break; } found_granular = true; + known.insert(k_str.to_string()); if let Ok(val) = serde_json::from_slice::(v.value()) { let sub_key = &k_str[prefix.len()..]; items_array.push(val.clone()); @@ -293,14 +353,14 @@ impl Store if found_granular { if let Ok(parsed) = serde_json::from_value::(serde_json::Value::Array(items_array)) { - return (parsed, false); + return (parsed, false, known); } if let Ok(parsed) = serde_json::from_value::(serde_json::Value::Object(items_map)) { - return (parsed, false); + return (parsed, false, known); } } - (T::default(), false) + (T::default(), false, std::collections::HashSet::new()) } fn extract_granular_keys(base_key: &str, val: &serde_json::Value) -> Vec { @@ -377,49 +437,43 @@ impl Store return; } - // Fast mutation under critical lock section, then immediately release the RwLock guard - let (old_snapshot, new_snapshot) = { + // Fast mutation under critical lock section, only ONE clone taken, then immediately release the RwLock guard + let new_snapshot = { let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner()); - let old = (*lock).clone(); f(&mut lock); - let new = (*lock).clone(); - (old, new) + (*lock).clone() }; // Expensive serialization and granular extraction run completely unblocked outside the lock - let old_keys = serde_json::to_value(&old_snapshot) - .map(|val| Self::extract_granular_keys(&self.key, &val)) - .unwrap_or_default(); - let full_bytes_res = serde_json::to_vec(&new_snapshot); let granular_entries = serde_json::to_value(&new_snapshot) .map(|val| Self::extract_granular_entries(&self.key, &val)) .unwrap_or_default(); - let new_keys: std::collections::HashSet<&str> = - granular_entries.iter().map(|(k, _)| k.as_str()).collect(); + let new_keys: std::collections::HashSet = + granular_entries.iter().map(|(k, _)| k.clone()).collect(); let mut removed_keys = Vec::new(); - for old_k in &old_keys { - if !new_keys.contains(old_k.as_str()) { - removed_keys.push(old_k.clone()); + { + let mut known = self.known_granular_keys.write().unwrap_or_else(|e| e.into_inner()); + for old_k in known.iter() { + if !new_keys.contains(old_k) { + removed_keys.push(old_k.clone()); + } } + *known = new_keys; } match full_bytes_res { Ok(data) => { - // Delete removed granular entries so they don't resurrect on restart - for del_key in removed_keys { - self.queue.push_delete(del_key, self.flushed.clone()); - } - - // Queue granular entries + let mut batch_inserts = Vec::with_capacity(granular_entries.len() + 1); for (g_key, g_bytes) in granular_entries { - self.queue.push(g_key, g_bytes, self.flushed.clone()); + batch_inserts.push((g_key, g_bytes)); } + batch_inserts.push((self.key.clone(), data.clone())); if self .queue - .push(self.key.clone(), data.clone(), self.flushed.clone()) + .push_batch(self.key.clone(), batch_inserts.clone(), removed_keys.clone(), self.flushed.clone()) .is_none() { tracing::warn!( @@ -433,7 +487,7 @@ impl Store handle.spawn(async move { let _ = tokio::time::timeout( std::time::Duration::from_secs(10), - queue.push_async(key, data, flushed), + queue.push_batch_async(key, batch_inserts, removed_keys, flushed), ) .await; }); @@ -460,47 +514,41 @@ impl Store return; } - let (old_snapshot, new_snapshot) = { + let new_snapshot = { let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner()); - let old = (*lock).clone(); f(&mut lock); - let new = (*lock).clone(); - (old, new) + (*lock).clone() }; - let old_keys = serde_json::to_value(&old_snapshot) - .map(|val| Self::extract_granular_keys(&self.key, &val)) - .unwrap_or_default(); - let full_bytes_res = serde_json::to_vec(&new_snapshot); let granular_entries = serde_json::to_value(&new_snapshot) .map(|val| Self::extract_granular_entries(&self.key, &val)) .unwrap_or_default(); - let new_keys: std::collections::HashSet<&str> = - granular_entries.iter().map(|(k, _)| k.as_str()).collect(); + let new_keys: std::collections::HashSet = + granular_entries.iter().map(|(k, _)| k.clone()).collect(); let mut removed_keys = Vec::new(); - for old_k in &old_keys { - if !new_keys.contains(old_k.as_str()) { - removed_keys.push(old_k.clone()); + { + let mut known = self.known_granular_keys.write().unwrap_or_else(|e| e.into_inner()); + for old_k in known.iter() { + if !new_keys.contains(old_k) { + removed_keys.push(old_k.clone()); + } } + *known = new_keys; } match full_bytes_res { Ok(data) => { - for del_key in removed_keys { - self.queue - .push_delete_async(del_key, self.flushed.clone()) - .await; - } + let mut batch_inserts = Vec::with_capacity(granular_entries.len() + 1); for (g_key, g_bytes) in granular_entries { - self.queue - .push_async(g_key, g_bytes, self.flushed.clone()) - .await; + batch_inserts.push((g_key, g_bytes)); } + batch_inserts.push((self.key.clone(), data)); + if let Some(rx) = self .queue - .push_async(self.key.clone(), data, self.flushed.clone()) + .push_batch_async(self.key.clone(), batch_inserts, removed_keys, self.flushed.clone()) .await { let _ = rx.await;