diff --git a/src/agent/job_monitor.rs b/src/agent/job_monitor.rs index 6497861a..3f038764 100644 --- a/src/agent/job_monitor.rs +++ b/src/agent/job_monitor.rs @@ -14,12 +14,15 @@ //! Agent Loop //! ``` +use std::sync::Arc; + use tokio::sync::{broadcast, mpsc}; use tokio::task::JoinHandle; use uuid::Uuid; use crate::channels::IncomingMessage; use crate::channels::web::types::SseEvent; +use crate::context::{ContextManager, JobState}; /// Route context for forwarding job monitor events back to the user's channel. #[derive(Debug, Clone)] @@ -40,10 +43,23 @@ pub struct JobMonitorRoute { /// Tool use/result and status events are intentionally skipped (too noisy for /// the main agent's context window). pub fn spawn_job_monitor( + job_id: Uuid, + event_rx: broadcast::Receiver<(Uuid, SseEvent)>, + inject_tx: mpsc::Sender, + route: JobMonitorRoute, +) -> JoinHandle<()> { + spawn_job_monitor_with_context(job_id, event_rx, inject_tx, route, None) +} + +/// Like `spawn_job_monitor`, but also transitions the job's in-memory state +/// when it receives a `JobResult` event. This ensures fire-and-forget sandbox +/// jobs don't stay `InProgress` forever in the `ContextManager`. +pub fn spawn_job_monitor_with_context( job_id: Uuid, mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>, inject_tx: mpsc::Sender, route: JobMonitorRoute, + context_manager: Option>, ) -> JoinHandle<()> { let short_id = job_id.to_string()[..8].to_string(); @@ -77,6 +93,26 @@ pub fn spawn_job_monitor( } } SseEvent::JobResult { status, .. } => { + // Transition in-memory state so the job frees its + // max_jobs slot and query tools show the final state. + if let Some(ref cm) = context_manager { + let target = if status == "completed" { + JobState::Completed + } else { + JobState::Failed + }; + let reason = if status != "completed" { + Some(format!("Container finished: {}", status)) + } else { + None + }; + let _ = cm + .update_context(job_id, |ctx| { + let _ = ctx.transition_to(target, reason); + }) + .await; + } + let mut msg = IncomingMessage::new( route.channel.clone(), route.user_id.clone(), @@ -121,6 +157,62 @@ pub fn spawn_job_monitor( }) } +/// Lightweight watcher that only transitions ContextManager state on job +/// completion. Used when monitor routing metadata is absent (no channel to +/// inject messages into) but we still need to free the `max_jobs` slot. +pub fn spawn_completion_watcher( + job_id: Uuid, + mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>, + context_manager: Arc, +) -> JoinHandle<()> { + let short_id = job_id.to_string()[..8].to_string(); + + tokio::spawn(async move { + loop { + match event_rx.recv().await { + Ok((ev_job_id, SseEvent::JobResult { status, .. })) if ev_job_id == job_id => { + let target = if status == "completed" { + JobState::Completed + } else { + JobState::Failed + }; + let reason = if status != "completed" { + Some(format!("Container finished: {}", status)) + } else { + None + }; + let _ = context_manager + .update_context(job_id, |ctx| { + let _ = ctx.transition_to(target, reason); + }) + .await; + tracing::debug!( + job_id = %short_id, + status = %status, + "Completion watcher exiting (job finished)" + ); + break; + } + Ok(_) => {} + Err(broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + job_id = %short_id, + skipped = n, + "Completion watcher lagged" + ); + } + Err(broadcast::error::RecvError::Closed) => { + tracing::debug!( + job_id = %short_id, + "Broadcast channel closed, stopping completion watcher" + ); + break; + } + } + } + }) +} + #[cfg(test)] mod tests { use super::*; @@ -294,4 +386,136 @@ mod tests { let msg = IncomingMessage::new("monitor", "system", "test").into_internal(); assert!(msg.is_internal); } + + // === Regression: fire-and-forget sandbox jobs must transition out of InProgress === + // Before this fix, spawn_job_monitor only forwarded SSE messages but never + // updated ContextManager. Background sandbox jobs stayed InProgress forever, + // permanently consuming a max_jobs slot. + + #[tokio::test] + async fn test_monitor_transitions_context_on_completion() { + use crate::context::{ContextManager, JobState}; + + let cm = Arc::new(ContextManager::new(5)); + let job_id = Uuid::new_v4(); + cm.register_sandbox_job(job_id, "user-1", "Build app", "desc") + .await + .unwrap(); + + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let handle = spawn_job_monitor_with_context( + job_id, + event_tx.subscribe(), + inject_tx, + test_route(), + Some(Arc::clone(&cm)), + ); + + // Send completion event + event_tx + .send(( + job_id, + SseEvent::JobResult { + job_id: job_id.to_string(), + status: "completed".to_string(), + session_id: None, + }, + )) + .unwrap(); + + // Drain the injected message + let _ = tokio::time::timeout(std::time::Duration::from_secs(1), inject_rx.recv()).await; + + // Wait for monitor to exit + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("monitor should exit") + .expect("monitor should not panic"); + + // Job should now be Completed, not InProgress + let ctx = cm.get_context(job_id).await.unwrap(); + assert_eq!(ctx.state, JobState::Completed); + } + + #[tokio::test] + async fn test_monitor_transitions_context_on_failure() { + use crate::context::{ContextManager, JobState}; + + let cm = Arc::new(ContextManager::new(5)); + let job_id = Uuid::new_v4(); + cm.register_sandbox_job(job_id, "user-1", "Build app", "desc") + .await + .unwrap(); + + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let (inject_tx, mut inject_rx) = mpsc::channel::(16); + + let handle = spawn_job_monitor_with_context( + job_id, + event_tx.subscribe(), + inject_tx, + test_route(), + Some(Arc::clone(&cm)), + ); + + // Send failure event + event_tx + .send(( + job_id, + SseEvent::JobResult { + job_id: job_id.to_string(), + status: "failed".to_string(), + session_id: None, + }, + )) + .unwrap(); + + let _ = tokio::time::timeout(std::time::Duration::from_secs(1), inject_rx.recv()).await; + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("monitor should exit") + .expect("monitor should not panic"); + + let ctx = cm.get_context(job_id).await.unwrap(); + assert_eq!(ctx.state, JobState::Failed); + } + + // === Regression: completion watcher (no route metadata) === + // When monitor_route_from_ctx() returns None, spawn_completion_watcher + // must still transition the job so the max_jobs slot is freed. + + #[tokio::test] + async fn test_completion_watcher_transitions_on_result() { + use crate::context::{ContextManager, JobState}; + + let cm = Arc::new(ContextManager::new(5)); + let job_id = Uuid::new_v4(); + cm.register_sandbox_job(job_id, "user-1", "Build app", "desc") + .await + .unwrap(); + + let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16); + let handle = spawn_completion_watcher(job_id, event_tx.subscribe(), Arc::clone(&cm)); + + event_tx + .send(( + job_id, + SseEvent::JobResult { + job_id: job_id.to_string(), + status: "completed".to_string(), + session_id: None, + }, + )) + .unwrap(); + + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("watcher should exit") + .expect("watcher should not panic"); + + let ctx = cm.get_context(job_id).await.unwrap(); + assert_eq!(ctx.state, JobState::Completed); + } } diff --git a/src/context/manager.rs b/src/context/manager.rs index f9bfedca..28343003 100644 --- a/src/context/manager.rs +++ b/src/context/manager.rs @@ -6,7 +6,7 @@ use std::time::Duration; use tokio::sync::RwLock; use uuid::Uuid; -use crate::context::{JobContext, Memory}; +use crate::context::{JobContext, JobState, Memory}; use crate::error::JobError; /// Manages contexts for multiple concurrent jobs. @@ -46,12 +46,41 @@ impl ContextManager { title: impl Into, description: impl Into, ) -> Result { - // Hold write lock for the entire check-insert to prevent TOCTOU races - // where two concurrent calls both pass the parallel_count check. + let context = JobContext::with_user(user_id, title, description); + let job_id = context.job_id; + self.insert_context(context).await?; + Ok(job_id) + } + + /// Register a sandbox job with a pre-determined ID. + /// + /// Unlike `create_job_for_user` (which generates its own UUID), this method + /// accepts an existing `job_id` — used by `execute_sandbox()` which creates + /// the UUID before the container so it can be shared with Docker labels and + /// DB persistence. + /// + /// The job starts in `InProgress` state since the container is about to be + /// created. Counts against `max_jobs` like any other job. + pub async fn register_sandbox_job( + &self, + job_id: Uuid, + user_id: impl Into, + title: impl Into, + description: impl Into, + ) -> Result<(), JobError> { + let mut context = JobContext::with_user(user_id, title, description); + context.job_id = job_id; + context.state = JobState::InProgress; + context.started_at = Some(chrono::Utc::now()); + self.insert_context(context).await + } + + /// Check max_jobs limit, insert context, and allocate memory. + /// + /// Holds the write lock for the entire check-insert to prevent TOCTOU + /// races where two concurrent calls both pass the parallel_count check. + async fn insert_context(&self, context: JobContext) -> Result<(), JobError> { let mut contexts = self.contexts.write().await; - // Only count jobs that consume execution slots (Pending, InProgress, Stuck). - // Completed and Submitted jobs are no longer actively executing and shouldn't - // block new job creation. let parallel_count = contexts .values() .filter(|c| c.state.is_parallel_blocking()) @@ -61,15 +90,16 @@ impl ContextManager { return Err(JobError::MaxJobsExceeded { max: self.max_jobs }); } - let context = JobContext::with_user(user_id, title, description); let job_id = context.job_id; contexts.insert(job_id, context); drop(contexts); - let memory = Memory::new(job_id); - self.memories.write().await.insert(job_id, memory); + self.memories + .write() + .await + .insert(job_id, Memory::new(job_id)); - Ok(job_id) + Ok(()) } /// Get a job context by ID. @@ -1262,4 +1292,87 @@ mod tests { } } } + + // === Regression: sandbox jobs must be visible to query tools === + // Before the fix, execute_sandbox() only persisted to DB but never + // registered in ContextManager, making sandbox jobs invisible to + // list_jobs, job_status, job_events, and resolve_job_id. + + #[tokio::test] + async fn register_sandbox_job_visible_to_queries() { + let manager = ContextManager::new(5); + let job_id = Uuid::new_v4(); + + manager + .register_sandbox_job( + job_id, + "user-42", + "Run tests", + "Execute test suite in sandbox", + ) + .await + .unwrap(); + + // Job should be retrievable by ID (used by job_status, job_events) + let ctx = manager.get_context(job_id).await.unwrap(); + assert_eq!(ctx.job_id, job_id); + assert_eq!(ctx.user_id, "user-42"); + assert_eq!(ctx.title, "Run tests"); + assert_eq!(ctx.state, JobState::InProgress); + assert!(ctx.started_at.is_some()); + + // Job should appear in all_jobs (used by resolve_job_id prefix matching) + let all = manager.all_jobs().await; + assert!(all.contains(&job_id)); + + // Job should appear in user-scoped listing (used by list_jobs) + let user_jobs = manager.all_jobs_for("user-42").await; + assert!(user_jobs.contains(&job_id)); + + // Job should appear in active jobs listing + let active = manager.active_jobs_for("user-42").await; + assert!(active.contains(&job_id)); + } + + #[tokio::test] + async fn register_sandbox_job_respects_max_jobs() { + let manager = ContextManager::new(2); + + // Fill up the slots with sandbox jobs + manager + .register_sandbox_job(Uuid::new_v4(), "user-1", "Job 1", "desc") + .await + .unwrap(); + manager + .register_sandbox_job(Uuid::new_v4(), "user-1", "Job 2", "desc") + .await + .unwrap(); + + // Third should fail + let result = manager + .register_sandbox_job(Uuid::new_v4(), "user-1", "Job 3", "desc") + .await; + assert!(matches!(result, Err(JobError::MaxJobsExceeded { max: 2 }))); + } + + #[tokio::test] + async fn register_sandbox_job_transitions_correctly() { + let manager = ContextManager::new(5); + let job_id = Uuid::new_v4(); + + manager + .register_sandbox_job(job_id, "user-1", "Task", "desc") + .await + .unwrap(); + + // Should be able to transition InProgress -> Completed + manager + .update_context(job_id, |ctx| ctx.transition_to(JobState::Completed, None)) + .await + .unwrap() + .unwrap(); + + let ctx = manager.get_context(job_id).await.unwrap(); + assert_eq!(ctx.state, JobState::Completed); + } } diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index ea7e5305..0933ee40 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -225,6 +225,41 @@ impl CreateJobTool { } } + /// Transition a sandbox job's state in the ContextManager (awaited). + /// + /// Best-effort: logs on failure (job may have been cleaned up already). + async fn update_context_state_async( + &self, + job_id: Uuid, + state: JobState, + reason: Option, + ) { + if let Err(e) = self + .context_manager + .update_context(job_id, |ctx| { + let _ = ctx.transition_to(state, reason); + }) + .await + { + tracing::debug!(job_id = %job_id, "sandbox context update skipped: {}", e); + } + } + + /// Fire-and-forget variant for use in sync contexts (e.g. `.map_err()` closures). + fn update_context_state(&self, job_id: Uuid, state: JobState, reason: Option) { + let cm = self.context_manager.clone(); + tokio::spawn(async move { + if let Err(e) = cm + .update_context(job_id, |ctx| { + let _ = ctx.transition_to(state, reason); + }) + .await + { + tracing::debug!(job_id = %job_id, "sandbox context update skipped: {}", e); + } + }); + } + /// Update sandbox job status in DB (fire-and-forget). fn update_status( &self, @@ -354,6 +389,16 @@ impl CreateJobTool { } }; + // Register in ContextManager so query tools (list_jobs, job_status, + // job_events, cancel_job) can find sandbox jobs. Without this, sandbox + // jobs exist only in the DB and are invisible to the agent. + self.context_manager + .register_sandbox_job(job_id, &ctx.user_id, task, task) + .await + .map_err(|e| { + ToolError::ExecutionFailed(format!("failed to register sandbox job: {}", e)) + })?; + // Persist the job to DB before creating the container. self.persist_job(SandboxJobRecord { id: job_id, @@ -397,6 +442,7 @@ impl CreateJobTool { None, Some(Utc::now()), ); + self.update_context_state(job_id, JobState::Failed, Some(e.to_string())); ToolError::ExecutionFailed(format!("failed to create container: {}", e)) })?; @@ -416,16 +462,20 @@ impl CreateJobTool { // monitor terminates. No JoinHandle is retained. if let (Some(etx), Some(itx)) = (&self.event_tx, &self.inject_tx) { if let Some(route) = monitor_route_from_ctx(ctx) { - crate::agent::job_monitor::spawn_job_monitor( + crate::agent::job_monitor::spawn_job_monitor_with_context( job_id, etx.subscribe(), itx.clone(), route, + Some(self.context_manager.clone()), ); } else { - tracing::debug!( - job_id = %job_id, - "Skipping job monitor injection due to missing route metadata" + // No routing metadata — can't inject messages, but still + // need to transition the job out of InProgress when done. + crate::agent::job_monitor::spawn_completion_watcher( + job_id, + etx.subscribe(), + self.context_manager.clone(), ); } } @@ -457,6 +507,12 @@ impl CreateJobTool { None, Some(Utc::now()), ); + self.update_context_state_async( + job_id, + JobState::Failed, + Some("Timed out (10 minutes)".to_string()), + ) + .await; return Err(ToolError::ExecutionFailed( "container execution timed out (10 minutes)".to_string(), )); @@ -491,6 +547,8 @@ impl CreateJobTool { None, Some(finished_at), ); + self.update_context_state_async(job_id, JobState::Completed, None) + .await; let result = serde_json::json!({ "job_id": job_id.to_string(), "status": "completed", @@ -508,6 +566,12 @@ impl CreateJobTool { None, Some(finished_at), ); + self.update_context_state_async( + job_id, + JobState::Failed, + Some(message.clone()), + ) + .await; return Err(ToolError::ExecutionFailed(format!( "container job failed: {}", message @@ -529,6 +593,12 @@ impl CreateJobTool { None, Some(Utc::now()), ); + self.update_context_state_async( + job_id, + JobState::Failed, + Some(message.clone()), + ) + .await; return Err(ToolError::ExecutionFailed(format!( "container job failed: {}", message @@ -544,6 +614,8 @@ impl CreateJobTool { None, Some(Utc::now()), ); + self.update_context_state_async(job_id, JobState::Completed, None) + .await; let result = serde_json::json!({ "job_id": job_id.to_string(), "status": "completed", @@ -1025,13 +1097,34 @@ impl Tool for JobStatusTool { } /// Tool for canceling a job. +/// +/// For sandbox jobs (registered via `register_sandbox_job`), cancellation also +/// stops the Docker container and updates the DB status — matching the behavior +/// of the web cancellation handler in `channels/web/handlers/jobs.rs`. pub struct CancelJobTool { context_manager: Arc, + job_manager: Option>, + store: Option>, } impl CancelJobTool { pub fn new(context_manager: Arc) -> Self { - Self { context_manager } + Self { + context_manager, + job_manager: None, + store: None, + } + } + + /// Inject sandbox dependencies so cancellation also stops containers. + pub fn with_sandbox( + mut self, + job_manager: Arc, + store: Option>, + ) -> Self { + self.job_manager = Some(job_manager); + self.store = store; + self } } @@ -1081,6 +1174,41 @@ impl Tool for CancelJobTool { .await { Ok(Ok(())) => { + // Stop the sandbox container if one exists for this job. + if let Some(ref jm) = self.job_manager + && let Err(e) = jm.stop_job(job_id).await + { + tracing::warn!( + job_id = %job_id, + "Failed to stop container during cancellation: {}", e + ); + } + + // Update DB status for sandbox jobs. Uses "failed" (not + // "cancelled") to match the web cancel handler convention — + // the sandbox DB schema treats cancellation as a failure variant. + if let Some(ref store) = self.store { + let store = store.clone(); + tokio::spawn(async move { + if let Err(e) = store + .update_sandbox_job_status( + job_id, + "failed", + Some(false), + Some("Cancelled by user"), + None, + Some(Utc::now()), + ) + .await + { + tracing::warn!( + job_id = %job_id, + "Failed to update sandbox job status on cancel: {}", e + ); + } + }); + } + let result = serde_json::json!({ "job_id": job_id.to_string(), "status": "cancelled", diff --git a/src/tools/registry.rs b/src/tools/registry.rs index a68e300b..c64b637f 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -367,6 +367,9 @@ impl ToolRegistry { if let Some(slot) = scheduler_slot { create_tool = create_tool.with_scheduler_slot(slot); } + // Clone before moving into create_tool so cancel_job can also use them. + let jm_for_cancel = job_manager.clone(); + let store_for_cancel = store.clone(); if let Some(jm) = job_manager { create_tool = create_tool.with_sandbox(jm, store.clone()); } @@ -379,7 +382,11 @@ impl ToolRegistry { self.register_sync(Arc::new(create_tool)); 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(Arc::clone(&context_manager)))); + let mut cancel_tool = CancelJobTool::new(Arc::clone(&context_manager)); + if let Some(jm) = jm_for_cancel { + cancel_tool = cancel_tool.with_sandbox(jm, store_for_cancel); + } + self.register_sync(Arc::new(cancel_tool)); // Base tools: create, list, status, cancel let mut job_tool_count = 4;