diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 5e4bf01a..3923530f 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -179,27 +179,33 @@ impl Scheduler { }) .unwrap_or(self.config.max_tokens_per_job); - // Apply both metadata and token budget in one closure (Issue #813: atomic update) - if let Some(meta) = metadata { + // Apply both metadata and token budget in one closure (Issue #813: atomic update). + // Use update_context_and_get to ensure atomicity: no gap where concurrent workers + // can modify the context between update and DB persist (Issue #807). + let ctx = if let Some(meta) = metadata { self.context_manager - .update_context(job_id, |ctx| { + .update_context_and_get(job_id, |ctx| { ctx.metadata = meta; if max_tokens > 0 { ctx.max_tokens = max_tokens; } }) - .await?; + .await? } else if max_tokens > 0 { self.context_manager - .update_context(job_id, |ctx| { + .update_context_and_get(job_id, |ctx| { ctx.max_tokens = max_tokens; }) - .await?; - } + .await? + } else { + // No metadata or token budget to set; get the initial context + self.context_manager.get_context(job_id).await? + }; - // Persist to DB before scheduling so the worker's FK references are valid + // Persist to DB before scheduling so the worker's FK references are valid. + // The context was read under the same lock as the update (atomic), preventing + // concurrent worker interference (Issue #807: non-transactional context updates). 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}"), @@ -832,6 +838,24 @@ mod tests { ); } + #[tokio::test] + async fn test_dispatch_job_no_metadata_no_user_tokens_edge_case() { + // Edge case coverage: when metadata=None AND max_tokens=0 (config), + // the else branch calls get_context() directly (not update_context_and_get). + // This test verifies that path works correctly (Issue #807: full branch coverage). + let sched = make_test_scheduler(0); // 0 = unlimited, but user provides None + let job_id = sched + .dispatch_job("user1", "test", "desc", None) // None metadata + .await + .unwrap(); // safety: test code + + let ctx = sched.context_manager.get_context(job_id).await.unwrap(); // safety: test code + // No metadata was set, should have default empty metadata + assert!(ctx.metadata.is_null() || ctx.metadata == serde_json::json!({})); // safety: test code + // No user tokens AND unlimited config means max_tokens stays at default + assert_eq!(ctx.max_tokens, 0, "unlimited config"); // safety: test code + } + #[test] fn test_scheduler_creation() { // Would need to mock dependencies for proper testing diff --git a/src/context/manager.rs b/src/context/manager.rs index 407a0eea..764f189a 100644 --- a/src/context/manager.rs +++ b/src/context/manager.rs @@ -87,6 +87,28 @@ impl ContextManager { Ok(f(context)) } + /// Atomically update a job context and return the updated context. + /// + /// This method holds the write lock for the entire update-and-read sequence, + /// preventing concurrent workers from interleaving modifications between the + /// update and the subsequent read (Issue #807: non-transactional context updates). + /// Use this when you need to update context and immediately persist it to DB. + pub async fn update_context_and_get( + &self, + job_id: Uuid, + f: F, + ) -> Result + where + F: FnOnce(&mut JobContext), + { + let mut contexts = self.contexts.write().await; + let context = contexts + .get_mut(&job_id) + .ok_or(JobError::NotFound { id: job_id })?; + f(context); + Ok(context.clone()) + } + /// Get job memory. pub async fn get_memory(&self, job_id: Uuid) -> Result { self.memories @@ -877,4 +899,70 @@ mod tests { assert_eq!(manager.all_jobs().await.len(), 10); } + + #[tokio::test] + async fn update_context_and_get_atomicity_regression_issue_807() { + // Regression test for Issue #807: non-transactional context updates. + // Verify that update_context_and_get returns the exact state that was set, + // without allowing concurrent workers to interleave modifications. + let manager = std::sync::Arc::new(ContextManager::new(100)); + let job_id = manager + .create_job("Atomicity Test", "verify no race condition") + .await + .unwrap(); // safety: test code + + // Update and get atomically, setting metadata + let metadata = serde_json::json!({ "priority": "high", "user_id": 42 }); + let returned_ctx = manager + .update_context_and_get(job_id, |ctx| { + ctx.metadata = metadata.clone(); + ctx.max_tokens = 5000; + }) + .await + .unwrap(); // safety: test code + + // Verify the returned context has the exact updates we set + assert_eq!(returned_ctx.metadata, metadata); // safety: test code + assert_eq!(returned_ctx.max_tokens, 5000); // safety: test code + + // Verify a fresh get returns the same state + let fresh_ctx = manager.get_context(job_id).await.unwrap(); // safety: test code + assert_eq!(fresh_ctx.metadata, metadata); // safety: test code + assert_eq!(fresh_ctx.max_tokens, 5000); // safety: test code + } + + #[tokio::test] + async fn update_context_and_get_no_concurrent_interleave() { + // Verify that concurrent updates cannot interleave during update_context_and_get. + // If the lock were released too early, a concurrent state transition could + // get mixed into the returned context. + let manager = std::sync::Arc::new(ContextManager::new(100)); + let job_id = manager + .create_job("Concurrent Race Test", "ensure atomicity") + .await + .unwrap(); // safety: test code + + let metadata = serde_json::json!({ "test": "race_condition" }); + let metadata_clone = metadata.clone(); + + // Spawn a task that will update_context_and_get + let mgr1 = std::sync::Arc::clone(&manager); + let returned_ctx_handle = tokio::spawn(async move { + mgr1.update_context_and_get(job_id, |ctx| { + ctx.metadata = metadata_clone; + ctx.max_tokens = 3000; + }) + .await + }); + + // The returned context should have *only* the metadata update, not any + // concurrent state transitions that might happen during the operation. + let returned_ctx = returned_ctx_handle.await.unwrap().unwrap(); // safety: test code + + // Verify atomicity: returned context has the metadata we set + assert_eq!(returned_ctx.metadata, metadata); // safety: test code + assert_eq!(returned_ctx.max_tokens, 3000); // safety: test code + // And it's in the initial state (Pending), not modified by concurrent workers + assert_eq!(returned_ctx.state, crate::context::JobState::Pending); // safety: test code + } }