use std::sync::Arc; use async_trait::async_trait; use serde_json::json; use crate::config::SchedulerSchedule; use crate::storage::{ SchedulerJobRecord, SchedulerJobState, SchedulerJobUpsert, SessionStore, }; use crate::tools::traits::{Tool, ToolResult}; pub struct SchedulerManageTool { store: Arc, } impl SchedulerManageTool { pub fn new(store: Arc) -> Self { Self { store } } } #[async_trait] impl Tool for SchedulerManageTool { fn name(&self) -> &str { "scheduler_manage" } fn description(&self) -> &str { "Manage DB-backed scheduled jobs. Supports actions: list, get, put, delete, pause, resume. Jobs persist in SQLite and are executed by the scheduler runtime." } fn parameters_schema(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "action": { "type": "string", "enum": ["list", "get", "put", "delete", "pause", "resume"] }, "id": { "type": "string", "description": "Job id" }, "enabled_only": { "type": "boolean", "description": "Only for list action" }, "kind": { "type": "string", "enum": ["internal_event", "outbound_message", "agent_task"] }, "schedule": { "type": "object", "description": "Schedule object, for example {type: 'interval', seconds: 300} or {type: 'cron', expression: '0 9 * * *'}" }, "target": { "type": "object" }, "payload": { "type": "object", "description": "Job payload. agent_task supports prompt, fresh_session, system_prompt, sender_id, metadata. outbound_message expects content. internal_event expects event." }, "max_runs": { "type": ["integer", "null"] } }, "required": ["action"] }) } async fn execute(&self, args: serde_json::Value) -> anyhow::Result { let action = match args.get("action").and_then(|value| value.as_str()) { Some(action) => action, None => return Ok(error_result("Missing required parameter: action")), }; let output = match action { "list" => { let enabled_only = args .get("enabled_only") .and_then(|value| value.as_bool()) .unwrap_or(false); let jobs = self.store.list_scheduler_jobs(enabled_only)?; json!(jobs.iter().map(record_to_json).collect::>()) } "get" => { let id = require_str(&args, "id")?; match self.store.get_scheduler_job(id)? { Some(record) => record_to_json(&record), None => return Ok(error_result(&format!("scheduler job '{}' not found", id))), } } "put" => { let input = build_upsert(&args)?; let record = self.store.upsert_scheduler_job(&input)?; record_to_json(&record) } "delete" => { let id = require_str(&args, "id")?; self.store.delete_scheduler_job(id)?; json!({"status": "deleted", "id": id}) } "pause" => { let id = require_str(&args, "id")?; let record = self .store .get_scheduler_job(id)? .ok_or_else(|| anyhow::anyhow!("scheduler job '{}' not found", id))?; let mut input = record_to_upsert(&record); input.enabled = false; input.state = SchedulerJobState::Paused; input.paused_at = Some(current_timestamp()); input.next_fire_at = None; let saved = self.store.upsert_scheduler_job(&input)?; record_to_json(&saved) } "resume" => { let id = require_str(&args, "id")?; let record = self .store .get_scheduler_job(id)? .ok_or_else(|| anyhow::anyhow!("scheduler job '{}' not found", id))?; let mut input = record_to_upsert(&record); input.enabled = true; input.state = SchedulerJobState::Scheduled; input.paused_at = None; input.completed_at = None; input.next_fire_at = None; let saved = self.store.upsert_scheduler_job(&input)?; record_to_json(&saved) } _ => return Ok(error_result("Unsupported action")), }; Ok(ToolResult { success: true, output: serde_json::to_string_pretty(&output)?, error: None, }) } } fn build_upsert(args: &serde_json::Value) -> anyhow::Result { let id = require_str(args, "id")?.to_string(); let kind = require_str(args, "kind")?.to_string(); let schedule_value = args .get("schedule") .cloned() .ok_or_else(|| anyhow::anyhow!("Missing required parameter: schedule"))?; let schedule: SchedulerSchedule = serde_json::from_value(schedule_value.clone())?; schedule.validate(&id)?; let (interval_secs, startup_delay_secs) = match &schedule { SchedulerSchedule::Interval { seconds, startup_delay_secs, } => (*seconds as i64, *startup_delay_secs as i64), _ => (0, 0), }; Ok(SchedulerJobUpsert { id, kind, schedule: serde_json::to_value(schedule)?, interval_secs, startup_delay_secs, target: args.get("target").cloned().unwrap_or_else(|| json!({})), payload: args.get("payload").cloned().unwrap_or_else(|| json!({})), enabled: args.get("enabled").and_then(|value| value.as_bool()).unwrap_or(true), state: if args.get("enabled").and_then(|value| value.as_bool()).unwrap_or(true) { SchedulerJobState::Scheduled } else { SchedulerJobState::Paused }, last_status: None, last_error: None, run_count: 0, max_runs: args.get("max_runs").and_then(|value| value.as_i64()), last_fired_at: None, next_fire_at: None, paused_at: None, completed_at: None, }) } fn record_to_json(record: &SchedulerJobRecord) -> serde_json::Value { json!({ "id": record.id, "kind": record.kind, "schedule": record.schedule, "target": record.target, "payload": record.payload, "enabled": record.enabled, "state": record.state, "last_status": record.last_status, "last_error": record.last_error, "run_count": record.run_count, "max_runs": record.max_runs, "last_fired_at": record.last_fired_at, "next_fire_at": record.next_fire_at, "paused_at": record.paused_at, "completed_at": record.completed_at, "created_at": record.created_at, "updated_at": record.updated_at, }) } fn record_to_upsert(record: &SchedulerJobRecord) -> SchedulerJobUpsert { SchedulerJobUpsert { id: record.id.clone(), kind: record.kind.clone(), schedule: record.schedule.clone(), interval_secs: record.interval_secs, startup_delay_secs: record.startup_delay_secs, target: record.target.clone(), payload: record.payload.clone(), enabled: record.enabled, state: record.state.clone(), last_status: record.last_status.clone(), last_error: record.last_error.clone(), run_count: record.run_count, max_runs: record.max_runs, last_fired_at: record.last_fired_at, next_fire_at: record.next_fire_at, paused_at: record.paused_at, completed_at: record.completed_at, } } fn require_str<'a>(args: &'a serde_json::Value, key: &str) -> anyhow::Result<&'a str> { args.get(key) .and_then(|value| value.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required parameter: {}", key)) } fn error_result(message: &str) -> ToolResult { ToolResult { success: false, output: String::new(), error: Some(message.to_string()), } } fn current_timestamp() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_millis() as i64 } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_scheduler_manage_put_and_get() { let store = Arc::new(SessionStore::in_memory().unwrap()); let tool = SchedulerManageTool::new(store); let put_result = tool .execute(json!({ "action": "put", "id": "heartbeat", "kind": "outbound_message", "schedule": { "type": "interval", "seconds": 60 }, "target": { "channel": "feishu", "chat_id": "oc_demo" }, "payload": { "content": "ping" } })) .await .unwrap(); assert!(put_result.success); let get_result = tool .execute(json!({ "action": "get", "id": "heartbeat" })) .await .unwrap(); assert!(get_result.success); assert!(get_result.output.contains("heartbeat")); assert!(get_result.output.contains("outbound_message")); } #[tokio::test] async fn test_scheduler_manage_put_agent_task() { let store = Arc::new(SessionStore::in_memory().unwrap()); let tool = SchedulerManageTool::new(store); let put_result = tool .execute(json!({ "action": "put", "id": "agent.daily_summary", "kind": "agent_task", "schedule": { "type": "cron", "expression": "0 9 * * *" }, "target": { "channel": "feishu", "chat_id": "oc_demo" }, "payload": { "prompt": "请总结今天待办" } })) .await .unwrap(); assert!(put_result.success); assert!(put_result.output.contains("agent_task")); } }