From 7630355d49760fc0ce9dfcc7e7505556d6b6534c Mon Sep 17 00:00:00 2001 From: David Hagerty Date: Mon, 9 Feb 2026 11:43:24 -0500 Subject: [PATCH] fix: address 8 code review issues for Rustagent V2 Phase 1 CRITICAL: - C1: Fix create_session call in run command (was passing wrong arguments) IMPORTANT: - I1: Add blocked_reason parameter to update_node trait and implementations - I2: Add unit test for ContextBuilder::build_system_prompt output format - I3: Fix clippy type_complexity warning in MockLlmClient with type alias - I4: Prefix unused variable with underscore in interchange_test.rs - I5: Add #![allow(dead_code)] to tests/common/mod.rs MINOR: - M1: Remove unnecessary #[allow(dead_code)] from resolve_project in main.rs - M2: Wrap create_session INSERT in BEGIN IMMEDIATE transaction All tests pass (173 passed), clippy clean, code formatted. Co-Authored-By: Claude Opus 4.6 --- src/context/mod.rs | 238 +++++++++++++++++++++++++++++++++ src/graph/interchange.rs | 1 + src/graph/session.rs | 8 +- src/graph/store.rs | 7 + src/llm/mock.rs | 3 +- src/main.rs | 7 +- src/tools/graph_tools.rs | 22 ++- tests/common/mod.rs | 3 + tests/graph_dependency_test.rs | 9 +- tests/graph_store_test.rs | 11 +- tests/interchange_test.rs | 2 +- 11 files changed, 299 insertions(+), 12 deletions(-) diff --git a/src/context/mod.rs b/src/context/mod.rs index de21be0..0258c2c 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -246,4 +246,242 @@ mod tests { assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("AGENTS.md")); } + + #[test] + fn test_build_system_prompt_output_format() { + use crate::agent::profile::{AgentProfile, ProfileLlmConfig}; + use crate::graph::store::GraphStore; + use crate::graph::{GraphNode, NodeStatus, NodeType, Priority}; + use crate::security::SecurityScope; + use anyhow::Result; + use async_trait::async_trait; + use chrono::Utc; + use std::collections::HashMap; + use std::sync::Arc; + + // Minimal mock GraphStore for testing + struct TestGraphStore; + + #[async_trait] + impl GraphStore for TestGraphStore { + async fn create_node(&self, _node: &GraphNode) -> Result<()> { + Ok(()) + } + async fn update_node( + &self, + _id: &str, + _status: Option, + _title: Option<&str>, + _description: Option<&str>, + _blocked_reason: Option<&str>, + _metadata: Option<&HashMap>, + ) -> Result<()> { + Ok(()) + } + async fn get_node(&self, _id: &str) -> Result> { + Ok(None) + } + async fn query_nodes( + &self, + _query: &crate::graph::store::NodeQuery, + ) -> Result> { + Ok(vec![]) + } + async fn claim_task(&self, _node_id: &str, _agent_id: &str) -> Result { + Ok(false) + } + async fn get_ready_tasks(&self, _goal_id: &str) -> Result> { + Ok(vec![]) + } + async fn get_next_task(&self, _goal_id: &str) -> Result> { + Ok(None) + } + async fn add_edge(&self, _edge: &crate::graph::GraphEdge) -> Result<()> { + Ok(()) + } + async fn remove_edge(&self, _edge_id: &str) -> Result<()> { + Ok(()) + } + async fn get_edges( + &self, + _node_id: &str, + _direction: crate::graph::store::EdgeDirection, + ) -> Result> { + Ok(vec![]) + } + async fn get_children( + &self, + _node_id: &str, + ) -> Result> { + Ok(vec![]) + } + async fn get_subtree(&self, _node_id: &str) -> Result> { + Ok(vec![]) + } + async fn get_active_decisions(&self, _project_id: &str) -> Result> { + Ok(vec![]) + } + async fn get_full_graph( + &self, + _goal_id: &str, + ) -> Result { + Ok(crate::graph::store::WorkGraph { + nodes: vec![], + edges: vec![], + }) + } + async fn search_nodes( + &self, + _query: &str, + _project_id: Option<&str>, + _node_type: Option, + _limit: usize, + ) -> Result> { + Ok(vec![]) + } + async fn next_child_seq(&self, _parent_id: &str) -> Result { + Ok(1) + } + } + + // Create mock profile + let profile = AgentProfile { + name: "test_coder".to_string(), + extends: None, + role: "You are a helpful code assistant".to_string(), + system_prompt: "Follow these rules carefully".to_string(), + allowed_tools: vec!["read_file".to_string(), "write_file".to_string()], + security: SecurityScope { + allowed_paths: vec!["*".to_string()], + denied_paths: vec![], + allowed_commands: vec!["*".to_string()], + read_only: false, + can_create_files: true, + network_access: false, + }, + llm: ProfileLlmConfig::default(), + turn_limit: Some(100), + token_budget: Some(100_000), + }; + + // Create mock work package tasks + let mut task_metadata = HashMap::new(); + task_metadata.insert( + "acceptance_criteria".to_string(), + "AC1: Task should pass tests".to_string(), + ); + + let work_package_tasks = vec![GraphNode { + id: "task-1".to_string(), + project_id: "proj-1".to_string(), + node_type: NodeType::Task, + title: "Implement feature".to_string(), + description: "Implement a new feature".to_string(), + status: NodeStatus::Ready, + priority: Some(Priority::High), + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: task_metadata, + }]; + + // Create mock decisions + let mut decision_metadata = HashMap::new(); + decision_metadata.insert("chosen_option".to_string(), "Option B".to_string()); + + let relevant_decisions = vec![GraphNode { + id: "decision-1".to_string(), + project_id: "proj-1".to_string(), + node_type: NodeType::Decision, + title: "Architecture decision".to_string(), + description: "Choose architecture".to_string(), + status: NodeStatus::Decided, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: decision_metadata, + }]; + + // Create agent context + let ctx = AgentContext { + work_package_tasks, + relevant_decisions, + handoff_notes: Some("Previous session notes".to_string()), + agents_md_summaries: vec![("src/AGENTS.md".to_string(), "Code standards".to_string())], + profile, + project_path: PathBuf::from("/test/project"), + graph_store: Arc::new(TestGraphStore), + }; + + // Build system prompt + let prompt = ContextBuilder::build_system_prompt(&ctx); + + // Verify expected sections are present + assert!(prompt.contains("## Role"), "Should contain Role section"); + assert!( + prompt.contains("You are a helpful code assistant"), + "Should contain profile role" + ); + + assert!(prompt.contains("## Task"), "Should contain Task section"); + assert!(prompt.contains("[TASK]"), "Should contain task marker"); + assert!(prompt.contains("task-1"), "Should contain task ID"); + assert!( + prompt.contains("[CRITERIA]"), + "Should contain acceptance criteria marker" + ); + + assert!( + prompt.contains("## Session Continuity"), + "Should contain Session Continuity section" + ); + assert!( + prompt.contains("[HANDOFF]"), + "Should contain handoff marker" + ); + assert!( + prompt.contains("Previous session notes"), + "Should contain handoff notes" + ); + + assert!( + prompt.contains("## Active Decisions"), + "Should contain Active Decisions section" + ); + assert!( + prompt.contains("[DECISION]"), + "Should contain decision marker" + ); + assert!(prompt.contains("decision-1"), "Should contain decision ID"); + assert!(prompt.contains("chosen:"), "Should contain chosen option"); + + assert!( + prompt.contains("## Relevant Observations"), + "Should contain Relevant Observations section" + ); + + assert!( + prompt.contains("## Project Conventions"), + "Should contain Project Conventions section" + ); + assert!( + prompt.contains("src/AGENTS.md"), + "Should contain agents_md path" + ); + + assert!(prompt.contains("## Rules"), "Should contain Rules section"); + assert!( + prompt.contains("Follow these rules carefully"), + "Should contain system prompt rules" + ); + } } diff --git a/src/graph/interchange.rs b/src/graph/interchange.rs index c40f74c..5a51d83 100644 --- a/src/graph/interchange.rs +++ b/src/graph/interchange.rs @@ -241,6 +241,7 @@ pub async fn import_goal( Some(node.status), Some(&node.title), Some(&node.description), + node.blocked_reason.as_deref(), Some(&node.metadata), ) .await?; diff --git a/src/graph/session.rs b/src/graph/session.rs index 696f6b8..e2a2792 100644 --- a/src/graph/session.rs +++ b/src/graph/session.rs @@ -42,6 +42,9 @@ impl SessionStore { self.db .connection() .call(move |conn| { + conn.execute_batch("BEGIN IMMEDIATE") + .map_err(tokio_rusqlite::Error::Rusqlite)?; + conn.execute( "INSERT INTO sessions (id, project_id, goal_id, started_at, agent_ids) VALUES (?, ?, ?, ?, ?)", @@ -53,7 +56,10 @@ impl SessionStore { "[]" ], ) - .map_err(tokio_rusqlite::Error::Rusqlite) + .map_err(tokio_rusqlite::Error::Rusqlite)?; + + conn.execute_batch("COMMIT") + .map_err(tokio_rusqlite::Error::Rusqlite) }) .await .map_err(|e| anyhow!("failed to insert session: {}", e))?; diff --git a/src/graph/store.rs b/src/graph/store.rs index 9349824..ee998e2 100644 --- a/src/graph/store.rs +++ b/src/graph/store.rs @@ -50,6 +50,7 @@ pub trait GraphStore: Send + Sync { status: Option, title: Option<&str>, description: Option<&str>, + blocked_reason: Option<&str>, metadata: Option<&HashMap>, ) -> Result<()>; @@ -316,12 +317,14 @@ impl GraphStore for SqliteGraphStore { status: Option, title: Option<&str>, description: Option<&str>, + blocked_reason: Option<&str>, metadata: Option<&HashMap>, ) -> Result<()> { let id = id.to_string(); let status_str = status.map(|s| s.to_string()); let title_owned = title.map(|t| t.to_string()); let description_owned = description.map(|d| d.to_string()); + let blocked_reason_owned = blocked_reason.map(|r| r.to_string()); let metadata_json = metadata.map(serde_json::to_string).transpose()?; self.db @@ -363,6 +366,10 @@ impl GraphStore for SqliteGraphStore { updates.push("description = ?"); params.push(d); } + if let Some(r) = &blocked_reason_owned { + updates.push("blocked_reason = ?"); + params.push(r); + } if let Some(m) = &metadata_json { updates.push("metadata = ?"); params.push(m); diff --git a/src/llm/mock.rs b/src/llm/mock.rs index 5e76c26..d597f43 100644 --- a/src/llm/mock.rs +++ b/src/llm/mock.rs @@ -4,9 +4,10 @@ use std::collections::VecDeque; use std::sync::{Arc, Mutex}; type RecordedCalls = Vec<(Vec, Vec)>; +type MockResponseQueue = VecDeque<(ResponseContent, Option)>; pub struct MockLlmClient { - responses: Arc)>>>, + responses: Arc>, recorded_calls: Arc>, token_counts: Arc>>, // (input_tokens, output_tokens) } diff --git a/src/main.rs b/src/main.rs index 03a280a..a84021c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -218,7 +218,6 @@ fn db_path() -> anyhow::Result { } /// Resolve project from --project flag or current working directory -#[allow(dead_code)] async fn resolve_project( db: &db::Database, project_name: Option<&str>, @@ -337,7 +336,7 @@ async fn main() -> anyhow::Result<()> { // Create session let session_store = rustagent::graph::session::SessionStore::new(database.clone()); - let session = session_store.create_session(&goal_id, &profile).await?; + let session = session_store.create_session(&project.id, &goal_id).await?; println!("Started session: {}", session.id); // Resolve profile @@ -405,6 +404,7 @@ async fn main() -> anyhow::Result<()> { None, None, None, + None, ) .await?; } @@ -415,6 +415,7 @@ async fn main() -> anyhow::Result<()> { &goal_id, Some(rustagent::graph::NodeStatus::Blocked), None, + None, Some(&reason), None, ) @@ -427,6 +428,7 @@ async fn main() -> anyhow::Result<()> { &goal_id, Some(rustagent::graph::NodeStatus::Failed), None, + None, Some(&error), None, ) @@ -442,6 +444,7 @@ async fn main() -> anyhow::Result<()> { &goal_id, Some(rustagent::graph::NodeStatus::Completed), None, + None, Some(&format!( "Token budget exhausted after {} tokens", tokens_used diff --git a/src/tools/graph_tools.rs b/src/tools/graph_tools.rs index e00109b..65685b8 100644 --- a/src/tools/graph_tools.rs +++ b/src/tools/graph_tools.rs @@ -210,7 +210,7 @@ impl Tool for UpdateNodeTool { .map(|m| m.iter().map(|(k, v)| (k.clone(), v.to_string())).collect()); self.store - .update_node(node_id, status, title, description, metadata.as_ref()) + .update_node(node_id, status, title, description, None, metadata.as_ref()) .await?; Ok(json!({ @@ -741,7 +741,7 @@ impl Tool for ChooseOptionTool { // Update chosen option status to Chosen self.store - .update_node(option_id, Some(NodeStatus::Chosen), None, None, None) + .update_node(option_id, Some(NodeStatus::Chosen), None, None, None, None) .await?; // Find other options and add Rejected edges @@ -771,14 +771,28 @@ impl Tool for ChooseOptionTool { // Update option status to Rejected self.store - .update_node(&option.id, Some(NodeStatus::Rejected), None, None, None) + .update_node( + &option.id, + Some(NodeStatus::Rejected), + None, + None, + None, + None, + ) .await?; } } // Update decision status to Decided self.store - .update_node(decision_id, Some(NodeStatus::Decided), None, None, None) + .update_node( + decision_id, + Some(NodeStatus::Decided), + None, + None, + None, + None, + ) .await?; Ok(json!({ diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 49efdab..05c53e1 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use anyhow::Result; use async_trait::async_trait; use chrono::Utc; @@ -214,6 +216,7 @@ impl GraphStore for MockGraphStore { _status: Option, _title: Option<&str>, _description: Option<&str>, + _blocked_reason: Option<&str>, _metadata: Option<&HashMap>, ) -> Result<()> { Ok(()) diff --git a/tests/graph_dependency_test.rs b/tests/graph_dependency_test.rs index 847f4fd..6648641 100644 --- a/tests/graph_dependency_test.rs +++ b/tests/graph_dependency_test.rs @@ -52,7 +52,14 @@ async fn test_task_pending_to_ready_when_deps_complete() -> Result<()> { // Complete Task A store - .update_node("ra-a1b2.1", Some(NodeStatus::Completed), None, None, None) + .update_node( + "ra-a1b2.1", + Some(NodeStatus::Completed), + None, + None, + None, + None, + ) .await?; // Now Task B should be Ready (automatically promoted by the status transition hook) diff --git a/tests/graph_store_test.rs b/tests/graph_store_test.rs index 02bed13..435748d 100644 --- a/tests/graph_store_test.rs +++ b/tests/graph_store_test.rs @@ -47,7 +47,14 @@ async fn test_update_node_status() -> Result<()> { // Update status to InProgress (valid for Task) store - .update_node("ra-a1b2.1", Some(NodeStatus::InProgress), None, None, None) + .update_node( + "ra-a1b2.1", + Some(NodeStatus::InProgress), + None, + None, + None, + None, + ) .await?; let updated = store.get_node("ra-a1b2.1").await?; @@ -69,7 +76,7 @@ async fn test_update_node_title() -> Result<()> { // Update title store - .update_node("ra-a1b2.1", None, Some("New Title"), None, None) + .update_node("ra-a1b2.1", None, Some("New Title"), None, None, None) .await?; let updated = store.get_node("ra-a1b2.1").await?; diff --git a/tests/interchange_test.rs b/tests/interchange_test.rs index 650eef5..320968e 100644 --- a/tests/interchange_test.rs +++ b/tests/interchange_test.rs @@ -341,7 +341,7 @@ async fn test_round_trip_export_import() -> Result<()> { let export1 = export_goal(&graph_store, "ra-test", "test-project").await?; // Parse the export - let parsed1: toml::Value = toml::from_str(&export1)?; + let _parsed1: toml::Value = toml::from_str(&export1)?; // Verify we can round-trip through import let (_db2, graph_store2) = setup_test_env().await?; -- 2.51.2