diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 41d13235..ca143690 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -64,13 +64,18 @@ pub struct Agent { impl Agent { /// Create a new agent. + /// + /// Optionally accepts a pre-created `ContextManager` for sharing with job tools. + /// If not provided, creates a new one. pub fn new( config: AgentConfig, deps: AgentDeps, channels: ChannelManager, heartbeat_config: Option, + context_manager: Option>, ) -> Self { - let context_manager = Arc::new(ContextManager::new(config.max_parallel_jobs)); + let context_manager = context_manager + .unwrap_or_else(|| Arc::new(ContextManager::new(config.max_parallel_jobs))); let scheduler = Arc::new(Scheduler::new( config.clone(), @@ -388,25 +393,21 @@ impl Agent { } } - // Route for job commands (bypass turn system) - // Build a temporary message with the content to route + // Handle explicit commands (starting with /) directly + // Everything else goes through the normal agentic loop with tools let temp_message = IncomingMessage { content: content.to_string(), ..message.clone() }; - let intent = self.router.route(&temp_message); - match &intent { - MessageIntent::CreateJob { .. } - | MessageIntent::CheckJobStatus { .. } - | MessageIntent::CancelJob { .. } - | MessageIntent::ListJobs { .. } - | MessageIntent::HelpJob { .. } - | MessageIntent::Command { .. } => { - return self.handle_job_or_command(intent, message).await; - } - _ => {} + + if let Some(intent) = self.router.route_command(&temp_message) { + // Explicit command like /status, /job, /list - handle directly + return self.handle_job_or_command(intent, message).await; } + // Natural language goes through the agentic loop + // Job tools (create_job, list_jobs, etc.) are in the tool registry + // Auto-compact if needed BEFORE adding new turn { let mut sess = session.lock().await; diff --git a/src/agent/router.rs b/src/agent/router.rs index 0e6e88d0..dccb4771 100644 --- a/src/agent/router.rs +++ b/src/agent/router.rs @@ -1,4 +1,8 @@ //! Message routing to appropriate handlers. +//! +//! The router handles explicit commands (starting with `/`). +//! Natural language intent classification is handled by `IntentClassifier` +//! which uses LLM + tools instead of brittle pattern matching. use crate::channels::IncomingMessage; @@ -27,7 +31,9 @@ pub enum MessageIntent { Unknown, } -/// Routes messages to appropriate handlers based on intent. +/// Routes messages to appropriate handlers based on explicit commands. +/// +/// For natural language messages, use `IntentClassifier` instead. pub struct Router { /// Command prefix (e.g., "/" or "!") command_prefix: String, @@ -47,17 +53,23 @@ impl Router { self } - /// Route a message to determine its intent. - pub fn route(&self, message: &IncomingMessage) -> MessageIntent { + /// Check if a message is an explicit command. + pub fn is_command(&self, message: &IncomingMessage) -> bool { + message.content.trim().starts_with(&self.command_prefix) + } + + /// Route an explicit command to determine its intent. + /// + /// Returns `None` if the message is not a command. + /// For non-commands, use `IntentClassifier::classify()` instead. + pub fn route_command(&self, message: &IncomingMessage) -> Option { let content = message.content.trim(); - // Check for commands if content.starts_with(&self.command_prefix) { - return self.parse_command(content); + Some(self.parse_command(content)) + } else { + None } - - // Try to extract intent from natural language - self.extract_intent(content) } fn parse_command(&self, content: &str) -> MessageIntent { @@ -111,61 +123,6 @@ impl Router { None => MessageIntent::Unknown, } } - - fn extract_intent(&self, content: &str) -> MessageIntent { - let lower = content.to_lowercase(); - - // Job creation patterns - must be explicit about creating a job - // More specific patterns to avoid capturing general conversation - let is_job_creation = lower.starts_with("create job ") - || lower.starts_with("new job ") - || lower.starts_with("schedule job ") - || lower.starts_with("run job ") - || (lower.contains("create") && lower.contains("job")); - - if is_job_creation { - return MessageIntent::CreateJob { - title: extract_title(content), - description: content.to_string(), - category: extract_category(content), - }; - } - - // Status check patterns - if lower.contains("status") - || lower.contains("how is") - || lower.contains("progress") - || lower.starts_with("check ") - { - return MessageIntent::CheckJobStatus { - job_id: extract_job_id(content), - }; - } - - // Cancel patterns - if lower.contains("cancel") || lower.contains("stop") || lower.contains("abort") { - if let Some(job_id) = extract_job_id(content) { - return MessageIntent::CancelJob { job_id }; - } - } - - // List patterns - if lower.starts_with("list") || lower.contains("show jobs") || lower.contains("my jobs") { - return MessageIntent::ListJobs { filter: None }; - } - - // Help patterns - if lower.contains("stuck") || lower.contains("not working") || lower.contains("fix") { - if let Some(job_id) = extract_job_id(content) { - return MessageIntent::HelpJob { job_id }; - } - } - - // Default to chat - MessageIntent::Chat { - content: content.to_string(), - } - } } impl Default for Router { @@ -174,53 +131,6 @@ impl Default for Router { } } -/// Extract a title from content. -fn extract_title(content: &str) -> String { - // Take first sentence or first N characters - let first_sentence = content.split('.').next().unwrap_or(content); - let title = first_sentence.chars().take(100).collect::(); - if title.len() < first_sentence.len() { - format!("{}...", title) - } else { - title - } -} - -/// Extract a category from content. -fn extract_category(content: &str) -> Option { - let lower = content.to_lowercase(); - - let categories = [ - ("code", "development"), - ("program", "development"), - ("website", "web"), - ("api", "development"), - ("data", "data"), - ("write", "writing"), - ("design", "design"), - ("research", "research"), - ]; - - for (keyword, category) in categories { - if lower.contains(keyword) { - return Some(category.to_string()); - } - } - - None -} - -/// Extract a job ID from content. -fn extract_job_id(content: &str) -> Option { - // Look for UUID patterns - let uuid_regex = regex::Regex::new( - r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}", - ) - .ok()?; - - uuid_regex.find(content).map(|m| m.as_str().to_string()) -} - #[cfg(test)] mod tests { use super::*; @@ -230,49 +140,61 @@ mod tests { let router = Router::new(); let msg = IncomingMessage::new("test", "user", "/status abc-123"); - let intent = router.route(&msg); + let intent = router.route_command(&msg); - assert!(matches!(intent, MessageIntent::CheckJobStatus { .. })); + assert!(matches!(intent, Some(MessageIntent::CheckJobStatus { .. }))); } #[test] - fn test_natural_language_routing() { + fn test_is_command() { let router = Router::new(); - // Explicit job creation with "create job" phrase - let msg = IncomingMessage::new("test", "user", "create job: build a website for me"); - let intent = router.route(&msg); - assert!(matches!(intent, MessageIntent::CreateJob { .. })); + let cmd_msg = IncomingMessage::new("test", "user", "/status"); + assert!(router.is_command(&cmd_msg)); - // Also matches when both "create" and "job" are present - let msg2 = IncomingMessage::new( - "test", - "user", - "I need to create a new job to build a website", - ); - let intent2 = router.route(&msg2); - assert!(matches!(intent2, MessageIntent::CreateJob { .. })); - - // General requests without explicit "job" fall through to Chat - let msg3 = IncomingMessage::new("test", "user", "Can you create a website for me?"); - let intent3 = router.route(&msg3); - assert!(matches!(intent3, MessageIntent::Chat { .. })); + let chat_msg = IncomingMessage::new("test", "user", "Hello there"); + assert!(!router.is_command(&chat_msg)); } #[test] - fn test_chat_fallback() { + fn test_non_command_returns_none() { let router = Router::new(); - let msg = IncomingMessage::new("test", "user", "Hello, how are you?"); - let intent = router.route(&msg); + // Natural language messages return None - they should use IntentClassifier + let msg = IncomingMessage::new("test", "user", "Can you create a website for me?"); + assert!(router.route_command(&msg).is_none()); - assert!(matches!(intent, MessageIntent::Chat { .. })); + let msg2 = IncomingMessage::new("test", "user", "Hello, how are you?"); + assert!(router.route_command(&msg2).is_none()); } #[test] - fn test_extract_job_id() { - let content = "Check status of job 550e8400-e29b-41d4-a716-446655440000"; - let id = extract_job_id(content); - assert_eq!(id, Some("550e8400-e29b-41d4-a716-446655440000".to_string())); + fn test_command_create_job() { + let router = Router::new(); + + let msg = IncomingMessage::new("test", "user", "/job build a website"); + let intent = router.route_command(&msg); + + match intent { + Some(MessageIntent::CreateJob { title, .. }) => { + assert_eq!(title, "build a website"); + } + _ => panic!("Expected CreateJob intent"), + } + } + + #[test] + fn test_command_list_jobs() { + let router = Router::new(); + + let msg = IncomingMessage::new("test", "user", "/list active"); + let intent = router.route_command(&msg); + + match intent { + Some(MessageIntent::ListJobs { filter }) => { + assert_eq!(filter, Some("active".to_string())); + } + _ => panic!("Expected ListJobs intent"), + } } } diff --git a/src/config.rs b/src/config.rs index bdcfe205..75022b2d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -665,7 +665,7 @@ impl SandboxModeConfig { key: "SANDBOX_ENABLED".to_string(), message: format!("must be 'true' or 'false': {e}"), })? - .unwrap_or(false), + .unwrap_or(true), policy: optional_env("SANDBOX_POLICY")?.unwrap_or_else(|| "readonly".to_string()), timeout_secs: parse_optional_env("SANDBOX_TIMEOUT_SECS", 120)?, memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?, diff --git a/src/main.rs b/src/main.rs index 1d5f8cd4..5e1de80e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,6 +13,7 @@ use near_agent::{ }, cli::{Cli, Command, run_tool_command}, config::Config, + context::ContextManager, history::Store, llm::{SessionConfig, create_llm_provider, create_session_manager}, safety::SafetyLayer, @@ -384,6 +385,12 @@ async fn main() -> anyhow::Result<()> { } } + // Create context manager (shared between job tools and agent) + let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs)); + + // Register job tools + tools.register_job_tools(Arc::clone(&context_manager)); + // Create and run the agent let deps = AgentDeps { store, @@ -397,6 +404,7 @@ async fn main() -> anyhow::Result<()> { deps, channels, Some(config.heartbeat.clone()), + Some(context_manager), ); tracing::info!("Agent initialized, starting main loop..."); diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs new file mode 100644 index 00000000..354fe14f --- /dev/null +++ b/src/tools/builtin/job.rs @@ -0,0 +1,417 @@ +//! Job management tools. +//! +//! These tools allow the LLM to manage jobs: +//! - Create new jobs/tasks +//! - List existing jobs +//! - Check job status +//! - Cancel running jobs + +use std::sync::Arc; + +use async_trait::async_trait; +use uuid::Uuid; + +use crate::context::{ContextManager, JobContext, JobState}; +use crate::tools::tool::{Tool, ToolError, ToolOutput}; + +/// Tool for creating a new job. +pub struct CreateJobTool { + context_manager: Arc, +} + +impl CreateJobTool { + pub fn new(context_manager: Arc) -> Self { + Self { context_manager } + } +} + +#[async_trait] +impl Tool for CreateJobTool { + fn name(&self) -> &str { + "create_job" + } + + fn description(&self) -> &str { + "Create a new job or task for the agent to work on. Use this when the user wants \ + you to do something substantial that should be tracked as a separate job." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "title": { + "type": "string", + "description": "A short title for the job (max 100 chars)" + }, + "description": { + "type": "string", + "description": "Full description of what needs to be done" + } + }, + "required": ["title", "description"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let title = params + .get("title") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'title' parameter".into()))?; + + let description = params + .get("description") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("missing 'description' parameter".into()) + })?; + + match self.context_manager.create_job(title, description).await { + Ok(job_id) => { + let result = serde_json::json!({ + "job_id": job_id.to_string(), + "title": title, + "status": "pending", + "message": format!("Created job '{}'", title) + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + Err(e) => { + let result = serde_json::json!({ + "error": e.to_string() + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + } + } + + fn requires_sanitization(&self) -> bool { + false + } +} + +/// Tool for listing jobs. +pub struct ListJobsTool { + context_manager: Arc, +} + +impl ListJobsTool { + pub fn new(context_manager: Arc) -> Self { + Self { context_manager } + } +} + +#[async_trait] +impl Tool for ListJobsTool { + fn name(&self) -> &str { + "list_jobs" + } + + fn description(&self) -> &str { + "List all jobs or filter by status. Shows job IDs, titles, and current status." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "filter": { + "type": "string", + "description": "Filter by status: 'active', 'completed', 'failed', 'all' (default: 'all')", + "enum": ["active", "completed", "failed", "all"] + } + } + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let filter = params + .get("filter") + .and_then(|v| v.as_str()) + .unwrap_or("all"); + + let job_ids = match filter { + "active" => self.context_manager.active_jobs().await, + _ => self.context_manager.all_jobs().await, + }; + + let mut jobs = Vec::new(); + for job_id in job_ids { + if let Ok(ctx) = self.context_manager.get_context(job_id).await { + let include = match filter { + "completed" => ctx.state == JobState::Completed, + "failed" => ctx.state == JobState::Failed, + "active" => ctx.state.is_active(), + _ => true, + }; + + if include { + jobs.push(serde_json::json!({ + "job_id": job_id.to_string(), + "title": ctx.title, + "status": format!("{:?}", ctx.state), + "created_at": ctx.created_at.to_rfc3339() + })); + } + } + } + + let summary = self.context_manager.summary().await; + + let result = serde_json::json!({ + "jobs": jobs, + "summary": { + "total": summary.total, + "pending": summary.pending, + "in_progress": summary.in_progress, + "completed": summary.completed, + "failed": summary.failed + } + }); + + Ok(ToolOutput::success(result, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false + } +} + +/// Tool for checking job status. +pub struct JobStatusTool { + context_manager: Arc, +} + +impl JobStatusTool { + pub fn new(context_manager: Arc) -> Self { + Self { context_manager } + } +} + +#[async_trait] +impl Tool for JobStatusTool { + fn name(&self) -> &str { + "job_status" + } + + fn description(&self) -> &str { + "Check the status and details of a specific job by its ID." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The UUID of the job to check" + } + }, + "required": ["job_id"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let job_id_str = params + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?; + + let job_id = Uuid::parse_str(job_id_str).map_err(|_| { + ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str)) + })?; + + match self.context_manager.get_context(job_id).await { + Ok(ctx) => { + let result = serde_json::json!({ + "job_id": job_id.to_string(), + "title": ctx.title, + "description": ctx.description, + "status": format!("{:?}", ctx.state), + "created_at": ctx.created_at.to_rfc3339(), + "started_at": ctx.started_at.map(|t| t.to_rfc3339()), + "completed_at": ctx.completed_at.map(|t| t.to_rfc3339()), + "actual_cost": ctx.actual_cost.to_string() + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + Err(e) => { + let result = serde_json::json!({ + "error": format!("Job not found: {}", e) + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + } + } + + fn requires_sanitization(&self) -> bool { + false + } +} + +/// Tool for canceling a job. +pub struct CancelJobTool { + context_manager: Arc, +} + +impl CancelJobTool { + pub fn new(context_manager: Arc) -> Self { + Self { context_manager } + } +} + +#[async_trait] +impl Tool for CancelJobTool { + fn name(&self) -> &str { + "cancel_job" + } + + fn description(&self) -> &str { + "Cancel a running or pending job. The job will be marked as cancelled and stopped." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The UUID of the job to cancel" + } + }, + "required": ["job_id"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let job_id_str = params + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("missing 'job_id' parameter".into()))?; + + let job_id = Uuid::parse_str(job_id_str).map_err(|_| { + ToolError::InvalidParameters(format!("invalid job ID format: {}", job_id_str)) + })?; + + // Transition to cancelled state + match self + .context_manager + .update_context(job_id, |ctx| { + ctx.transition_to(JobState::Cancelled, Some("Cancelled by user".to_string())) + }) + .await + { + Ok(Ok(())) => { + let result = serde_json::json!({ + "job_id": job_id.to_string(), + "status": "cancelled", + "message": "Job cancelled successfully" + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + Ok(Err(reason)) => { + let result = serde_json::json!({ + "error": format!("Cannot cancel job: {}", reason) + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + Err(e) => { + let result = serde_json::json!({ + "error": format!("Job not found: {}", e) + }); + Ok(ToolOutput::success(result, start.elapsed())) + } + } + } + + fn requires_approval(&self) -> bool { + true // Canceling a job should require approval + } + + fn requires_sanitization(&self) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_create_job_tool() { + let manager = Arc::new(ContextManager::new(5)); + let tool = CreateJobTool::new(manager.clone()); + + let params = serde_json::json!({ + "title": "Test Job", + "description": "A test job description" + }); + + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await.unwrap(); + + let job_id = result.result.get("job_id").unwrap().as_str().unwrap(); + assert!(!job_id.is_empty()); + } + + #[tokio::test] + async fn test_list_jobs_tool() { + let manager = Arc::new(ContextManager::new(5)); + + // Create some jobs + manager.create_job("Job 1", "Desc 1").await.unwrap(); + manager.create_job("Job 2", "Desc 2").await.unwrap(); + + let tool = ListJobsTool::new(manager); + + let params = serde_json::json!({}); + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await.unwrap(); + + let jobs = result.result.get("jobs").unwrap().as_array().unwrap(); + assert_eq!(jobs.len(), 2); + } + + #[tokio::test] + async fn test_job_status_tool() { + let manager = Arc::new(ContextManager::new(5)); + let job_id = manager.create_job("Test Job", "Description").await.unwrap(); + + let tool = JobStatusTool::new(manager); + + let params = serde_json::json!({ + "job_id": job_id.to_string() + }); + let ctx = JobContext::default(); + let result = tool.execute(params, &ctx).await.unwrap(); + + assert_eq!( + result.result.get("title").unwrap().as_str().unwrap(), + "Test Job" + ); + } +} diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index 6d9c8c64..b922ca66 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -4,6 +4,7 @@ mod echo; mod ecommerce; mod file; mod http; +mod job; mod json; mod marketplace; mod memory; @@ -16,6 +17,7 @@ pub use echo::EchoTool; pub use ecommerce::EcommerceTool; pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool}; pub use http::HttpTool; +pub use job::{CancelJobTool, CreateJobTool, JobStatusTool, ListJobsTool}; pub use json::JsonTool; pub use marketplace::MarketplaceTool; pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool}; diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 805b29fb..9e87b751 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -5,12 +5,14 @@ use std::sync::Arc; use tokio::sync::RwLock; +use crate::context::ContextManager; use crate::llm::{LlmProvider, ToolDefinition}; use crate::safety::SafetyLayer; use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder}; use crate::tools::builtin::{ - ApplyPatchTool, EchoTool, HttpTool, JsonTool, ListDirTool, MemoryReadTool, MemorySearchTool, - MemoryTreeTool, MemoryWriteTool, ReadFileTool, ShellTool, TimeTool, WriteFileTool, + ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, HttpTool, JobStatusTool, JsonTool, + ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, + ReadFileTool, ShellTool, TimeTool, WriteFileTool, }; use crate::tools::tool::Tool; use crate::tools::wasm::{ @@ -144,6 +146,19 @@ impl ToolRegistry { tracing::info!("Registered 4 memory tools"); } + /// Register job management tools. + /// + /// Job tools allow the LLM to create, list, check status, and cancel jobs. + /// These enable natural language job management without hardcoded intent parsing. + pub fn register_job_tools(&self, context_manager: Arc) { + self.register_sync(Arc::new(CreateJobTool::new(Arc::clone(&context_manager)))); + self.register_sync(Arc::new(ListJobsTool::new(Arc::clone(&context_manager)))); + self.register_sync(Arc::new(JobStatusTool::new(Arc::clone(&context_manager)))); + self.register_sync(Arc::new(CancelJobTool::new(context_manager))); + + tracing::info!("Registered 4 job management tools"); + } + /// Register the software builder tool. /// /// The builder tool allows the agent to create new software including WASM tools,