mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Replace hardcoded intent patterns with job tools
Remove the brittle natural language pattern matching from the router and add job management tools to the normal tool registry instead. - Add job tools: create_job, list_jobs, job_status, cancel_job - Router now only handles explicit /commands - Natural language goes through agentic loop with all tools - LLM naturally picks appropriate tools based on user intent - Share ContextManager between job tools and Agent Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
27ffc12f6c
commit
7baf9e379d
+14
-13
@@ -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<HeartbeatConfig>,
|
||||
context_manager: Option<Arc<ContextManager>>,
|
||||
) -> 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,24 +393,20 @@ 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 { .. } => {
|
||||
|
||||
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
|
||||
{
|
||||
|
||||
+62
-140
@@ -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 {
|
||||
let content = message.content.trim();
|
||||
|
||||
// Check for commands
|
||||
if content.starts_with(&self.command_prefix) {
|
||||
return self.parse_command(content);
|
||||
/// Check if a message is an explicit command.
|
||||
pub fn is_command(&self, message: &IncomingMessage) -> bool {
|
||||
message.content.trim().starts_with(&self.command_prefix)
|
||||
}
|
||||
|
||||
// Try to extract intent from natural language
|
||||
self.extract_intent(content)
|
||||
/// 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<MessageIntent> {
|
||||
let content = message.content.trim();
|
||||
|
||||
if content.starts_with(&self.command_prefix) {
|
||||
Some(self.parse_command(content))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
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::<String>();
|
||||
if title.len() < first_sentence.len() {
|
||||
format!("{}...", title)
|
||||
} else {
|
||||
title
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a category from content.
|
||||
fn extract_category(content: &str) -> Option<String> {
|
||||
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<String> {
|
||||
// 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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+1
-1
@@ -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)?,
|
||||
|
||||
@@ -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...");
|
||||
|
||||
@@ -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<ContextManager>,
|
||||
}
|
||||
|
||||
impl CreateJobTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> 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<ToolOutput, ToolError> {
|
||||
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<ContextManager>,
|
||||
}
|
||||
|
||||
impl ListJobsTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> 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<ToolOutput, ToolError> {
|
||||
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<ContextManager>,
|
||||
}
|
||||
|
||||
impl JobStatusTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> 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<ToolOutput, ToolError> {
|
||||
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<ContextManager>,
|
||||
}
|
||||
|
||||
impl CancelJobTool {
|
||||
pub fn new(context_manager: Arc<ContextManager>) -> 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<ToolOutput, ToolError> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
|
||||
+17
-2
@@ -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<ContextManager>) {
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user