//! Job scheduler for parallel execution. use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use tokio::sync::{RwLock, mpsc, oneshot}; use tokio::task::JoinHandle; use uuid::Uuid; use crate::agent::task::{Task, TaskContext, TaskOutput}; use crate::agent::worker::{Worker, WorkerDeps}; use crate::channels::web::types::SseEvent; use crate::config::AgentConfig; use crate::context::{ContextManager, JobContext, JobState}; use crate::db::Database; use crate::error::{Error, JobError}; use crate::hooks::HookRegistry; use crate::llm::LlmProvider; use crate::safety::SafetyLayer; use crate::tools::{ApprovalContext, ToolRegistry}; /// Message to send to a worker. #[derive(Debug)] pub enum WorkerMessage { /// Start working on the job. Start, /// Stop the job. Stop, /// Check health. Ping, /// Inject a follow-up user message into the worker's reasoning context. UserMessage(String), } /// Status of a scheduled job. #[derive(Debug)] pub struct ScheduledJob { pub handle: JoinHandle<()>, pub tx: mpsc::Sender, } /// Status of a scheduled sub-task. struct ScheduledSubtask { handle: JoinHandle>, } /// Schedules and manages parallel job execution. pub struct Scheduler { config: AgentConfig, context_manager: Arc, llm: Arc, safety: Arc, tools: Arc, store: Option>, hooks: Arc, /// SSE broadcast sender for live job event streaming. sse_tx: Option>, /// HTTP interceptor for trace recording/replay (propagated to workers). http_interceptor: Option>, /// Running jobs (main LLM-driven jobs). jobs: Arc>>, /// Running sub-tasks (tool executions, background tasks). subtasks: Arc>>, } impl Scheduler { /// Create a new scheduler. pub fn new( config: AgentConfig, context_manager: Arc, llm: Arc, safety: Arc, tools: Arc, store: Option>, hooks: Arc, ) -> Self { Self { config, context_manager, llm, safety, tools, store, hooks, sse_tx: None, http_interceptor: None, jobs: Arc::new(RwLock::new(HashMap::new())), subtasks: Arc::new(RwLock::new(HashMap::new())), } } /// Set the SSE broadcast sender for live job event streaming. pub fn set_sse_sender(&mut self, tx: tokio::sync::broadcast::Sender) { self.sse_tx = Some(tx); } /// Set the HTTP interceptor for trace recording/replay. pub fn set_http_interceptor( &mut self, interceptor: Arc, ) { self.http_interceptor = Some(interceptor); } /// Create, persist, and schedule a job in one shot. /// /// This is the preferred entry point for dispatching new jobs. It: /// 1. Creates the job context via `ContextManager` /// 2. Optionally applies metadata (e.g. `max_iterations`) /// 3. Persists the job to the database (so FK references from /// `job_actions` / `llm_calls` work immediately) /// 4. Schedules the job for worker execution /// /// Returns the new job ID. pub async fn dispatch_job( &self, user_id: &str, title: &str, description: &str, metadata: Option, ) -> Result { self.dispatch_job_inner(user_id, title, description, metadata, None) .await } /// Dispatch a job with an explicit approval context for autonomous execution. /// /// Same as `dispatch_job`, but the worker will use the given `ApprovalContext` /// to determine which tools are pre-approved (instead of blocking all non-`Never` tools). pub async fn dispatch_job_with_context( &self, user_id: &str, title: &str, description: &str, metadata: Option, approval_context: ApprovalContext, ) -> Result { self.dispatch_job_inner( user_id, title, description, metadata, Some(approval_context), ) .await } /// Shared implementation for `dispatch_job` and `dispatch_job_with_context`. async fn dispatch_job_inner( &self, user_id: &str, title: &str, description: &str, metadata: Option, approval_context: Option, ) -> Result { let job_id = self .context_manager .create_job_for_user(user_id, title, description) .await?; // Apply token budget from config, allowing per-job metadata override. let max_tokens = metadata .as_ref() .and_then(|m| m.get("max_tokens")) .and_then(|v| v.as_u64()) .unwrap_or(self.config.max_tokens_per_job); // Apply metadata if provided if let Some(meta) = metadata { self.context_manager .update_context(job_id, |ctx| { ctx.metadata = meta; }) .await?; } // Set token budget (separate update to avoid overwriting metadata) if max_tokens > 0 { self.context_manager .update_context(job_id, |ctx| { ctx.max_tokens = max_tokens; }) .await?; } // Persist to DB before scheduling so the worker's FK references are valid if let Some(ref store) = self.store { let ctx = self.context_manager.get_context(job_id).await?; store.save_job(&ctx).await.map_err(|e| JobError::Failed { id: job_id, reason: format!("failed to persist job: {e}"), })?; } self.schedule_with_context(job_id, approval_context).await?; Ok(job_id) } /// Schedule a job for execution. pub async fn schedule(&self, job_id: Uuid) -> Result<(), JobError> { self.schedule_with_context(job_id, None).await } /// Schedule a job with an optional approval context. async fn schedule_with_context( &self, job_id: Uuid, approval_context: Option, ) -> Result<(), JobError> { // Hold write lock for the entire check-insert sequence to prevent // TOCTOU races where two concurrent calls both pass the checks. { let mut jobs = self.jobs.write().await; if jobs.contains_key(&job_id) { return Ok(()); } if jobs.len() >= self.config.max_parallel_jobs { return Err(JobError::MaxJobsExceeded { max: self.config.max_parallel_jobs, }); } // Transition job to in_progress self.context_manager .update_context(job_id, |ctx| { ctx.transition_to( JobState::InProgress, Some("Scheduled for execution".to_string()), ) }) .await? .map_err(|s| JobError::ContextError { id: job_id, reason: s, })?; // Create worker channel let (tx, rx) = mpsc::channel(16); // Create worker with shared dependencies let deps = WorkerDeps { context_manager: self.context_manager.clone(), llm: self.llm.clone(), safety: self.safety.clone(), tools: self.tools.clone(), store: self.store.clone(), hooks: self.hooks.clone(), timeout: self.config.job_timeout, use_planning: self.config.use_planning, sse_tx: self.sse_tx.clone(), approval_context, http_interceptor: self.http_interceptor.clone(), }; let worker = Worker::new(job_id, deps); // Spawn worker task let handle = tokio::spawn(async move { if let Err(e) = worker.run(rx).await { tracing::error!("Worker for job {} failed: {}", job_id, e); } }); // Start the worker if tx.send(WorkerMessage::Start).await.is_err() { tracing::error!(job_id = %job_id, "Worker died before receiving Start message"); } // Insert while still holding the write lock jobs.insert(job_id, ScheduledJob { handle, tx }); } // Cleanup task for this job to avoid capacity leaks let jobs = Arc::clone(&self.jobs); tokio::spawn(async move { loop { let finished = { let jobs_read = jobs.read().await; match jobs_read.get(&job_id) { Some(scheduled) => scheduled.handle.is_finished(), None => true, } }; if finished { jobs.write().await.remove(&job_id); break; } tokio::time::sleep(Duration::from_secs(1)).await; } }); tracing::info!("Scheduled job {} for execution", job_id); Ok(()) } /// Schedule a sub-task from within a worker. /// /// Sub-tasks are lightweight tasks that don't go through the full job lifecycle. /// They're used for parallel tool execution and background computations. /// /// Returns a oneshot receiver to get the result. pub async fn spawn_subtask( &self, parent_id: Uuid, task: Task, ) -> Result>, JobError> { let task_id = Uuid::new_v4(); let (result_tx, result_rx) = oneshot::channel(); let handle = match task { Task::Job { .. } => { // Jobs should go through schedule(), not spawn_subtask return Err(JobError::ContextError { id: parent_id, reason: "Use schedule() for Job tasks, not spawn_subtask()".to_string(), }); } Task::ToolExec { parent_id: tool_parent_id, tool_name, params, } => { let tools = self.tools.clone(); let context_manager = self.context_manager.clone(); let safety = self.safety.clone(); // TODO: propagate parent job's ApprovalContext here when subtasks // are used in autonomous/routine paths (currently only used in tests). tokio::spawn(async move { let result = Self::execute_tool_task( tools, context_manager, safety, None, tool_parent_id, &tool_name, params, ) .await; // Send result (ignore if receiver dropped) let _ = result_tx.send(result); }) } Task::Background { id: _, handler } => { let ctx = TaskContext::new(task_id).with_parent(parent_id); tokio::spawn(async move { let result = handler.run(ctx).await; let _ = result_tx.send(result); }) } }; // Track the subtask self.subtasks.write().await.insert( task_id, ScheduledSubtask { handle: tokio::spawn(async move { // Wrap the handle to get its result match handle.await { Ok(()) => Err(Error::Job(JobError::ContextError { id: task_id, reason: "Subtask completed but result not captured".to_string(), })), Err(e) => Err(Error::Job(JobError::ContextError { id: task_id, reason: format!("Subtask panicked: {}", e), })), } }), }, ); // Cleanup task for subtask tracking let subtasks = Arc::clone(&self.subtasks); tokio::spawn(async move { loop { let finished = { let subtasks_read = subtasks.read().await; match subtasks_read.get(&task_id) { Some(scheduled) => scheduled.handle.is_finished(), None => true, } }; if finished { subtasks.write().await.remove(&task_id); break; } tokio::time::sleep(Duration::from_secs(1)).await; } }); tracing::debug!( parent_id = %parent_id, task_id = %task_id, "Spawned subtask" ); Ok(result_rx) } /// Schedule multiple tasks in parallel and wait for all to complete. /// /// Returns results in the same order as the input tasks. pub async fn spawn_batch( &self, parent_id: Uuid, tasks: Vec, ) -> Vec> { if tasks.is_empty() { return Vec::new(); } let mut receivers = Vec::with_capacity(tasks.len()); // Spawn all tasks for task in tasks { match self.spawn_subtask(parent_id, task).await { Ok(rx) => receivers.push(Some(rx)), Err(e) => { // Store the error directly receivers.push(None); tracing::warn!( parent_id = %parent_id, error = %e, "Failed to spawn subtask in batch" ); } } } // Collect results let mut results = Vec::with_capacity(receivers.len()); for rx in receivers { let result = match rx { Some(receiver) => match receiver.await { Ok(task_result) => task_result, Err(_) => Err(Error::Job(JobError::ContextError { id: parent_id, reason: "Subtask channel closed unexpectedly".to_string(), })), }, None => Err(Error::Job(JobError::ContextError { id: parent_id, reason: "Subtask failed to spawn".to_string(), })), }; results.push(result); } results } /// Execute a single tool as a subtask. async fn execute_tool_task( tools: Arc, context_manager: Arc, safety: Arc, approval_context: Option, job_id: Uuid, tool_name: &str, params: serde_json::Value, ) -> Result { let start = std::time::Instant::now(); // Get the tool let tool = tools.get(tool_name).await.ok_or_else(|| { Error::Tool(crate::error::ToolError::NotFound { name: tool_name.to_string(), }) })?; // Get job context let job_ctx: JobContext = context_manager.get_context(job_id).await?; if job_ctx.state == JobState::Cancelled { return Err(crate::error::ToolError::ExecutionFailed { name: tool_name.to_string(), reason: "Job is cancelled".to_string(), } .into()); } let requirement = tool.requires_approval(¶ms); let blocked = ApprovalContext::is_blocked_or_default(&approval_context, tool_name, requirement); if blocked { return Err(crate::error::ToolError::AuthRequired { name: tool_name.to_string(), } .into()); } // Validate tool parameters let validation = safety.validator().validate_tool_params(¶ms); if !validation.is_valid { let details = validation .errors .iter() .map(|e| format!("{}: {}", e.field, e.message)) .collect::>() .join("; "); return Err(crate::error::ToolError::InvalidParameters { name: tool_name.to_string(), reason: format!("Invalid tool parameters: {}", details), } .into()); } // Execute with per-tool timeout let tool_timeout = tool.execution_timeout(); let result = tokio::time::timeout(tool_timeout, async { tool.execute(params, &job_ctx).await }) .await .map_err(|_| { Error::Tool(crate::error::ToolError::Timeout { name: tool_name.to_string(), timeout: tool_timeout, }) })? .map_err(|e| { Error::Tool(crate::error::ToolError::ExecutionFailed { name: tool_name.to_string(), reason: e.to_string(), }) })?; Ok(TaskOutput::new(result.result, start.elapsed())) } /// Stop a running job. pub async fn stop(&self, job_id: Uuid) -> Result<(), JobError> { let mut jobs = self.jobs.write().await; if let Some(scheduled) = jobs.remove(&job_id) { // Send stop signal let _ = scheduled.tx.send(WorkerMessage::Stop).await; // Give it a moment to clean up tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; // Abort if still running if !scheduled.handle.is_finished() { scheduled.handle.abort(); } // Update job state self.context_manager .update_context(job_id, |ctx| { if let Err(e) = ctx.transition_to( JobState::Cancelled, Some("Stopped by scheduler".to_string()), ) { tracing::warn!( job_id = %job_id, error = %e, "Failed to transition job to Cancelled state" ); } }) .await?; // Persist cancellation (fire-and-forget) if let Some(ref store) = self.store { let store = store.clone(); tokio::spawn(async move { if let Err(e) = store .update_job_status( job_id, JobState::Cancelled, Some("Stopped by scheduler"), ) .await { tracing::warn!("Failed to persist cancellation for job {}: {}", job_id, e); } }); } tracing::info!("Stopped job {}", job_id); } Ok(()) } /// Send a follow-up user message to a running job. /// /// Returns `Ok(())` if the message was queued, `Err` if the job is not running. pub async fn send_message(&self, job_id: Uuid, content: String) -> Result<(), JobError> { // Clone the sender while holding the lock, then release before the // async send to avoid blocking scheduler writes during backpressure. let tx = { let jobs = self.jobs.read().await; let scheduled = jobs.get(&job_id).ok_or(JobError::NotFound { id: job_id })?; scheduled.tx.clone() }; tx.send(WorkerMessage::UserMessage(content)) .await .map_err(|_| JobError::Failed { id: job_id, reason: "Worker channel closed".to_string(), })?; Ok(()) } /// Check if a job is running. pub async fn is_running(&self, job_id: Uuid) -> bool { self.jobs.read().await.contains_key(&job_id) } /// Get count of running jobs. pub async fn running_count(&self) -> usize { self.jobs.read().await.len() } /// Get count of running subtasks. pub async fn subtask_count(&self) -> usize { self.subtasks.read().await.len() } /// Get all running job IDs. pub async fn running_jobs(&self) -> Vec { self.jobs.read().await.keys().cloned().collect() } /// Clean up finished jobs and subtasks. pub async fn cleanup_finished(&self) { // Clean up jobs { let mut jobs = self.jobs.write().await; let mut finished = Vec::new(); for (id, scheduled) in jobs.iter() { if scheduled.handle.is_finished() { finished.push(*id); } } for id in finished { jobs.remove(&id); tracing::debug!("Cleaned up finished job {}", id); } } // Clean up subtasks { let mut subtasks = self.subtasks.write().await; let mut finished = Vec::new(); for (id, scheduled) in subtasks.iter() { if scheduled.handle.is_finished() { finished.push(*id); } } for id in finished { subtasks.remove(&id); tracing::trace!("Cleaned up finished subtask {}", id); } } } /// Stop all jobs. pub async fn stop_all(&self) { let job_ids: Vec = self.jobs.read().await.keys().cloned().collect(); for job_id in job_ids { let _ = self.stop(job_id).await; } // Abort all subtasks let mut subtasks = self.subtasks.write().await; for (_, scheduled) in subtasks.drain() { scheduled.handle.abort(); } } /// Get access to the tools registry. pub fn tools(&self) -> &Arc { &self.tools } /// Get access to the context manager. pub fn context_manager(&self) -> &Arc { &self.context_manager } } #[cfg(test)] mod tests { use super::*; use crate::config::SafetyConfig; use crate::safety::SafetyLayer; use crate::tools::{ApprovalRequirement, Tool, ToolError, ToolOutput}; #[test] fn test_scheduler_creation() { // Would need to mock dependencies for proper testing } #[tokio::test] async fn test_spawn_batch_empty() { // This test would need mock dependencies. // For now just verify the empty case doesn't panic. } /// A tool that returns `UnlessAutoApproved`. struct SoftApprovalTool; #[async_trait::async_trait] impl Tool for SoftApprovalTool { fn name(&self) -> &str { "soft_gate" } fn description(&self) -> &str { "needs soft approval" } fn parameters_schema(&self) -> serde_json::Value { serde_json::json!({"type": "object", "properties": {}}) } async fn execute( &self, _params: serde_json::Value, _ctx: &JobContext, ) -> Result { Ok(ToolOutput::text( "soft_ok", std::time::Instant::now().elapsed(), )) } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::UnlessAutoApproved } fn requires_sanitization(&self) -> bool { false } } /// A tool that returns `Always`. struct HardApprovalTool; #[async_trait::async_trait] impl Tool for HardApprovalTool { fn name(&self) -> &str { "hard_gate" } fn description(&self) -> &str { "needs hard approval" } fn parameters_schema(&self) -> serde_json::Value { serde_json::json!({"type": "object", "properties": {}}) } async fn execute( &self, _params: serde_json::Value, _ctx: &JobContext, ) -> Result { Ok(ToolOutput::text( "hard_ok", std::time::Instant::now().elapsed(), )) } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::Always } fn requires_sanitization(&self) -> bool { false } } async fn setup_tools_and_job() -> ( Arc, Arc, Arc, Uuid, ) { let registry = ToolRegistry::new(); registry.register(Arc::new(SoftApprovalTool)).await; registry.register(Arc::new(HardApprovalTool)).await; let cm = Arc::new(ContextManager::new(5)); let job_id = cm.create_job("test", "approval test").await.unwrap(); cm.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None)) .await .unwrap() .unwrap(); let safety = Arc::new(SafetyLayer::new(&SafetyConfig { max_output_length: 100_000, injection_check_enabled: false, })); (Arc::new(registry), cm, safety, job_id) } #[tokio::test] async fn test_execute_tool_task_blocks_without_context() { let (tools, cm, safety, job_id) = setup_tools_and_job().await; // Without approval context, UnlessAutoApproved is blocked let result = Scheduler::execute_tool_task( tools.clone(), cm.clone(), safety.clone(), None, job_id, "soft_gate", serde_json::json!({}), ) .await; assert!( result.is_err(), "soft_gate should be blocked without context" ); // Always is also blocked let result = Scheduler::execute_tool_task( tools, cm, safety, None, job_id, "hard_gate", serde_json::json!({}), ) .await; assert!( result.is_err(), "hard_gate should be blocked without context" ); } #[tokio::test] async fn test_execute_tool_task_autonomous_unblocks_soft() { let (tools, cm, safety, job_id) = setup_tools_and_job().await; // Autonomous context auto-approves UnlessAutoApproved let result = Scheduler::execute_tool_task( tools.clone(), cm.clone(), safety.clone(), Some(ApprovalContext::autonomous()), job_id, "soft_gate", serde_json::json!({}), ) .await; assert!( result.is_ok(), "soft_gate should pass with autonomous context" ); // But still blocks Always let result = Scheduler::execute_tool_task( tools, cm, safety, Some(ApprovalContext::autonomous()), job_id, "hard_gate", serde_json::json!({}), ) .await; assert!( result.is_err(), "hard_gate should still be blocked without explicit permission" ); } #[tokio::test] async fn test_execute_tool_task_autonomous_with_permissions() { let (tools, cm, safety, job_id) = setup_tools_and_job().await; // Autonomous context with explicit permission for hard_gate let ctx = ApprovalContext::autonomous_with_tools(["hard_gate".to_string()]); let result = Scheduler::execute_tool_task( tools.clone(), cm.clone(), safety.clone(), Some(ctx.clone()), job_id, "soft_gate", serde_json::json!({}), ) .await; assert!(result.is_ok(), "soft_gate should pass"); let result = Scheduler::execute_tool_task( tools, cm, safety, Some(ctx), job_id, "hard_gate", serde_json::json!({}), ) .await; assert!( result.is_ok(), "hard_gate should pass with explicit permission" ); } }