gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: register sandbox jobs in ContextManager for query tool visibility (#1426)

Browse Source 
* fix: register sandbox jobs in ContextManager for query tool visibility

Sandbox jobs created via execute_sandbox() were persisted to the database
but never registered in the in-memory ContextManager. Since all query tools
(list_jobs, job_status, job_events, cancel_job) only search the
ContextManager, sandbox jobs were invisible to the agent despite running
successfully in Docker containers.

Changes:
- Add register_sandbox_job() to ContextManager (pre-determined UUID,
  starts InProgress, respects max_jobs)
- Extract insert_context() helper to deduplicate create_job_for_user
  and register_sandbox_job
- Add update_context_state / update_context_state_async to sync
  ContextManager state on sandbox job completion/failure
- Extend job_monitor with spawn_job_monitor_with_context() and
  spawn_completion_watcher() so fire-and-forget jobs transition out
  of InProgress when the container finishes
- Make CancelJobTool sandbox-aware (stops container + updates DB)
- Wire sandbox deps into CancelJobTool in register_job_tools()
- 8 regression tests across context manager and job monitor

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: add missing allow_always field in PendingApproval test literal

Upstream commit 09e1c97 added the allow_always field to PendingApproval
but missed updating the test struct literal, breaking compilation.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Vincent Leraitre
2026-03-19 23:22:34 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent b952d229f9
commit e82f4bd2e5
4 changed files with 488 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/job_monitor.rs
+224
View File
@@ -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<IncomingMessage>,
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<IncomingMessage>,
route: JobMonitorRoute,
context_manager: Option<Arc<ContextManager>>,
) -> 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<ContextManager>,
) -> 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::<IncomingMessage>(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::<IncomingMessage>(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);
}
}
src/context/manager.rs
+123 -10
View File
@@ -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<String>,
description: impl Into<String>,
) -> Result<Uuid, JobError> {
// 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<String>,
title: impl Into<String>,
description: impl Into<String>,
) -> 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);
}
}
src/tools/builtin/job.rs
+133 -5
View File
@@ -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<String>,
) {
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<String>) {
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<ContextManager>,
job_manager: Option<Arc<ContainerJobManager>>,
store: Option<Arc<dyn Database>>,
}
impl CancelJobTool {
pub fn new(context_manager: Arc<ContextManager>) -> 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<ContainerJobManager>,
store: Option<Arc<dyn Database>>,
) -> 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",
src/tools/registry.rs
+8 -1
View File
@@ -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;