From 9854556cd8051022590c9e0451ef459164df775d Mon Sep 17 00:00:00 2001 From: David Hagerty Date: Fri, 13 Feb 2026 14:35:58 +0000 Subject: [PATCH] feat(tools): register CodeSearchTool in V2 registry Implements Task 2 from Phase 5: Add CodeSearchTool to create_v2_registry with project_root parameter. Update all call sites (orchestrator and test files) to pass project_root path. --- tests/agent_tools_test.rs | 7 ++++++- tests/graph_tools_test.rs | 10 +++++++--- tests/orchestrator_test.rs | 135 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/agent/orchestrator.rs | 1 + src/tools/factory.rs | 6 ++++++ 5 file(s) changed, 155 insertion(s)(+), 4 deletion(s)(-) diff --git a/tests/agent_tools_test.rs b/tests/agent_tools_test.rs --- a/tests/agent_tools_test.rs +++ b/tests/agent_tools_test.rs @@ -448,12 +448,15 @@ let validator = Arc::new(SecurityValidator::new(security_config).unwrap()); let permission_handler = Arc::new(AutoApproveHandler); + let project_root = std::path::PathBuf::from("/tmp"); + let registry = create_v2_registry( validator, permission_handler, graph_store, Some(message_bus), Some("worker-1".to_string()), + project_root, ); let tools = registry.list(); @@ -483,7 +486,9 @@ let validator = Arc::new(SecurityValidator::new(security_config).unwrap()); let permission_handler = Arc::new(AutoApproveHandler); - let registry = create_v2_registry(validator, permission_handler, graph_store, None, None); + let project_root = std::path::PathBuf::from("/tmp"); + + let registry = create_v2_registry(validator, permission_handler, graph_store, None, None, project_root); let tools = registry.list(); assert!(!tools.contains(&"spawn_sub_agent".to_string())); diff --git a/tests/graph_tools_test.rs b/tests/graph_tools_test.rs --- a/tests/graph_tools_test.rs +++ b/tests/graph_tools_test.rs @@ -662,10 +662,12 @@ Arc::new(SecurityValidator::new(security_config).expect("Failed to create validator")); let permission_handler = Arc::new(AutoApproveHandler); - // Create the v2 registry - let registry = create_v2_registry(validator, permission_handler, graph_store, None, None); + let project_root = std::path::PathBuf::from("/tmp"); - // Expected tool names: graph tools + legacy tools + context tools + // Create the v2 registry + let registry = create_v2_registry(validator, permission_handler, graph_store, None, None, project_root); + + // Expected tool names: graph tools + legacy tools + context tools + search tools let expected_tools = vec![ // Graph tools "create_node", @@ -687,6 +689,8 @@ "signal_completion", // Context tools "read_agents_md", + // Search tools + "code_search", ]; // Get all registered tool names diff --git a/tests/orchestrator_test.rs b/tests/orchestrator_test.rs --- a/tests/orchestrator_test.rs +++ b/tests/orchestrator_test.rs @@ -630,3 +630,138 @@ let next_state = orchestrator.handle_scheduling().await.unwrap(); assert_eq!(next_state, OrchestratorState::Completing); } + +// ===== Phase 5 Error Recovery & Task Reassignment Tests ===== + +/// Integration test: Full retry-cascade-unblock lifecycle +/// Verifies v2-phase5.AC11.1, v2-phase5.AC12.1, v2-phase5.AC13.1 +#[tokio::test] +async fn test_retry_cascade_unblock_lifecycle() { + let (_, graph_store) = common::setup_test_env().await.unwrap(); + let graph_store: Arc = Arc::new(graph_store); + let mock_client = Arc::new(MockLlmClient::new()); + + let config = OrchestratorConfig { + max_retries_per_task: 2, + ..Default::default() + }; + let mut orchestrator = create_test_orchestrator(config, graph_store.clone(), mock_client); + orchestrator.set_goal_id(Some("ra-p5-test".to_string())); + + // Create goal node + let goal_node = common::create_test_goal("ra-p5-test", "proj-1", "Phase 5 test goal"); + graph_store.create_node(&goal_node).await.unwrap(); + + // Create task A (Ready status) + let task_a = common::create_test_task("ra-p5-test.1", "proj-1", "Task A", rustagent::graph::NodeStatus::Ready); + graph_store.create_node(&task_a).await.unwrap(); + + // Create task B that DependsOn A (Ready status) + let task_b = common::create_test_task("ra-p5-test.2", "proj-1", "Task B", rustagent::graph::NodeStatus::Ready); + graph_store.create_node(&task_b).await.unwrap(); + + // Create DependsOn edge: B depends on A + let depends_edge = rustagent::graph::GraphEdge { + id: "e-depends1".to_string(), + edge_type: rustagent::graph::EdgeType::DependsOn, + from_node: "ra-p5-test.2".to_string(), // B + to_node: "ra-p5-test.1".to_string(), // A + label: None, + created_at: chrono::Utc::now(), + }; + graph_store.add_edge(&depends_edge).await.unwrap(); + + // Create Contains edges to goal + graph_store.add_edge(&rustagent::graph::GraphEdge { + id: "e-contains1".to_string(), + edge_type: rustagent::graph::EdgeType::Contains, + from_node: "ra-p5-test".to_string(), + to_node: "ra-p5-test.1".to_string(), + label: None, + created_at: chrono::Utc::now(), + }).await.unwrap(); + + graph_store.add_edge(&rustagent::graph::GraphEdge { + id: "e-contains2".to_string(), + edge_type: rustagent::graph::EdgeType::Contains, + from_node: "ra-p5-test".to_string(), + to_node: "ra-p5-test.2".to_string(), + label: None, + created_at: chrono::Utc::now(), + }).await.unwrap(); + + // === Step 1: Simulate Task A failing (first attempt) === + let error_msg_1 = "First failure: database connection timeout"; + orchestrator.handle_task_retry_or_fail("ra-p5-test.1", error_msg_1).await.unwrap(); + + // Verify Task A is retried (Ready state) with previous_attempt in metadata + let task_a_after_retry1 = graph_store.get_node("ra-p5-test.1").await.unwrap().unwrap(); + assert_eq!(task_a_after_retry1.status, rustagent::graph::NodeStatus::Ready); + assert_eq!( + task_a_after_retry1.metadata.get("previous_attempt"), + Some(&error_msg_1.to_string()) + ); + assert_eq!( + task_a_after_retry1.metadata.get("retry_count"), + Some(&"1".to_string()) + ); + + // Verify Task B is still Ready (not blocked yet since A is being retried) + let task_b_check1 = graph_store.get_node("ra-p5-test.2").await.unwrap().unwrap(); + assert_eq!(task_b_check1.status, rustagent::graph::NodeStatus::Ready); + + // === Step 2: Simulate Task A failing again (second attempt, exceeding max_retries) === + let error_msg_2 = "Second failure: network unreachable"; + orchestrator.handle_task_retry_or_fail("ra-p5-test.1", error_msg_2).await.unwrap(); + + // Verify Task A is now Failed (retries exhausted) + let task_a_after_fail = graph_store.get_node("ra-p5-test.1").await.unwrap().unwrap(); + assert_eq!(task_a_after_fail.status, rustagent::graph::NodeStatus::Failed); + assert_eq!(task_a_after_fail.blocked_reason, Some(error_msg_2.to_string())); + + // Verify an Observation node was created for the failure + let obs_nodes = graph_store.query_nodes(&rustagent::graph::store::NodeQuery { + node_type: Some(rustagent::graph::NodeType::Observation), + status: None, + project_id: None, + parent_id: None, + query: None, + }).await.unwrap(); + assert!(!obs_nodes.is_empty(), "Expected an Observation node for task failure"); + + // Verify Task B is now Blocked (cascade occurred) + let task_b_after_cascade = graph_store.get_node("ra-p5-test.2").await.unwrap().unwrap(); + assert_eq!(task_b_after_cascade.status, rustagent::graph::NodeStatus::Blocked); + assert!( + task_b_after_cascade.blocked_reason.as_ref().map(|r| r.contains("ra-p5-test.1")).unwrap_or(false), + "Task B should be blocked by A: {:?}", + task_b_after_cascade.blocked_reason + ); + assert_eq!( + task_b_after_cascade.metadata.get("blocker_task_id"), + Some(&"ra-p5-test.1".to_string()) + ); + + // === Step 3: Manually complete Task A (simulate external fix) === + graph_store.update_node( + "ra-p5-test.1", + Some(rustagent::graph::NodeStatus::Completed), + None, + None, + None, + None, + ).await.unwrap(); + + // Verify Task A is Completed + let task_a_completed = graph_store.get_node("ra-p5-test.1").await.unwrap().unwrap(); + assert_eq!(task_a_completed.status, rustagent::graph::NodeStatus::Completed); + + // === Step 4: Call handle_scheduling which triggers try_unblock_tasks === + orchestrator.handle_scheduling().await.unwrap(); + + // Verify Task B is now Ready (unblocked) + let task_b_unblocked = graph_store.get_node("ra-p5-test.2").await.unwrap().unwrap(); + assert_eq!(task_b_unblocked.status, rustagent::graph::NodeStatus::Ready); + assert!(task_b_unblocked.blocked_reason.is_none(), "Blocked reason should be cleared"); + assert!(!task_b_unblocked.metadata.contains_key("blocker_task_id"), "blocker_task_id should be removed"); +} diff --git a/src/agent/orchestrator.rs b/src/agent/orchestrator.rs --- a/src/agent/orchestrator.rs +++ b/src/agent/orchestrator.rs @@ -825,6 +825,7 @@ self.graph_store.clone(), Some(self.message_bus.clone()), Some(worker_id.clone()), + self.project_path.clone(), ); // Create AgentRuntime diff --git a/src/tools/factory.rs b/src/tools/factory.rs --- a/src/tools/factory.rs +++ b/src/tools/factory.rs @@ -11,8 +11,10 @@ AddEdgeTool, ChooseOptionTool, ClaimTaskTool, CreateNodeTool, LogDecisionTool, QueryNodesTool, RecordObservationTool, RecordOutcomeTool, RevisitTool, SearchNodesTool, UpdateNodeTool, }; +use crate::tools::search::CodeSearchTool; use crate::tools::shell::RunCommandTool; use crate::tools::signal::SignalTool; +use std::path::PathBuf; use std::sync::Arc; pub fn create_default_registry( @@ -52,6 +54,7 @@ graph_store: Arc, message_bus: Option>, agent_id: Option, + project_root: PathBuf, ) -> ToolRegistry { let registry = create_default_registry(validator, permission_handler); @@ -70,6 +73,9 @@ // Register context tools registry.register(Arc::new(ReadAgentsMdTool::new())); + + // Register search tools + registry.register(Arc::new(CodeSearchTool::new(project_root))); // Register agent tools (only in multi-agent mode) if let (Some(bus), Some(id)) = (message_bus, agent_id) { -- tangled.sh