feat(jobs): push background-job completion notifications via per-context queue
- add NotificationQueue/SystemNotification: every context owns a fresh queue (children never inherit the parent's, avoiding first-drainer-wins races between transcripts) - job tasks push job_completed/job_failed events on completion, failure, and timeout; a panic skips the push and is surfaced by the guardrail's finished-handle enumeration and collect's JoinError mapping instead - events for jobs already collected or cancelled are dropped at drain time by filtering against live supervisor registration - replace inject_escalation_notification with single-pass merge_system_channel: pending_escalations (root-only) ordered before system_notifications (any depth) on the last tool result of a batch; byte-identical output when notifications are empty, proven by the unmodified pre-merger characterization tests
This commit is contained in:
@@ -34,6 +34,7 @@ use crate::rag::Rag;
|
|||||||
use crate::supervisor::Supervisor;
|
use crate::supervisor::Supervisor;
|
||||||
use crate::supervisor::escalation::EscalationQueue;
|
use crate::supervisor::escalation::EscalationQueue;
|
||||||
use crate::supervisor::mailbox::Inbox;
|
use crate::supervisor::mailbox::Inbox;
|
||||||
|
use crate::supervisor::notification::NotificationQueue;
|
||||||
use crate::utils::{
|
use crate::utils::{
|
||||||
AbortSignal, abortable_run_with_spinner, edit_file, fuzzy_filter, get_env_name,
|
AbortSignal, abortable_run_with_spinner, edit_file, fuzzy_filter, get_env_name,
|
||||||
list_file_names, now, render_prompt, temp_file,
|
list_file_names, now, render_prompt, temp_file,
|
||||||
@@ -330,6 +331,7 @@ pub struct RequestContext {
|
|||||||
pub self_agent_id: Option<String>,
|
pub self_agent_id: Option<String>,
|
||||||
pub inbox: Option<Arc<Inbox>>,
|
pub inbox: Option<Arc<Inbox>>,
|
||||||
pub escalation_queue: Option<Arc<EscalationQueue>>,
|
pub escalation_queue: Option<Arc<EscalationQueue>>,
|
||||||
|
pub notification_queue: Arc<NotificationQueue>,
|
||||||
pub current_depth: usize,
|
pub current_depth: usize,
|
||||||
pub auto_continue_count: usize,
|
pub auto_continue_count: usize,
|
||||||
pub pending_agents_guardrail_count: u32,
|
pub pending_agents_guardrail_count: u32,
|
||||||
@@ -364,6 +366,7 @@ impl RequestContext {
|
|||||||
self_agent_id: None,
|
self_agent_id: None,
|
||||||
inbox: None,
|
inbox: None,
|
||||||
escalation_queue: None,
|
escalation_queue: None,
|
||||||
|
notification_queue: Arc::new(NotificationQueue::new()),
|
||||||
current_depth: 0,
|
current_depth: 0,
|
||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
pending_agents_guardrail_count: 0,
|
pending_agents_guardrail_count: 0,
|
||||||
@@ -424,6 +427,7 @@ impl RequestContext {
|
|||||||
self_agent_id: None,
|
self_agent_id: None,
|
||||||
inbox: None,
|
inbox: None,
|
||||||
escalation_queue: None,
|
escalation_queue: None,
|
||||||
|
notification_queue: Arc::new(NotificationQueue::new()),
|
||||||
current_depth: 0,
|
current_depth: 0,
|
||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
pending_agents_guardrail_count: 0,
|
pending_agents_guardrail_count: 0,
|
||||||
@@ -471,6 +475,7 @@ impl RequestContext {
|
|||||||
self_agent_id: self.self_agent_id.clone(),
|
self_agent_id: self.self_agent_id.clone(),
|
||||||
inbox: self.inbox.clone(),
|
inbox: self.inbox.clone(),
|
||||||
escalation_queue: self.escalation_queue.clone(),
|
escalation_queue: self.escalation_queue.clone(),
|
||||||
|
notification_queue: self.notification_queue.clone(),
|
||||||
current_depth: self.current_depth,
|
current_depth: self.current_depth,
|
||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
pending_agents_guardrail_count: 0,
|
pending_agents_guardrail_count: 0,
|
||||||
@@ -516,6 +521,7 @@ impl RequestContext {
|
|||||||
self_agent_id: Some(self_agent_id),
|
self_agent_id: Some(self_agent_id),
|
||||||
inbox: Some(inbox),
|
inbox: Some(inbox),
|
||||||
escalation_queue: parent.escalation_queue.clone(),
|
escalation_queue: parent.escalation_queue.clone(),
|
||||||
|
notification_queue: Arc::new(NotificationQueue::new()),
|
||||||
current_depth,
|
current_depth,
|
||||||
auto_continue_count: 0,
|
auto_continue_count: 0,
|
||||||
pending_agents_guardrail_count: 0,
|
pending_agents_guardrail_count: 0,
|
||||||
@@ -4202,6 +4208,7 @@ impl RequestContext {
|
|||||||
self.supervisor = supervisor;
|
self.supervisor = supervisor;
|
||||||
self.inbox = None;
|
self.inbox = None;
|
||||||
self.escalation_queue = None;
|
self.escalation_queue = None;
|
||||||
|
self.notification_queue = Arc::new(NotificationQueue::new());
|
||||||
self.self_agent_id = None;
|
self.self_agent_id = None;
|
||||||
self.parent_supervisor = None;
|
self.parent_supervisor = None;
|
||||||
self.current_depth = 0;
|
self.current_depth = 0;
|
||||||
@@ -4244,6 +4251,7 @@ impl RequestContext {
|
|||||||
self.self_agent_id = None;
|
self.self_agent_id = None;
|
||||||
self.inbox = None;
|
self.inbox = None;
|
||||||
self.escalation_queue = None;
|
self.escalation_queue = None;
|
||||||
|
self.notification_queue = Arc::new(NotificationQueue::new());
|
||||||
self.current_depth = 0;
|
self.current_depth = 0;
|
||||||
self.auto_continue_count = 0;
|
self.auto_continue_count = 0;
|
||||||
self.pending_agents_guardrail_count = 0;
|
self.pending_agents_guardrail_count = 0;
|
||||||
@@ -5463,6 +5471,32 @@ mod tests {
|
|||||||
assert!(ctx.root_escalation_queue().is_none());
|
assert!(ctx.root_escalation_queue().is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn new_for_child_gets_fresh_notification_queue() {
|
||||||
|
let parent = create_test_ctx();
|
||||||
|
let child = RequestContext::new_for_child(
|
||||||
|
Arc::clone(&parent.app),
|
||||||
|
&parent,
|
||||||
|
1,
|
||||||
|
Arc::new(Inbox::new()),
|
||||||
|
"agent_test_1".to_string(),
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!Arc::ptr_eq(&parent.notification_queue, &child.notification_queue),
|
||||||
|
"each child owns its notifications; a shared queue would race drains"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fork_for_branch_shares_notification_queue() {
|
||||||
|
let ctx = create_test_ctx();
|
||||||
|
let branch = ctx.fork_for_branch();
|
||||||
|
assert!(Arc::ptr_eq(
|
||||||
|
&ctx.notification_queue,
|
||||||
|
&branch.notification_queue
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
fn app_state_with_mcp_config(mcp_server_support: bool, server_names: &[&str]) -> Arc<AppState> {
|
fn app_state_with_mcp_config(mcp_server_support: bool, server_names: &[&str]) -> Arc<AppState> {
|
||||||
app_state_with_mcp_command(mcp_server_support, server_names, "echo")
|
app_state_with_mcp_command(mcp_server_support, server_names, "echo")
|
||||||
}
|
}
|
||||||
|
|||||||
+145
-6
@@ -14,6 +14,7 @@ use crate::mcp::{
|
|||||||
MCP_PROMPT_META_FUNCTION_NAME_PREFIX, MCP_READ_META_FUNCTION_NAME_PREFIX,
|
MCP_PROMPT_META_FUNCTION_NAME_PREFIX, MCP_READ_META_FUNCTION_NAME_PREFIX,
|
||||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||||
};
|
};
|
||||||
|
use crate::supervisor::notification::job_notification;
|
||||||
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor};
|
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor};
|
||||||
use crate::utils::{create_abort_signal, muted_warning_text, temp_file, wait_abort_signal};
|
use crate::utils::{create_abort_signal, muted_warning_text, temp_file, wait_abort_signal};
|
||||||
|
|
||||||
@@ -432,27 +433,39 @@ async fn handle_start(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
|||||||
current_depth: ctx.current_depth,
|
current_depth: ctx.current_depth,
|
||||||
};
|
};
|
||||||
let task_state = Arc::clone(&state);
|
let task_state = Arc::clone(&state);
|
||||||
|
let task_notifications = Arc::clone(&ctx.notification_queue);
|
||||||
|
let notify_id = job_id.clone();
|
||||||
|
let notify_tool = tool.clone();
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let result = run_mcp_job(job_ctx, server, inner_tool, inner_args).await;
|
let result = run_mcp_job(job_ctx, server, inner_tool, inner_args).await;
|
||||||
task_state.lock().status = match &result {
|
let success = result.is_ok();
|
||||||
Ok(_) => JobStatus::Completed,
|
task_state.lock().status = if success {
|
||||||
Err(_) => JobStatus::Failed,
|
JobStatus::Completed
|
||||||
|
} else {
|
||||||
|
JobStatus::Failed
|
||||||
};
|
};
|
||||||
|
task_notifications.push(job_notification(¬ify_id, ¬ify_tool, success));
|
||||||
result
|
result
|
||||||
})
|
})
|
||||||
} else {
|
} else {
|
||||||
let snapshot = build_env_snapshot(ctx, &tool, &arguments)?;
|
let snapshot = build_env_snapshot(ctx, &tool, &arguments)?;
|
||||||
let task_state = Arc::clone(&state);
|
let task_state = Arc::clone(&state);
|
||||||
let task_buf = Arc::clone(&output_buf);
|
let task_buf = Arc::clone(&output_buf);
|
||||||
|
let task_notifications = Arc::clone(&ctx.notification_queue);
|
||||||
|
let notify_id = job_id.clone();
|
||||||
|
let notify_tool = tool.clone();
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let result = run_process_job(snapshot, Arc::clone(&task_state), task_buf).await;
|
let result = run_process_job(snapshot, Arc::clone(&task_state), task_buf).await;
|
||||||
|
let success = matches!(&result, Ok(job_result) if job_result.exit_code == Some(0));
|
||||||
let mut job_state = task_state.lock();
|
let mut job_state = task_state.lock();
|
||||||
job_state.pgid = None;
|
job_state.pgid = None;
|
||||||
job_state.status = match &result {
|
job_state.status = if success {
|
||||||
Ok(job_result) if job_result.exit_code == Some(0) => JobStatus::Completed,
|
JobStatus::Completed
|
||||||
_ => JobStatus::Failed,
|
} else {
|
||||||
|
JobStatus::Failed
|
||||||
};
|
};
|
||||||
drop(job_state);
|
drop(job_state);
|
||||||
|
task_notifications.push(job_notification(¬ify_id, ¬ify_tool, success));
|
||||||
result
|
result
|
||||||
})
|
})
|
||||||
};
|
};
|
||||||
@@ -1533,6 +1546,132 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[test]
|
||||||
|
fn job_completion_pushes_notification_for_own_context() {
|
||||||
|
run_async(async {
|
||||||
|
let mut ctx = plain_ctx();
|
||||||
|
ctx.declared_function_names.insert("echo".into());
|
||||||
|
|
||||||
|
let started = handle_start(&mut ctx, &json!({"tool": "echo", "arguments": {}}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let job_id = started["job_id"].as_str().unwrap().to_string();
|
||||||
|
|
||||||
|
let supervisor = ctx.supervisor.clone().unwrap();
|
||||||
|
let deadline = time::Instant::now() + Duration::from_secs(5);
|
||||||
|
while time::Instant::now() < deadline {
|
||||||
|
let finished = supervisor
|
||||||
|
.read()
|
||||||
|
.job(&job_id)
|
||||||
|
.is_none_or(|job| job.join_handle.is_finished());
|
||||||
|
if finished {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
time::sleep(Duration::from_millis(10)).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let events = ctx.notification_queue.drain();
|
||||||
|
assert_eq!(events.len(), 1);
|
||||||
|
assert_eq!(events[0].event, "job_completed");
|
||||||
|
assert_eq!(events[0].id, job_id);
|
||||||
|
assert_eq!(events[0].tool_or_agent, "echo");
|
||||||
|
assert_eq!(events[0].status, "success");
|
||||||
|
assert_eq!(
|
||||||
|
events[0].next_action,
|
||||||
|
format!("job__collect --id {job_id} for output")
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[test]
|
||||||
|
fn job_failure_pushes_failed_notification() {
|
||||||
|
run_async(async {
|
||||||
|
let mut ctx = plain_ctx();
|
||||||
|
ctx.declared_function_names.insert("false".into());
|
||||||
|
|
||||||
|
let started = handle_start(&mut ctx, &json!({"tool": "false", "arguments": {}}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let job_id = started["job_id"].as_str().unwrap().to_string();
|
||||||
|
|
||||||
|
let supervisor = ctx.supervisor.clone().unwrap();
|
||||||
|
let deadline = time::Instant::now() + Duration::from_secs(5);
|
||||||
|
while time::Instant::now() < deadline {
|
||||||
|
let finished = supervisor
|
||||||
|
.read()
|
||||||
|
.job(&job_id)
|
||||||
|
.is_none_or(|job| job.join_handle.is_finished());
|
||||||
|
if finished {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
time::sleep(Duration::from_millis(10)).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let events = ctx.notification_queue.drain();
|
||||||
|
assert_eq!(events.len(), 1);
|
||||||
|
assert_eq!(events[0].event, "job_failed");
|
||||||
|
assert_eq!(events[0].status, "failed");
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cancelled_job_notification_is_suppressed_at_drain() {
|
||||||
|
run_async(async {
|
||||||
|
let ctx = ctx_with_job_supervisor(4);
|
||||||
|
let join_handle = tokio::spawn(async {
|
||||||
|
time::sleep(Duration::from_secs(30)).await;
|
||||||
|
Ok(JobResult {
|
||||||
|
output: Value::Null,
|
||||||
|
exit_code: Some(0),
|
||||||
|
output_bytes_captured: 0,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
let handle = JobHandle {
|
||||||
|
id: "j1".to_string(),
|
||||||
|
tool: "execute_command".to_string(),
|
||||||
|
started_at: Instant::now(),
|
||||||
|
join_handle,
|
||||||
|
abort_signal: create_abort_signal(),
|
||||||
|
state: Arc::new(Mutex::new(JobState {
|
||||||
|
status: JobStatus::Running,
|
||||||
|
pgid: None,
|
||||||
|
})),
|
||||||
|
output_buf: Arc::new(Mutex::new(RingBuf::default())),
|
||||||
|
no_change_checks: 0,
|
||||||
|
};
|
||||||
|
ctx.supervisor
|
||||||
|
.as_ref()
|
||||||
|
.unwrap()
|
||||||
|
.write()
|
||||||
|
.register(handle)
|
||||||
|
.unwrap();
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("j1", "execute_command", false));
|
||||||
|
|
||||||
|
let result = handle_cancel(&ctx, &json!({"id": "j1"})).await.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(result["status"], "cancelled");
|
||||||
|
assert!(
|
||||||
|
super::super::drain_live_notifications(&ctx).is_empty(),
|
||||||
|
"events for a cancelled job must never reach the model"
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn already_collected_job_notification_is_suppressed_at_drain() {
|
||||||
|
let ctx = ctx_with_job_supervisor(4);
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("j1", "execute_command", true));
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
super::super::drain_live_notifications(&ctx).is_empty(),
|
||||||
|
"events for an already-collected job must be dropped"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn job_handlers_teach_cross_kind_for_agent_ids() {
|
fn job_handlers_teach_cross_kind_for_agent_ids() {
|
||||||
let ctx = ctx_with_job_supervisor(4);
|
let ctx = ctx_with_job_supervisor(4);
|
||||||
|
|||||||
+263
-21
@@ -341,12 +341,17 @@ pub async fn eval_tool_calls(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if ctx.current_depth == 0
|
if let Some(last) = output.last_mut() {
|
||||||
&& let Some(queue) = ctx.root_escalation_queue()
|
let escalations = if ctx.current_depth == 0 {
|
||||||
&& queue.has_pending()
|
ctx.root_escalation_queue()
|
||||||
&& let Some(last) = output.last_mut()
|
.filter(|queue| queue.has_pending())
|
||||||
{
|
.map(|queue| queue.pending_summary())
|
||||||
inject_escalation_notification(last, queue.pending_summary());
|
.unwrap_or_default()
|
||||||
|
} else {
|
||||||
|
vec![]
|
||||||
|
};
|
||||||
|
let notifications = drain_live_notifications(ctx);
|
||||||
|
merge_system_channel(last, escalations, notifications);
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(output)
|
Ok(output)
|
||||||
@@ -365,22 +370,75 @@ fn normalize_tool_result(result: Value) -> Value {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn inject_escalation_notification(last: &mut ToolResult, summary: Vec<Value>) {
|
/// Drains this context's own notification queue and drops events whose
|
||||||
let instruction = "Child agents are BLOCKED waiting for your reply. \
|
/// handle is no longer registered with the supervisor (already collected or
|
||||||
Call agent__reply_escalation for each pending escalation to unblock them.";
|
/// cancelled), so the model is never pointed at a dead id.
|
||||||
match &mut last.output {
|
fn drain_live_notifications(ctx: &RequestContext) -> Vec<Value> {
|
||||||
Value::Object(map) => {
|
let events = ctx.notification_queue.drain();
|
||||||
map.insert("pending_escalations".into(), json!(summary));
|
if events.is_empty() {
|
||||||
map.insert("escalation_instruction".into(), json!(instruction));
|
return vec![];
|
||||||
}
|
|
||||||
other => {
|
|
||||||
*other = json!({
|
|
||||||
"output": other.take(),
|
|
||||||
"pending_escalations": summary,
|
|
||||||
"escalation_instruction": instruction,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
let Some(supervisor) = ctx.supervisor.as_ref() else {
|
||||||
|
return vec![];
|
||||||
|
};
|
||||||
|
let sup = supervisor.read();
|
||||||
|
events
|
||||||
|
.into_iter()
|
||||||
|
.filter(|event| sup.has_job(&event.id) || sup.has_agent(&event.id))
|
||||||
|
.map(|event| event.to_value())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Single-pass merge of both system channels onto the last tool result of a
|
||||||
|
/// batch: pending escalations (children are blocked; listed first) and
|
||||||
|
/// background-task completion notifications. A single pass is mandatory —
|
||||||
|
/// two independent mergers would each apply the non-object wrap and nest the
|
||||||
|
/// output twice. With both channels empty this is a no-op, and with only
|
||||||
|
/// escalations it produces exactly the pre-notification output shape.
|
||||||
|
fn merge_system_channel(last: &mut ToolResult, escalations: Vec<Value>, notifications: Vec<Value>) {
|
||||||
|
if escalations.is_empty() && notifications.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let escalation_instruction = "Child agents are BLOCKED waiting for your reply. \
|
||||||
|
Call agent__reply_escalation for each pending escalation to unblock them.";
|
||||||
|
let notification_instruction =
|
||||||
|
"Background tasks have finished; collect each result with its next_action command.";
|
||||||
|
|
||||||
|
let map = match &mut last.output {
|
||||||
|
Value::Object(map) => map,
|
||||||
|
other => {
|
||||||
|
let mut map = serde_json::Map::new();
|
||||||
|
map.insert("output".into(), other.take());
|
||||||
|
*other = Value::Object(map);
|
||||||
|
match other {
|
||||||
|
Value::Object(map) => map,
|
||||||
|
_ => unreachable!(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if !escalations.is_empty() {
|
||||||
|
map.insert("pending_escalations".into(), json!(escalations));
|
||||||
|
map.insert(
|
||||||
|
"escalation_instruction".into(),
|
||||||
|
json!(escalation_instruction),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if !notifications.is_empty() {
|
||||||
|
map.insert("system_notifications".into(), json!(notifications));
|
||||||
|
map.insert(
|
||||||
|
"notification_instruction".into(),
|
||||||
|
json!(notification_instruction),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Escalation-only entry point retained so the characterization tests that
|
||||||
|
/// pinned the pre-merger output shape keep proving, unmodified, that
|
||||||
|
/// `merge_system_channel` with no notifications is byte-identical to the
|
||||||
|
/// injection behavior they were written against.
|
||||||
|
#[cfg(test)]
|
||||||
|
fn inject_escalation_notification(last: &mut ToolResult, summary: Vec<Value>) {
|
||||||
|
merge_system_channel(last, summary, vec![]);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||||
@@ -2495,6 +2553,7 @@ mod tests {
|
|||||||
};
|
};
|
||||||
use crate::config::{Agent, AgentConfig, AppConfig, AppState, WorkingMode};
|
use crate::config::{Agent, AgentConfig, AppConfig, AppState, WorkingMode};
|
||||||
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||||
|
use crate::supervisor::notification::job_notification;
|
||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
use base64::engine::general_purpose::STANDARD;
|
use base64::engine::general_purpose::STANDARD;
|
||||||
use rmcp::model::{CallToolResult, ContentBlock};
|
use rmcp::model::{CallToolResult, ContentBlock};
|
||||||
@@ -2579,6 +2638,189 @@ mod tests {
|
|||||||
assert!(result.output["escalation_instruction"].is_string());
|
assert!(result.output["escalation_instruction"].is_string());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn ctx_with_registered_job(id: &str) -> RequestContext {
|
||||||
|
let rt = tokio::runtime::Builder::new_current_thread()
|
||||||
|
.enable_all()
|
||||||
|
.build()
|
||||||
|
.unwrap();
|
||||||
|
let join_handle = rt.spawn(async {
|
||||||
|
Ok(crate::supervisor::JobResult {
|
||||||
|
output: Value::Null,
|
||||||
|
exit_code: Some(0),
|
||||||
|
output_bytes_captured: 0,
|
||||||
|
})
|
||||||
|
});
|
||||||
|
std::mem::forget(rt);
|
||||||
|
let handle = crate::supervisor::JobHandle {
|
||||||
|
id: id.to_string(),
|
||||||
|
tool: "execute_command".to_string(),
|
||||||
|
started_at: std::time::Instant::now(),
|
||||||
|
join_handle,
|
||||||
|
abort_signal: crate::utils::create_abort_signal(),
|
||||||
|
state: Arc::new(parking_lot::Mutex::new(crate::supervisor::JobState {
|
||||||
|
status: crate::supervisor::JobStatus::Completed,
|
||||||
|
pgid: None,
|
||||||
|
})),
|
||||||
|
output_buf: Arc::new(parking_lot::Mutex::new(jobs::RingBuf::default())),
|
||||||
|
no_change_checks: 0,
|
||||||
|
};
|
||||||
|
let mut sup = crate::supervisor::Supervisor::new(0, 3).with_max_concurrent_jobs(4);
|
||||||
|
sup.register(handle).unwrap();
|
||||||
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
|
ctx.supervisor = Some(Arc::new(parking_lot::RwLock::new(sup)));
|
||||||
|
ctx
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_system_channel_noop_when_both_channels_empty() {
|
||||||
|
let mut object_result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"}));
|
||||||
|
merge_system_channel(&mut object_result, vec![], vec![]);
|
||||||
|
assert_eq!(object_result.output, json!({"status": "ok"}));
|
||||||
|
|
||||||
|
let mut plain_result = ToolResult::new(call("t", Some("id-1")), json!("DONE"));
|
||||||
|
merge_system_channel(&mut plain_result, vec![], vec![]);
|
||||||
|
assert_eq!(plain_result.output, json!("DONE"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_system_channel_escalations_only_matches_legacy_wrap_bytes() {
|
||||||
|
let summary = vec![json!({"escalation_id": "esc_1"})];
|
||||||
|
let expected = json!({
|
||||||
|
"output": "DONE",
|
||||||
|
"pending_escalations": summary,
|
||||||
|
"escalation_instruction": "Child agents are BLOCKED waiting for your reply. \
|
||||||
|
Call agent__reply_escalation for each pending escalation to unblock them.",
|
||||||
|
});
|
||||||
|
|
||||||
|
let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE"));
|
||||||
|
merge_system_channel(&mut result, summary, vec![]);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
serde_json::to_string(&result.output).unwrap(),
|
||||||
|
serde_json::to_string(&expected).unwrap()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_system_channel_adds_notifications_without_escalation_keys() {
|
||||||
|
let mut result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"}));
|
||||||
|
merge_system_channel(&mut result, vec![], vec![json!({"id": "job_1"})]);
|
||||||
|
assert_eq!(result.output["status"], "ok");
|
||||||
|
assert_eq!(
|
||||||
|
result.output["system_notifications"],
|
||||||
|
json!([{"id": "job_1"}])
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
result.output["notification_instruction"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("next_action")
|
||||||
|
);
|
||||||
|
assert!(result.output.get("pending_escalations").is_none());
|
||||||
|
assert!(result.output.get("escalation_instruction").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_system_channel_wraps_non_object_once_with_both_channels() {
|
||||||
|
let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE"));
|
||||||
|
merge_system_channel(
|
||||||
|
&mut result,
|
||||||
|
vec![json!({"escalation_id": "esc_1"})],
|
||||||
|
vec![json!({"id": "job_1"})],
|
||||||
|
);
|
||||||
|
assert_eq!(result.output["output"], json!("DONE"));
|
||||||
|
assert_eq!(
|
||||||
|
result.output["pending_escalations"][0]["escalation_id"],
|
||||||
|
"esc_1"
|
||||||
|
);
|
||||||
|
assert_eq!(result.output["system_notifications"][0]["id"], "job_1");
|
||||||
|
let keys: Vec<&str> = result
|
||||||
|
.output
|
||||||
|
.as_object()
|
||||||
|
.unwrap()
|
||||||
|
.keys()
|
||||||
|
.map(|k| k.as_str())
|
||||||
|
.collect();
|
||||||
|
assert_eq!(
|
||||||
|
keys,
|
||||||
|
vec![
|
||||||
|
"output",
|
||||||
|
"pending_escalations",
|
||||||
|
"escalation_instruction",
|
||||||
|
"system_notifications",
|
||||||
|
"notification_instruction"
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drain_live_notifications_drops_unregistered_ids() {
|
||||||
|
let ctx = ctx_with_registered_job("job_live");
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("job_live", "execute_command", true));
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("job_gone", "execute_command", true));
|
||||||
|
|
||||||
|
let live = drain_live_notifications(&ctx);
|
||||||
|
|
||||||
|
assert_eq!(live.len(), 1);
|
||||||
|
assert_eq!(live[0]["id"], "job_live");
|
||||||
|
assert!(
|
||||||
|
ctx.notification_queue.drain().is_empty(),
|
||||||
|
"drain must consume the queue"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drain_live_notifications_without_supervisor_drops_everything() {
|
||||||
|
let ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("job_x", "execute_command", true));
|
||||||
|
assert!(drain_live_notifications(&ctx).is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn eval_tool_calls_merges_notifications_at_depth_without_escalations() {
|
||||||
|
let mut ctx = ctx_with_registered_job("job_n1");
|
||||||
|
ctx.current_depth = 1;
|
||||||
|
let queue = ctx.ensure_root_escalation_queue();
|
||||||
|
submit_escalation(&queue, "esc_1");
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("job_n1", "execute_command", true));
|
||||||
|
|
||||||
|
let calls = vec![call("unknown_tool", Some("id-1"))];
|
||||||
|
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
|
||||||
|
|
||||||
|
let out = &results[0].output;
|
||||||
|
assert_eq!(out["system_notifications"][0]["id"], "job_n1");
|
||||||
|
assert_eq!(out["system_notifications"][0]["event"], "job_completed");
|
||||||
|
assert!(out["notification_instruction"].is_string());
|
||||||
|
assert!(
|
||||||
|
out.get("pending_escalations").is_none(),
|
||||||
|
"escalations are root-only"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn eval_tool_calls_merges_both_channels_onto_last_result() {
|
||||||
|
let mut ctx = ctx_with_registered_job("job_n1");
|
||||||
|
let queue = ctx.ensure_root_escalation_queue();
|
||||||
|
submit_escalation(&queue, "esc_1");
|
||||||
|
ctx.notification_queue
|
||||||
|
.push(job_notification("job_n1", "execute_command", false));
|
||||||
|
|
||||||
|
let calls = vec![call("unknown_tool", Some("id-1"))];
|
||||||
|
let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap();
|
||||||
|
|
||||||
|
let out = &results[0].output;
|
||||||
|
assert_eq!(out["pending_escalations"][0]["escalation_id"], "esc_1");
|
||||||
|
assert_eq!(out["system_notifications"][0]["event"], "job_failed");
|
||||||
|
assert!(
|
||||||
|
out.get("output").is_none(),
|
||||||
|
"object outputs are extended in place, never wrapped"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn eval_tool_calls_soft_fails_unknown_tool() {
|
fn eval_tool_calls_soft_fails_unknown_tool() {
|
||||||
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
pub mod escalation;
|
pub mod escalation;
|
||||||
pub mod mailbox;
|
pub mod mailbox;
|
||||||
|
pub mod notification;
|
||||||
pub mod taskqueue;
|
pub mod taskqueue;
|
||||||
|
|
||||||
use crate::function::jobs::RingBuf;
|
use crate::function::jobs::RingBuf;
|
||||||
|
|||||||
@@ -0,0 +1,128 @@
|
|||||||
|
use fmt::{Debug, Formatter};
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::fmt;
|
||||||
|
|
||||||
|
/// One background-task completion event, delivered to the context that
|
||||||
|
/// started the task by merging a `system_notifications` entry onto the last
|
||||||
|
/// tool result of a batch.
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct SystemNotification {
|
||||||
|
pub event: &'static str,
|
||||||
|
pub id: String,
|
||||||
|
pub tool_or_agent: String,
|
||||||
|
pub status: &'static str,
|
||||||
|
pub next_action: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SystemNotification {
|
||||||
|
pub fn to_value(&self) -> Value {
|
||||||
|
json!({
|
||||||
|
"event": self.event,
|
||||||
|
"id": self.id,
|
||||||
|
"tool_or_agent": self.tool_or_agent,
|
||||||
|
"status": self.status,
|
||||||
|
"next_action": self.next_action,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn job_notification(id: &str, tool: &str, success: bool) -> SystemNotification {
|
||||||
|
SystemNotification {
|
||||||
|
event: if success {
|
||||||
|
"job_completed"
|
||||||
|
} else {
|
||||||
|
"job_failed"
|
||||||
|
},
|
||||||
|
id: id.to_string(),
|
||||||
|
tool_or_agent: tool.to_string(),
|
||||||
|
status: if success { "success" } else { "failed" },
|
||||||
|
next_action: format!("job__collect --id {id} for output"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Completion events for background work started by ONE context. Unlike the
|
||||||
|
/// escalation queue (shared, root-owned), every context owns a fresh queue:
|
||||||
|
/// a queue shared between parent and child would race their drains and
|
||||||
|
/// deliver one context's events into the other's transcript.
|
||||||
|
pub struct NotificationQueue {
|
||||||
|
pending: parking_lot::Mutex<Vec<SystemNotification>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl NotificationQueue {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
pending: parking_lot::Mutex::new(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn push(&self, notification: SystemNotification) {
|
||||||
|
self.pending.lock().push(notification);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn drain(&self) -> Vec<SystemNotification> {
|
||||||
|
std::mem::take(&mut *self.pending.lock())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for NotificationQueue {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Debug for NotificationQueue {
|
||||||
|
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
|
||||||
|
let count = self.pending.lock().len();
|
||||||
|
f.debug_struct("NotificationQueue")
|
||||||
|
.field("pending_count", &count)
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn job_notification_success_shape() {
|
||||||
|
let event = job_notification("job_a1b2", "execute_command", true);
|
||||||
|
assert_eq!(
|
||||||
|
event.to_value(),
|
||||||
|
json!({
|
||||||
|
"event": "job_completed",
|
||||||
|
"id": "job_a1b2",
|
||||||
|
"tool_or_agent": "execute_command",
|
||||||
|
"status": "success",
|
||||||
|
"next_action": "job__collect --id job_a1b2 for output",
|
||||||
|
})
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn job_notification_failure_shape() {
|
||||||
|
let event = job_notification("job_a1b2", "execute_command", false);
|
||||||
|
assert_eq!(event.event, "job_failed");
|
||||||
|
assert_eq!(event.status, "failed");
|
||||||
|
assert_eq!(event.next_action, "job__collect --id job_a1b2 for output");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drain_empties_queue_and_preserves_order() {
|
||||||
|
let queue = NotificationQueue::new();
|
||||||
|
queue.push(job_notification("job_1", "execute_command", true));
|
||||||
|
queue.push(job_notification("job_2", "execute_command", false));
|
||||||
|
|
||||||
|
let drained = queue.drain();
|
||||||
|
|
||||||
|
assert_eq!(drained.len(), 2);
|
||||||
|
assert_eq!(drained[0].id, "job_1");
|
||||||
|
assert_eq!(drained[1].id, "job_2");
|
||||||
|
assert!(queue.drain().is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drain_on_empty_queue_is_a_noop() {
|
||||||
|
let queue = NotificationQueue::default();
|
||||||
|
assert!(queue.drain().is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user