diff --git a/src/tools/graph_tools.rs b/src/tools/graph_tools.rs new file mode 100644 index 0000000..e00109b --- /dev/null +++ b/src/tools/graph_tools.rs @@ -0,0 +1,1167 @@ +use crate::graph::store::{GraphStore, NodeQuery}; +use crate::graph::{ + EdgeType, GraphEdge, GraphNode, NodeStatus, NodeType, generate_child_id, generate_edge_id, + generate_goal_id, +}; +use anyhow::{Context, Result}; +use async_trait::async_trait; +use chrono::Utc; +use serde_json::{Value, json}; +use std::collections::HashMap; +use std::sync::Arc; + +use super::Tool; + +/// Low-level tool for creating nodes +pub struct CreateNodeTool { + store: Arc, +} + +impl CreateNodeTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for CreateNodeTool { + fn name(&self) -> &str { + "create_node" + } + + fn description(&self) -> &str { + "Create a new node in the work graph" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "node_type": { + "type": "string", + "enum": ["goal", "task", "decision", "option", "outcome", "observation", "revisit"], + "description": "Type of node to create" + }, + "title": { + "type": "string", + "description": "Title of the node" + }, + "description": { + "type": "string", + "description": "Description of the node" + }, + "project_id": { + "type": "string", + "description": "Project ID (required)" + }, + "parent_id": { + "type": "string", + "description": "Parent node ID (optional, creates as child if provided)" + }, + "priority": { + "type": "string", + "enum": ["critical", "high", "medium", "low"], + "description": "Priority level (optional)" + }, + "metadata": { + "type": "object", + "description": "Additional metadata (optional)" + } + }, + "required": ["node_type", "title", "description", "project_id"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let node_type_str = params["node_type"] + .as_str() + .context("Missing 'node_type' parameter")?; + let node_type: NodeType = node_type_str.parse()?; + + let title = params["title"] + .as_str() + .context("Missing 'title' parameter")? + .to_string(); + + let description = params["description"] + .as_str() + .context("Missing 'description' parameter")? + .to_string(); + + let project_id = params["project_id"] + .as_str() + .context("Missing 'project_id' parameter")? + .to_string(); + + let parent_id_opt = params["parent_id"].as_str().map(|s| s.to_string()); + + // Generate ID based on parent + let id = if let Some(p_id) = &parent_id_opt { + let seq = self.store.next_child_seq(p_id).await?; + generate_child_id(p_id, seq) + } else { + generate_goal_id() + }; + + // Parse priority if provided + let priority = params["priority"].as_str().and_then(|p| p.parse().ok()); + + // Parse metadata if provided + let metadata: HashMap = params["metadata"] + .as_object() + .map(|m| m.iter().map(|(k, v)| (k.clone(), v.to_string())).collect()) + .unwrap_or_default(); + + // Create the node + let node = GraphNode { + id: id.clone(), + project_id, + node_type, + title, + description, + status: NodeStatus::Pending, + priority, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata, + }; + + self.store.create_node(&node).await?; + + // If parent was provided, a Contains edge is created automatically in create_node + Ok(json!({ + "id": id, + "message": "Node created successfully" + }) + .to_string()) + } +} + +/// Low-level tool for updating nodes +pub struct UpdateNodeTool { + store: Arc, +} + +impl UpdateNodeTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for UpdateNodeTool { + fn name(&self) -> &str { + "update_node" + } + + fn description(&self) -> &str { + "Update a node in the work graph" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "node_id": { + "type": "string", + "description": "ID of node to update" + }, + "status": { + "type": "string", + "enum": ["pending", "active", "completed", "cancelled", "ready", "claimed", + "in_progress", "review", "blocked", "failed", "decided", "superseded", + "abandoned", "chosen", "rejected"], + "description": "New status (optional)" + }, + "title": { + "type": "string", + "description": "New title (optional)" + }, + "description": { + "type": "string", + "description": "New description (optional)" + }, + "metadata": { + "type": "object", + "description": "New metadata (optional)" + } + }, + "required": ["node_id"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let node_id = params["node_id"] + .as_str() + .context("Missing 'node_id' parameter")?; + + let status = params["status"].as_str().and_then(|s| s.parse().ok()); + + let title = params["title"].as_str(); + let description = params["description"].as_str(); + + let metadata: Option> = params["metadata"] + .as_object() + .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()) + .await?; + + Ok(json!({ + "message": "Node updated successfully" + }) + .to_string()) + } +} + +/// Low-level tool for adding edges +pub struct AddEdgeTool { + store: Arc, +} + +impl AddEdgeTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for AddEdgeTool { + fn name(&self) -> &str { + "add_edge" + } + + fn description(&self) -> &str { + "Add an edge between two nodes in the work graph" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "edge_type": { + "type": "string", + "enum": ["contains", "depends_on", "leads_to", "chosen", "rejected", "supersedes", "informs"], + "description": "Type of edge" + }, + "from_node": { + "type": "string", + "description": "Source node ID" + }, + "to_node": { + "type": "string", + "description": "Target node ID" + }, + "label": { + "type": "string", + "description": "Edge label (optional)" + } + }, + "required": ["edge_type", "from_node", "to_node"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let edge_type_str = params["edge_type"] + .as_str() + .context("Missing 'edge_type' parameter")?; + let edge_type: EdgeType = edge_type_str.parse()?; + + let from_node = params["from_node"] + .as_str() + .context("Missing 'from_node' parameter")? + .to_string(); + + let to_node = params["to_node"] + .as_str() + .context("Missing 'to_node' parameter")? + .to_string(); + + let label = params["label"].as_str().map(|s| s.to_string()); + + let edge = GraphEdge { + id: generate_edge_id(), + edge_type, + from_node, + to_node, + label, + created_at: Utc::now(), + }; + + self.store.add_edge(&edge).await?; + + Ok(json!({ + "id": edge.id, + "message": "Edge created successfully" + }) + .to_string()) + } +} + +/// Low-level tool for querying nodes +pub struct QueryNodesTool { + store: Arc, +} + +impl QueryNodesTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for QueryNodesTool { + fn name(&self) -> &str { + "query_nodes" + } + + fn description(&self) -> &str { + "Query nodes in the work graph with flexible filters" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "node_type": { + "type": "string", + "enum": ["goal", "task", "decision", "option", "outcome", "observation", "revisit"], + "description": "Filter by node type (optional)" + }, + "status": { + "type": "string", + "enum": ["pending", "active", "completed", "cancelled", "ready", "claimed", + "in_progress", "review", "blocked", "failed", "decided", "superseded", + "abandoned", "chosen", "rejected"], + "description": "Filter by status (optional)" + }, + "project_id": { + "type": "string", + "description": "Filter by project (optional)" + }, + "parent_id": { + "type": "string", + "description": "Filter by parent node (optional)" + } + }, + "required": [] + }) + } + + async fn execute(&self, params: Value) -> Result { + let query = NodeQuery { + node_type: params["node_type"].as_str().and_then(|s| s.parse().ok()), + status: params["status"].as_str().and_then(|s| s.parse().ok()), + project_id: params["project_id"].as_str().map(|s| s.to_string()), + parent_id: params["parent_id"].as_str().map(|s| s.to_string()), + query: None, + }; + + let nodes = self.store.query_nodes(&query).await?; + + Ok(json!(nodes).to_string()) + } +} + +/// Low-level tool for full-text search +pub struct SearchNodesTool { + store: Arc, +} + +impl SearchNodesTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for SearchNodesTool { + fn name(&self) -> &str { + "search_nodes" + } + + fn description(&self) -> &str { + "Full-text search for nodes in the work graph" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query" + }, + "project_id": { + "type": "string", + "description": "Filter by project (optional)" + }, + "node_type": { + "type": "string", + "enum": ["goal", "task", "decision", "option", "outcome", "observation", "revisit"], + "description": "Filter by node type (optional)" + }, + "limit": { + "type": "integer", + "description": "Maximum results to return (default: 50)" + } + }, + "required": ["query"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let query = params["query"] + .as_str() + .context("Missing 'query' parameter")?; + + let project_id = params["project_id"].as_str(); + let node_type = params["node_type"].as_str().and_then(|s| s.parse().ok()); + let limit = params["limit"].as_u64().map(|n| n as usize).unwrap_or(50); + + let results = self + .store + .search_nodes(query, project_id, node_type, limit) + .await?; + + Ok(json!(results).to_string()) + } +} + +/// High-level tool for claiming a task +pub struct ClaimTaskTool { + store: Arc, +} + +impl ClaimTaskTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for ClaimTaskTool { + fn name(&self) -> &str { + "claim_task" + } + + fn description(&self) -> &str { + "Atomically claim a task for execution" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "node_id": { + "type": "string", + "description": "Task node ID" + }, + "agent_id": { + "type": "string", + "description": "Agent ID claiming the task" + } + }, + "required": ["node_id", "agent_id"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let node_id = params["node_id"] + .as_str() + .context("Missing 'node_id' parameter")?; + + let agent_id = params["agent_id"] + .as_str() + .context("Missing 'agent_id' parameter")?; + + let claimed = self.store.claim_task(node_id, agent_id).await?; + + Ok(json!({ + "claimed": claimed, + "message": if claimed { + "Task claimed successfully" + } else { + "Task was not in Ready state (may have been claimed by another agent)" + } + }) + .to_string()) + } +} + +/// High-level tool for logging a decision with options +pub struct LogDecisionTool { + store: Arc, +} + +impl LogDecisionTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for LogDecisionTool { + fn name(&self) -> &str { + "log_decision" + } + + fn description(&self) -> &str { + "Log a decision point with multiple options" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "title": { + "type": "string", + "description": "Decision title" + }, + "description": { + "type": "string", + "description": "Decision description" + }, + "project_id": { + "type": "string", + "description": "Project ID" + }, + "parent_id": { + "type": "string", + "description": "Parent node ID (optional)" + }, + "options": { + "type": "array", + "items": { + "type": "object", + "properties": { + "title": { "type": "string" }, + "description": { "type": "string" }, + "pros": { "type": "string" }, + "cons": { "type": "string" } + }, + "required": ["title", "description"] + }, + "description": "Options for this decision" + } + }, + "required": ["title", "description", "project_id", "options"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let title = params["title"] + .as_str() + .context("Missing 'title' parameter")? + .to_string(); + + let description = params["description"] + .as_str() + .context("Missing 'description' parameter")? + .to_string(); + + let project_id = params["project_id"] + .as_str() + .context("Missing 'project_id' parameter")? + .to_string(); + + let parent_id_opt = params["parent_id"].as_str().map(|s| s.to_string()); + + let options_arr = params["options"] + .as_array() + .context("Missing 'options' parameter")?; + + // Create the decision node + let decision_id = if let Some(p_id) = &parent_id_opt { + let seq = self.store.next_child_seq(p_id).await?; + generate_child_id(p_id, seq) + } else { + generate_goal_id() + }; + + let decision_node = GraphNode { + id: decision_id.clone(), + project_id: project_id.clone(), + node_type: NodeType::Decision, + title, + description, + status: NodeStatus::Active, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: HashMap::new(), + }; + + self.store.create_node(&decision_node).await?; + + // Create option nodes and LeadsTo edges + let mut option_ids = vec![]; + for option in options_arr { + let opt_title = option["title"] + .as_str() + .context("Missing option title")? + .to_string(); + + let opt_desc = option["description"] + .as_str() + .context("Missing option description")? + .to_string(); + + let seq = self.store.next_child_seq(&decision_id).await?; + let option_id = generate_child_id(&decision_id, seq); + + let mut opt_metadata = HashMap::new(); + if let Some(pros) = option["pros"].as_str() { + opt_metadata.insert("pros".to_string(), pros.to_string()); + } + if let Some(cons) = option["cons"].as_str() { + opt_metadata.insert("cons".to_string(), cons.to_string()); + } + + let option_node = GraphNode { + id: option_id.clone(), + project_id: project_id.clone(), + node_type: NodeType::Option, + title: opt_title, + description: opt_desc, + status: NodeStatus::Active, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: opt_metadata, + }; + + self.store.create_node(&option_node).await?; + + // Create LeadsTo edge from decision to option + let edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::LeadsTo, + from_node: decision_id.clone(), + to_node: option_id.clone(), + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&edge).await?; + option_ids.push(option_id); + } + + Ok(json!({ + "decision_id": decision_id, + "option_ids": option_ids, + "message": "Decision and options created successfully" + }) + .to_string()) + } +} + +/// High-level tool for choosing an option +pub struct ChooseOptionTool { + store: Arc, +} + +impl ChooseOptionTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for ChooseOptionTool { + fn name(&self) -> &str { + "choose_option" + } + + fn description(&self) -> &str { + "Choose an option for a decision" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "decision_id": { + "type": "string", + "description": "Decision node ID" + }, + "option_id": { + "type": "string", + "description": "Option node ID to choose" + }, + "rationale": { + "type": "string", + "description": "Rationale for the choice" + } + }, + "required": ["decision_id", "option_id", "rationale"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let decision_id = params["decision_id"] + .as_str() + .context("Missing 'decision_id' parameter")?; + + let option_id = params["option_id"] + .as_str() + .context("Missing 'option_id' parameter")?; + + let rationale = params["rationale"] + .as_str() + .context("Missing 'rationale' parameter")?; + + // Add Chosen edge from decision to chosen option + let chosen_edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::Chosen, + from_node: decision_id.to_string(), + to_node: option_id.to_string(), + label: Some(rationale.to_string()), + created_at: Utc::now(), + }; + + self.store.add_edge(&chosen_edge).await?; + + // Update chosen option status to Chosen + self.store + .update_node(option_id, Some(NodeStatus::Chosen), None, None, None) + .await?; + + // Find other options and add Rejected edges + // Query all options under this decision + let query = NodeQuery { + node_type: Some(NodeType::Option), + status: None, + project_id: None, + parent_id: Some(decision_id.to_string()), + query: None, + }; + + let options = self.store.query_nodes(&query).await?; + for option in options { + if option.id != option_id { + // Add Rejected edge from decision to this option + let rejected_edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::Rejected, + from_node: decision_id.to_string(), + to_node: option.id.clone(), + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&rejected_edge).await?; + + // Update option status to Rejected + self.store + .update_node(&option.id, Some(NodeStatus::Rejected), None, None, None) + .await?; + } + } + + // Update decision status to Decided + self.store + .update_node(decision_id, Some(NodeStatus::Decided), None, None, None) + .await?; + + Ok(json!({ + "message": "Option chosen successfully" + }) + .to_string()) + } +} + +/// High-level tool for recording an outcome +pub struct RecordOutcomeTool { + store: Arc, +} + +impl RecordOutcomeTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for RecordOutcomeTool { + fn name(&self) -> &str { + "record_outcome" + } + + fn description(&self) -> &str { + "Record the outcome of a task or decision" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "parent_id": { + "type": "string", + "description": "Parent node ID (task or decision)" + }, + "title": { + "type": "string", + "description": "Outcome title" + }, + "description": { + "type": "string", + "description": "Outcome description" + }, + "project_id": { + "type": "string", + "description": "Project ID" + }, + "success": { + "type": "boolean", + "description": "Whether the outcome was successful" + } + }, + "required": ["parent_id", "title", "description", "project_id", "success"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let parent_id = params["parent_id"] + .as_str() + .context("Missing 'parent_id' parameter")?; + + let title = params["title"] + .as_str() + .context("Missing 'title' parameter")? + .to_string(); + + let description = params["description"] + .as_str() + .context("Missing 'description' parameter")? + .to_string(); + + let project_id = params["project_id"] + .as_str() + .context("Missing 'project_id' parameter")? + .to_string(); + + let success = params["success"] + .as_bool() + .context("Missing 'success' parameter")?; + + // Create outcome node + let seq = self.store.next_child_seq(parent_id).await?; + let outcome_id = generate_child_id(parent_id, seq); + + let mut metadata = HashMap::new(); + metadata.insert("success".to_string(), success.to_string()); + + let outcome_node = GraphNode { + id: outcome_id.clone(), + project_id, + node_type: NodeType::Outcome, + title, + description, + status: NodeStatus::Completed, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: Some(Utc::now()), + blocked_reason: None, + metadata, + }; + + self.store.create_node(&outcome_node).await?; + + // Create LeadsTo edge from parent to outcome + let edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::LeadsTo, + from_node: parent_id.to_string(), + to_node: outcome_id.clone(), + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&edge).await?; + + Ok(json!({ + "outcome_id": outcome_id, + "message": "Outcome recorded successfully" + }) + .to_string()) + } +} + +/// High-level tool for recording an observation +pub struct RecordObservationTool { + store: Arc, +} + +impl RecordObservationTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for RecordObservationTool { + fn name(&self) -> &str { + "record_observation" + } + + fn description(&self) -> &str { + "Record an observation about the work" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "title": { + "type": "string", + "description": "Observation title" + }, + "description": { + "type": "string", + "description": "Observation description" + }, + "project_id": { + "type": "string", + "description": "Project ID" + }, + "related_node_id": { + "type": "string", + "description": "Related node ID for Informs edge (optional)" + } + }, + "required": ["title", "description", "project_id"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let title = params["title"] + .as_str() + .context("Missing 'title' parameter")? + .to_string(); + + let description = params["description"] + .as_str() + .context("Missing 'description' parameter")? + .to_string(); + + let project_id = params["project_id"] + .as_str() + .context("Missing 'project_id' parameter")? + .to_string(); + + let related_node_id = params["related_node_id"].as_str().map(|s| s.to_string()); + + // Create observation node + let observation_id = generate_goal_id(); + + let observation_node = GraphNode { + id: observation_id.clone(), + project_id, + node_type: NodeType::Observation, + title, + description, + status: NodeStatus::Active, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: HashMap::new(), + }; + + self.store.create_node(&observation_node).await?; + + // Create Informs edge if related node provided + if let Some(related_id) = related_node_id { + let edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::Informs, + from_node: observation_id.clone(), + to_node: related_id, + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&edge).await?; + } + + Ok(json!({ + "observation_id": observation_id, + "message": "Observation recorded successfully" + }) + .to_string()) + } +} + +/// High-level tool for revisiting a decision +pub struct RevisitTool { + store: Arc, +} + +impl RevisitTool { + pub fn new(store: Arc) -> Self { + Self { store } + } +} + +#[async_trait] +impl Tool for RevisitTool { + fn name(&self) -> &str { + "revisit" + } + + fn description(&self) -> &str { + "Revisit and potentially revise a past decision based on an outcome" + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "outcome_id": { + "type": "string", + "description": "Outcome node ID" + }, + "project_id": { + "type": "string", + "description": "Project ID" + }, + "reason": { + "type": "string", + "description": "Reason for revisiting" + }, + "new_decision_title": { + "type": "string", + "description": "Title for new decision if creating one (optional)" + } + }, + "required": ["outcome_id", "project_id", "reason"] + }) + } + + async fn execute(&self, params: Value) -> Result { + let outcome_id = params["outcome_id"] + .as_str() + .context("Missing 'outcome_id' parameter")?; + + let project_id = params["project_id"] + .as_str() + .context("Missing 'project_id' parameter")? + .to_string(); + + let reason = params["reason"] + .as_str() + .context("Missing 'reason' parameter")? + .to_string(); + + let new_decision_title = params["new_decision_title"].as_str(); + + // Create revisit node + let revisit_id = generate_goal_id(); + + let revisit_node = GraphNode { + id: revisit_id.clone(), + project_id: project_id.clone(), + node_type: NodeType::Revisit, + title: format!("Revisit of {}", outcome_id), + description: reason, + status: NodeStatus::Active, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: HashMap::new(), + }; + + self.store.create_node(&revisit_node).await?; + + // Create LeadsTo edge from outcome to revisit + let edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::LeadsTo, + from_node: outcome_id.to_string(), + to_node: revisit_id.clone(), + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&edge).await?; + + let mut result = json!({ + "revisit_id": revisit_id, + "message": "Revisit recorded successfully" + }); + + // Create new decision if title provided + if let Some(title) = new_decision_title { + let decision_id = generate_goal_id(); + + let decision_node = GraphNode { + id: decision_id.clone(), + project_id, + node_type: NodeType::Decision, + title: title.to_string(), + description: "Decision created from revisit".to_string(), + status: NodeStatus::Pending, + priority: None, + assigned_to: None, + created_by: None, + labels: vec![], + created_at: Utc::now(), + started_at: None, + completed_at: None, + blocked_reason: None, + metadata: HashMap::new(), + }; + + self.store.create_node(&decision_node).await?; + + // Create LeadsTo edge from revisit to new decision + let decision_edge = GraphEdge { + id: generate_edge_id(), + edge_type: EdgeType::LeadsTo, + from_node: revisit_id, + to_node: decision_id.clone(), + label: None, + created_at: Utc::now(), + }; + + self.store.add_edge(&decision_edge).await?; + + if let Some(obj) = result.as_object_mut() { + obj.insert("decision_id".to_string(), json!(decision_id)); + } + } + + Ok(result.to_string()) + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index f8938dc..42060f6 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -82,6 +82,7 @@ impl Default for ToolRegistry { pub mod factory; pub mod file; +pub mod graph_tools; pub mod permission_check; pub mod shell; pub mod signal; diff --git a/tests/graph_tools_test.rs b/tests/graph_tools_test.rs new file mode 100644 index 0000000..f1602c5 --- /dev/null +++ b/tests/graph_tools_test.rs @@ -0,0 +1,639 @@ +use rustagent::db::Database; +use rustagent::graph::store::{GraphStore, SqliteGraphStore}; +use rustagent::graph::{EdgeType, NodeStatus, NodeType}; +use rustagent::tools::Tool; +use rustagent::tools::graph_tools::*; +use serde_json::{Value, json}; +use std::sync::Arc; + +/// Create a test database in memory with a test project +async fn setup_test_db() -> anyhow::Result<(Database, Arc)> { + let db = Database::open_in_memory().await?; + + // Insert a test project + db.connection() + .call(|conn| { + let now = chrono::Utc::now().to_rfc3339(); + conn.execute( + "INSERT INTO projects (id, name, path, registered_at, config_overrides, metadata) + VALUES (?, ?, ?, ?, ?, ?)", + rusqlite::params![ + "proj-1", + "proj-1", + "/tmp/proj-1", + &now, + None::, + "{}" + ], + )?; + Ok(()) + }) + .await?; + + let store = Arc::new(SqliteGraphStore::new(db.clone())); + Ok((db, store)) +} + +#[tokio::test] +async fn test_create_node_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let tool = CreateNodeTool::new(store.clone()); + + let params = json!({ + "node_type": "task", + "title": "Test Task", + "description": "A test task", + "project_id": "proj-1" + }); + + let result = tool.execute(params).await?; + let parsed: Value = serde_json::from_str(&result)?; + + assert!(parsed["id"].as_str().is_some()); + assert!(parsed["id"].as_str().unwrap().starts_with("ra-")); + assert_eq!(parsed["message"], "Node created successfully"); + + // Verify node was created + let node_id = parsed["id"].as_str().unwrap(); + let node = store.get_node(node_id).await?.expect("Node not found"); + + assert_eq!(node.title, "Test Task"); + assert_eq!(node.node_type, NodeType::Task); + assert_eq!(node.status, NodeStatus::Pending); + + Ok(()) +} + +#[tokio::test] +async fn test_create_child_node_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let tool = CreateNodeTool::new(store.clone()); + + // Create parent + let parent_params = json!({ + "node_type": "goal", + "title": "Parent Goal", + "description": "A parent goal", + "project_id": "proj-1" + }); + + let parent_result = tool.execute(parent_params).await?; + let parent_parsed: Value = serde_json::from_str(&parent_result)?; + let parent_id = parent_parsed["id"].as_str().unwrap(); + + // Create child with parent_id + let child_params = json!({ + "node_type": "task", + "title": "Child Task", + "description": "A child task", + "project_id": "proj-1", + "parent_id": parent_id + }); + + let child_result = tool.execute(child_params).await?; + let child_parsed: Value = serde_json::from_str(&child_result)?; + let child_id = child_parsed["id"].as_str().unwrap(); + + // Verify child ID has parent prefix + assert!(child_id.starts_with(parent_id)); + assert!(child_id.contains(".")); + + // Verify Contains edge was created + let edges = store + .get_edges(parent_id, rustagent::graph::store::EdgeDirection::Outgoing) + .await?; + + assert!(!edges.is_empty()); + assert_eq!(edges[0].0.edge_type, EdgeType::Contains); + + Ok(()) +} + +#[tokio::test] +async fn test_update_node_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let update_tool = UpdateNodeTool::new(store.clone()); + + // Create a node + let create_params = json!({ + "node_type": "task", + "title": "Original Title", + "description": "Original description", + "project_id": "proj-1" + }); + + let create_result = create_tool.execute(create_params).await?; + let parsed: Value = serde_json::from_str(&create_result)?; + let node_id = parsed["id"].as_str().unwrap(); + + // Update the node + let update_params = json!({ + "node_id": node_id, + "title": "Updated Title", + "description": "Updated description", + "status": "ready" + }); + + update_tool.execute(update_params).await?; + + // Verify update + let node = store.get_node(node_id).await?.expect("Node not found"); + + assert_eq!(node.title, "Updated Title"); + assert_eq!(node.description, "Updated description"); + assert_eq!(node.status, NodeStatus::Ready); + + Ok(()) +} + +#[tokio::test] +async fn test_add_edge_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let edge_tool = AddEdgeTool::new(store.clone()); + + // Create two nodes + let params1 = json!({ + "node_type": "task", + "title": "Task 1", + "description": "First task", + "project_id": "proj-1" + }); + + let result1 = create_tool.execute(params1).await?; + let parsed1: Value = serde_json::from_str(&result1)?; + let node1_id = parsed1["id"].as_str().unwrap(); + + let params2 = json!({ + "node_type": "task", + "title": "Task 2", + "description": "Second task", + "project_id": "proj-1" + }); + + let result2 = create_tool.execute(params2).await?; + let parsed2: Value = serde_json::from_str(&result2)?; + let node2_id = parsed2["id"].as_str().unwrap(); + + // Add DependsOn edge + let edge_params = json!({ + "edge_type": "depends_on", + "from_node": node2_id, + "to_node": node1_id, + "label": "blocks" + }); + + let edge_result = edge_tool.execute(edge_params).await?; + let edge_parsed: Value = serde_json::from_str(&edge_result)?; + + assert!(edge_parsed["id"].as_str().is_some()); + assert_eq!(edge_parsed["message"], "Edge created successfully"); + + // Verify edge + let edges = store + .get_edges(node2_id, rustagent::graph::store::EdgeDirection::Outgoing) + .await?; + + assert_eq!(edges.len(), 1); + assert_eq!(edges[0].0.edge_type, EdgeType::DependsOn); + + Ok(()) +} + +#[tokio::test] +async fn test_query_nodes_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let query_tool = QueryNodesTool::new(store.clone()); + + // Create a few nodes + for i in 1..=3 { + let params = json!({ + "node_type": "task", + "title": format!("Task {}", i), + "description": "Test task", + "project_id": "proj-1" + }); + create_tool.execute(params).await?; + } + + // Query all tasks + let query_params = json!({ + "node_type": "task", + "project_id": "proj-1" + }); + + let result = query_tool.execute(query_params).await?; + let parsed: Vec = serde_json::from_str(&result)?; + + assert_eq!(parsed.len(), 3); + assert_eq!(parsed[0]["node_type"], "task"); + + Ok(()) +} + +#[tokio::test] +async fn test_search_nodes_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let search_tool = SearchNodesTool::new(store.clone()); + + // Create nodes with specific titles + let params = json!({ + "node_type": "task", + "title": "Authentication Task", + "description": "Handle user authentication", + "project_id": "proj-1" + }); + create_tool.execute(params).await?; + + // Search for "authentication" + let search_params = json!({ + "query": "authentication" + }); + + let result = search_tool.execute(search_params).await?; + let parsed: Vec = serde_json::from_str(&result)?; + + assert!(!parsed.is_empty()); + assert_eq!(parsed[0]["title"], "Authentication Task"); + + Ok(()) +} + +#[tokio::test] +async fn test_claim_task_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let update_tool = UpdateNodeTool::new(store.clone()); + let claim_tool = ClaimTaskTool::new(store.clone()); + + // Create a task + let create_params = json!({ + "node_type": "task", + "title": "Test Task", + "description": "To be claimed", + "project_id": "proj-1" + }); + + let create_result = create_tool.execute(create_params).await?; + let parsed: Value = serde_json::from_str(&create_result)?; + let task_id = parsed["id"].as_str().unwrap(); + + // Update status to Ready + let update_params = json!({ + "node_id": task_id, + "status": "ready" + }); + update_tool.execute(update_params).await?; + + // Claim the task + let claim_params = json!({ + "node_id": task_id, + "agent_id": "agent-1" + }); + + let claim_result = claim_tool.execute(claim_params).await?; + let claim_parsed: Value = serde_json::from_str(&claim_result)?; + + assert_eq!(claim_parsed["claimed"], true); + + // Verify node status + let node = store.get_node(task_id).await?.expect("Node not found"); + + assert_eq!(node.status, NodeStatus::Claimed); + assert_eq!(node.assigned_to, Some("agent-1".to_string())); + + Ok(()) +} + +#[tokio::test] +async fn test_log_decision_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let tool = LogDecisionTool::new(store.clone()); + + let params = json!({ + "title": "Architecture Decision", + "description": "Choose between microservices or monolith", + "project_id": "proj-1", + "options": [ + { + "title": "Microservices", + "description": "Multiple independent services", + "pros": "Scalability, independence", + "cons": "Complexity, latency" + }, + { + "title": "Monolith", + "description": "Single unified application", + "pros": "Simplicity, performance", + "cons": "Scalability limitations" + } + ] + }); + + let result = tool.execute(params).await?; + let parsed: Value = serde_json::from_str(&result)?; + + assert!(parsed["decision_id"].as_str().is_some()); + assert!(parsed["option_ids"].is_array()); + assert_eq!(parsed["option_ids"].as_array().unwrap().len(), 2); + + // Verify decision node was created + let decision_id = parsed["decision_id"].as_str().unwrap(); + let decision = store + .get_node(decision_id) + .await? + .expect("Decision not found"); + + assert_eq!(decision.node_type, NodeType::Decision); + assert_eq!(decision.status, NodeStatus::Active); + + // Verify option nodes were created + let option_ids = parsed["option_ids"].as_array().unwrap(); + for option_id_val in option_ids { + let option_id = option_id_val.as_str().unwrap(); + let option = store.get_node(option_id).await?.expect("Option not found"); + + assert_eq!(option.node_type, NodeType::Option); + assert_eq!(option.status, NodeStatus::Active); + } + + Ok(()) +} + +#[tokio::test] +async fn test_choose_option_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let decision_tool = LogDecisionTool::new(store.clone()); + let choose_tool = ChooseOptionTool::new(store.clone()); + + // Create a decision with options + let decision_params = json!({ + "title": "Test Decision", + "description": "Test", + "project_id": "proj-1", + "options": [ + { + "title": "Option A", + "description": "First option" + }, + { + "title": "Option B", + "description": "Second option" + } + ] + }); + + let decision_result = decision_tool.execute(decision_params).await?; + let decision_parsed: Value = serde_json::from_str(&decision_result)?; + + let decision_id = decision_parsed["decision_id"].as_str().unwrap(); + let option_ids = decision_parsed["option_ids"].as_array().unwrap(); + let chosen_option_id = option_ids[0].as_str().unwrap(); + + // Choose an option + let choose_params = json!({ + "decision_id": decision_id, + "option_id": chosen_option_id, + "rationale": "Best fit for our needs" + }); + + choose_tool.execute(choose_params).await?; + + // Verify decision status changed to Decided + let decision = store + .get_node(decision_id) + .await? + .expect("Decision not found"); + + assert_eq!(decision.status, NodeStatus::Decided); + + // Verify chosen option has Chosen status + let chosen_option = store + .get_node(chosen_option_id) + .await? + .expect("Option not found"); + + assert_eq!(chosen_option.status, NodeStatus::Chosen); + + // Verify other options are Rejected + let other_option_id = option_ids[1].as_str().unwrap(); + let other_option = store + .get_node(other_option_id) + .await? + .expect("Option not found"); + + assert_eq!(other_option.status, NodeStatus::Rejected); + + Ok(()) +} + +#[tokio::test] +async fn test_record_outcome_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let outcome_tool = RecordOutcomeTool::new(store.clone()); + + // Create a task + let task_params = json!({ + "node_type": "task", + "title": "Test Task", + "description": "Task to record outcome for", + "project_id": "proj-1" + }); + + let task_result = create_tool.execute(task_params).await?; + let task_parsed: Value = serde_json::from_str(&task_result)?; + let task_id = task_parsed["id"].as_str().unwrap(); + + // Record outcome + let outcome_params = json!({ + "parent_id": task_id, + "title": "Task Completed", + "description": "Successfully completed the task", + "project_id": "proj-1", + "success": true + }); + + let outcome_result = outcome_tool.execute(outcome_params).await?; + let outcome_parsed: Value = serde_json::from_str(&outcome_result)?; + + assert!(outcome_parsed["outcome_id"].as_str().is_some()); + + // Verify outcome node + let outcome_id = outcome_parsed["outcome_id"].as_str().unwrap(); + let outcome = store + .get_node(outcome_id) + .await? + .expect("Outcome not found"); + + assert_eq!(outcome.node_type, NodeType::Outcome); + assert_eq!(outcome.status, NodeStatus::Completed); + assert_eq!(outcome.metadata.get("success"), Some(&"true".to_string())); + + Ok(()) +} + +#[tokio::test] +async fn test_record_observation_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let obs_tool = RecordObservationTool::new(store.clone()); + + // Create a task to observe + let task_params = json!({ + "node_type": "task", + "title": "Test Task", + "description": "Task to observe", + "project_id": "proj-1" + }); + + let task_result = create_tool.execute(task_params).await?; + let task_parsed: Value = serde_json::from_str(&task_result)?; + let task_id = task_parsed["id"].as_str().unwrap(); + + // Record observation related to task + let obs_params = json!({ + "title": "Performance Issue Observed", + "description": "Task took longer than expected", + "project_id": "proj-1", + "related_node_id": task_id + }); + + let obs_result = obs_tool.execute(obs_params).await?; + let obs_parsed: Value = serde_json::from_str(&obs_result)?; + + assert!(obs_parsed["observation_id"].as_str().is_some()); + + // Verify observation node + let obs_id = obs_parsed["observation_id"].as_str().unwrap(); + let obs = store + .get_node(obs_id) + .await? + .expect("Observation not found"); + + assert_eq!(obs.node_type, NodeType::Observation); + assert_eq!(obs.status, NodeStatus::Active); + + // Verify Informs edge + let edges = store + .get_edges(obs_id, rustagent::graph::store::EdgeDirection::Outgoing) + .await?; + + assert_eq!(edges.len(), 1); + assert_eq!(edges[0].0.edge_type, EdgeType::Informs); + + Ok(()) +} + +#[tokio::test] +async fn test_revisit_tool() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + let create_tool = CreateNodeTool::new(store.clone()); + let outcome_tool = RecordOutcomeTool::new(store.clone()); + let revisit_tool = RevisitTool::new(store.clone()); + + // Create a task and outcome + let task_params = json!({ + "node_type": "task", + "title": "Test Task", + "description": "Task", + "project_id": "proj-1" + }); + + let task_result = create_tool.execute(task_params).await?; + let task_parsed: Value = serde_json::from_str(&task_result)?; + let task_id = task_parsed["id"].as_str().unwrap(); + + let outcome_params = json!({ + "parent_id": task_id, + "title": "Outcome", + "description": "Task outcome", + "project_id": "proj-1", + "success": true + }); + + let outcome_result = outcome_tool.execute(outcome_params).await?; + let outcome_parsed: Value = serde_json::from_str(&outcome_result)?; + let outcome_id = outcome_parsed["outcome_id"].as_str().unwrap(); + + // Revisit with new decision + let revisit_params = json!({ + "outcome_id": outcome_id, + "project_id": "proj-1", + "reason": "Results not as expected", + "new_decision_title": "Reconsider approach" + }); + + let revisit_result = revisit_tool.execute(revisit_params).await?; + let revisit_parsed: Value = serde_json::from_str(&revisit_result)?; + + assert!(revisit_parsed["revisit_id"].as_str().is_some()); + assert!(revisit_parsed["decision_id"].as_str().is_some()); + + // Verify revisit node + let revisit_id = revisit_parsed["revisit_id"].as_str().unwrap(); + let revisit = store + .get_node(revisit_id) + .await? + .expect("Revisit not found"); + + assert_eq!(revisit.node_type, NodeType::Revisit); + assert_eq!(revisit.status, NodeStatus::Active); + + // Verify new decision was created + let decision_id = revisit_parsed["decision_id"].as_str().unwrap(); + let decision = store + .get_node(decision_id) + .await? + .expect("Decision not found"); + + assert_eq!(decision.node_type, NodeType::Decision); + + Ok(()) +} + +#[tokio::test] +async fn test_tool_name_and_description() -> anyhow::Result<()> { + let (_db, store) = setup_test_db().await?; + + let tools: Vec<(Box, &str)> = vec![ + (Box::new(CreateNodeTool::new(store.clone())), "create_node"), + (Box::new(UpdateNodeTool::new(store.clone())), "update_node"), + (Box::new(AddEdgeTool::new(store.clone())), "add_edge"), + (Box::new(QueryNodesTool::new(store.clone())), "query_nodes"), + ( + Box::new(SearchNodesTool::new(store.clone())), + "search_nodes", + ), + (Box::new(ClaimTaskTool::new(store.clone())), "claim_task"), + ( + Box::new(LogDecisionTool::new(store.clone())), + "log_decision", + ), + ( + Box::new(ChooseOptionTool::new(store.clone())), + "choose_option", + ), + ( + Box::new(RecordOutcomeTool::new(store.clone())), + "record_outcome", + ), + ( + Box::new(RecordObservationTool::new(store.clone())), + "record_observation", + ), + (Box::new(RevisitTool::new(store.clone())), "revisit"), + ]; + + for (tool, expected_name) in tools { + assert_eq!(tool.name(), expected_name); + assert!(!tool.description().is_empty()); + let params = tool.parameters(); + assert!(params.is_object()); + } + + Ok(()) +}