feat: Added explicit guardrail handling for pending agents

This commit is contained in:
2026-06-11 20:20:14 -06:00
parent 6ebd32d47c
commit 00939e4634
7 changed files with 270 additions and 21 deletions
+30
View File
@@ -93,6 +93,36 @@ pub(in crate::config) const DEFAULT_SPAWN_INSTRUCTIONS: &str = indoc! {"
agent__collect --id agent_explore_e5f6g7h8 agent__collect --id agent_explore_e5f6g7h8
``` ```
### CRITICAL: Never end your turn with pending agents
Spawned agents do NOT report back on their own. They run in the background until you
actively reclaim them with `agent__collect` (to get their output) or `agent__cancel`
(to discard them). If you spawn agents and then emit a final message without reclaiming
them, the system will detect the unreclaimed agents and reject the turn-end, injecting
a reminder forcing you to handle them. After several such reminders, the system will
auto-cancel them and warn you that work was lost.
The correct flow when you have nothing else to do:
```
# WRONG - do NOT do this:
agent__spawn --agent explore --prompt \"...\"
agent__spawn --agent explore --prompt \"...\"
# ... emit text like \"I will synthesize once they report back.\" and stop
# ^ The agents will be abandoned. Their output will be lost.
# RIGHT - always do this:
agent__spawn --agent explore --prompt \"...\"
agent__spawn --agent explore --prompt \"...\"
agent__collect --id <first_id> # blocks until done
agent__collect --id <second_id> # blocks until done
# ... NOW you can synthesize and end your turn
```
`agent__collect` is a **blocking wait**: it pauses your execution until the agent
completes, then returns the output as a tool result. Use it freely — it is the
correct primitive for \"I'm done with my own work and just need the agents' results\".
### Parallel Spawning (DEFAULT for multi-agent work) ### Parallel Spawning (DEFAULT for multi-agent work)
When a task needs multiple agents, **spawn them all at once**, then collect: When a task needs multiple agents, **spawn them all at once**, then collect:
+7 -1
View File
@@ -120,6 +120,7 @@ pub struct RequestContext {
pub escalation_queue: Option<Arc<EscalationQueue>>, pub escalation_queue: Option<Arc<EscalationQueue>>,
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 todo_list: TodoList, pub todo_list: TodoList,
pub skill_registry: SkillRegistry, pub skill_registry: SkillRegistry,
pub last_continuation_response: Option<String>, pub last_continuation_response: Option<String>,
@@ -149,6 +150,7 @@ impl RequestContext {
escalation_queue: None, escalation_queue: None,
current_depth: 0, current_depth: 0,
auto_continue_count: 0, auto_continue_count: 0,
pending_agents_guardrail_count: 0,
todo_list: TodoList::default(), todo_list: TodoList::default(),
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
@@ -204,6 +206,7 @@ impl RequestContext {
escalation_queue: None, escalation_queue: None,
current_depth: 0, current_depth: 0,
auto_continue_count: 0, auto_continue_count: 0,
pending_agents_guardrail_count: 0,
todo_list: TodoList::default(), todo_list: TodoList::default(),
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
@@ -246,6 +249,7 @@ impl RequestContext {
escalation_queue: self.escalation_queue.clone(), escalation_queue: self.escalation_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,
todo_list: self.todo_list.clone(), todo_list: self.todo_list.clone(),
skill_registry: self.skill_registry.clone(), skill_registry: self.skill_registry.clone(),
last_continuation_response: None, last_continuation_response: None,
@@ -286,6 +290,7 @@ impl RequestContext {
escalation_queue: parent.escalation_queue.clone(), escalation_queue: parent.escalation_queue.clone(),
current_depth, current_depth,
auto_continue_count: 0, auto_continue_count: 0,
pending_agents_guardrail_count: 0,
todo_list: TodoList::default(), todo_list: TodoList::default(),
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
@@ -2787,7 +2792,7 @@ impl RequestContext {
if self.agent.take().is_some() { if self.agent.take().is_some() {
if let Some(supervisor) = self.supervisor.clone() { if let Some(supervisor) = self.supervisor.clone() {
supervisor.read().cancel_all(); supervisor.read().cancel_recursive();
} }
self.supervisor = None; self.supervisor = None;
self.parent_supervisor = None; self.parent_supervisor = None;
@@ -2796,6 +2801,7 @@ impl RequestContext {
self.escalation_queue = None; self.escalation_queue = None;
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.todo_list = TodoList::default(); self.todo_list = TodoList::default();
self.rag.take(); self.rag.take();
self.discontinuous_last_message(); self.discontinuous_last_message();
+152 -18
View File
@@ -3,7 +3,7 @@ use crate::client::{Model, ModelType, call_chat_completions};
use crate::config::{Agent, AppState, Input, RequestContext, Role, RoleLike}; use crate::config::{Agent, AppState, Input, RequestContext, Role, RoleLike};
use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox}; use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor}; use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor};
use crate::utils::{AbortSignal, create_abort_signal}; use crate::utils::{AbortSignal, create_abort_signal, wait_abort_signal};
use crate::graph; use crate::graph;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -16,10 +16,69 @@ use std::pin::Pin;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use tokio::time; use tokio::time;
use tokio::time::Instant;
use uuid::Uuid; use uuid::Uuid;
pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__"; pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__";
pub const PENDING_AGENTS_GUARDRAIL_MAX: u32 = 3;
pub enum GuardrailAction {
NoAction,
Inject(String),
ForceTerminate(Vec<String>),
}
pub fn pending_agent_ids(ctx: &RequestContext) -> Vec<String> {
let Some(sup) = ctx.supervisor.as_ref() else {
return Vec::new();
};
let sup = sup.read();
sup.list_agents()
.into_iter()
.filter_map(|(id, _)| match sup.is_finished(id) {
Some(false) => Some(id.to_string()),
_ => None,
})
.collect()
}
pub fn build_pending_agents_guardrail_prompt(ids: &[String]) -> String {
let count = ids.len();
let id_list = ids
.iter()
.map(|id| format!("- {id}"))
.collect::<Vec<_>>()
.join("\n");
format!(
"[SYSTEM GUARDRAIL] You attempted to end your turn while {count} spawned background agent(s) \
are still running:\n{id_list}\n\nThese agents will be abandoned if your turn ends now. You MUST \
reclaim each one before ending your turn. For each agent: call `agent__collect` (blocks until \
done, returns output) or `agent__cancel` (discards). Do NOT emit a text-only response \
expecting them to 'report back' — they will not."
)
}
pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailAction {
let pending = pending_agent_ids(ctx);
if pending.is_empty() {
ctx.pending_agents_guardrail_count = 0;
return GuardrailAction::NoAction;
}
if ctx.pending_agents_guardrail_count >= PENDING_AGENTS_GUARDRAIL_MAX {
if let Some(sup) = ctx.supervisor.as_ref().cloned() {
sup.read().cancel_recursive();
}
ctx.pending_agents_guardrail_count = 0;
return GuardrailAction::ForceTerminate(pending);
}
ctx.pending_agents_guardrail_count += 1;
GuardrailAction::Inject(build_pending_agents_guardrail_prompt(&pending))
}
pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> { pub fn escalation_function_declarations() -> Vec<FunctionDeclaration> {
vec![FunctionDeclaration { vec![FunctionDeclaration {
name: format!("{SUPERVISOR_FUNCTION_PREFIX}reply_escalation"), name: format!("{SUPERVISOR_FUNCTION_PREFIX}reply_escalation"),
@@ -55,7 +114,11 @@ pub fn supervisor_function_declarations() -> Vec<FunctionDeclaration> {
vec![ vec![
FunctionDeclaration { FunctionDeclaration {
name: format!("{SUPERVISOR_FUNCTION_PREFIX}spawn"), name: format!("{SUPERVISOR_FUNCTION_PREFIX}spawn"),
description: "Spawn a subagent to run in the background. Returns a task_id for tracking. The agent runs in parallel. You can continue working while it executes.".to_string(), description: "Spawn a subagent to run in the background. Returns an `id` immediately so you can continue \
working in parallel. CRITICAL: every spawned agent MUST be reclaimed before you end your \
turn — call `agent__collect` to retrieve its output, or `agent__cancel` if you no longer \
need it. Ending your turn with pending agents will abandon their work and the system will \
reject the turn-end.".to_string(),
parameters: JsonSchema { parameters: JsonSchema {
type_value: Some("object".to_string()), type_value: Some("object".to_string()),
properties: Some(IndexMap::from([ properties: Some(IndexMap::from([
@@ -109,7 +172,11 @@ pub fn supervisor_function_declarations() -> Vec<FunctionDeclaration> {
}, },
FunctionDeclaration { FunctionDeclaration {
name: format!("{SUPERVISOR_FUNCTION_PREFIX}collect"), name: format!("{SUPERVISOR_FUNCTION_PREFIX}collect"),
description: "Wait for a spawned agent to finish and return its result. Blocks until the agent completes.".to_string(), description: "Block until the named spawned agent finishes and return its result. This is your primary \
wait primitive — it pauses your execution until the agent completes (or you are interrupted). \
Call this for every agent you spawned before ending your turn. Do NOT end your turn assuming \
agents will 'report back later' — they will not; they will be abandoned. If you no longer \
need an agent's result, call `agent__cancel` instead.".to_string(),
parameters: JsonSchema { parameters: JsonSchema {
type_value: Some("object".to_string()), type_value: Some("object".to_string()),
properties: Some(IndexMap::from([( properties: Some(IndexMap::from([(
@@ -137,7 +204,10 @@ pub fn supervisor_function_declarations() -> Vec<FunctionDeclaration> {
}, },
FunctionDeclaration { FunctionDeclaration {
name: format!("{SUPERVISOR_FUNCTION_PREFIX}cancel"), name: format!("{SUPERVISOR_FUNCTION_PREFIX}cancel"),
description: "Cancel a running subagent by its ID.".to_string(), description: "Cancel a running subagent by its ID. Use this when an agent's output is no longer needed \
(e.g. you changed direction, or you're about to end your turn and don't want to wait). \
Cancellation cascades: all of the cancelled agent's own descendants are also cancelled. This \
call waits briefly for the agent to actually finish cleanup before returning.".to_string(),
parameters: JsonSchema { parameters: JsonSchema {
type_value: Some("object".to_string()), type_value: Some("object".to_string()),
properties: Some(IndexMap::from([( properties: Some(IndexMap::from([(
@@ -315,7 +385,7 @@ pub async fn handle_supervisor_tool(
"check" => handle_check(ctx, args).await, "check" => handle_check(ctx, args).await,
"collect" => handle_collect(ctx, args).await, "collect" => handle_collect(ctx, args).await,
"list" => handle_list(ctx), "list" => handle_list(ctx),
"cancel" => handle_cancel(ctx, args), "cancel" => handle_cancel(ctx, args).await,
"send_message" => handle_send_message(ctx, args), "send_message" => handle_send_message(ctx, args),
"check_inbox" => handle_check_inbox(ctx), "check_inbox" => handle_check_inbox(ctx),
"task_create" => handle_task_create(ctx, args), "task_create" => handle_task_create(ctx, args),
@@ -370,14 +440,28 @@ pub fn run_child_agent(
} }
if tool_results.is_empty() { if tool_results.is_empty() {
break; match check_pending_agents_guardrail(&mut child_ctx) {
GuardrailAction::NoAction => break,
GuardrailAction::ForceTerminate(ids) => {
log::warn!(
"Pending-agent guardrail force-cancelled {} agent(s) after max reminders: {:?}",
ids.len(),
ids
);
break;
}
GuardrailAction::Inject(prompt) => {
input = Input::from_str(&child_ctx, &prompt, None)?;
continue;
}
}
} }
input = input.merge_tool_results(output, tool_results); input = input.merge_tool_results(output, tool_results);
} }
if let Some(supervisor) = child_ctx.supervisor.clone() { if let Some(supervisor) = child_ctx.supervisor.clone() {
supervisor.read().cancel_all(); supervisor.read().cancel_recursive();
} }
Ok(accumulated_output) Ok(accumulated_output)
@@ -642,6 +726,7 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let spawn_agent_id = agent_id.clone(); let spawn_agent_id = agent_id.clone();
let spawn_agent_name = agent_name.clone(); let spawn_agent_name = agent_name.clone();
let spawn_abort = child_abort.clone(); let spawn_abort = child_abort.clone();
let child_supervisor = child_ctx.supervisor.clone();
let join_handle = tokio::spawn(async move { let join_handle = tokio::spawn(async move {
let result = run_child_agent(child_ctx, input, spawn_abort).await; let result = run_child_agent(child_ctx, input, spawn_abort).await;
@@ -669,6 +754,7 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
inbox: child_inbox, inbox: child_inbox,
abort_signal: child_abort, abort_signal: child_abort,
join_handle, join_handle,
child_supervisor,
}; };
let supervisor = ctx let supervisor = ctx
@@ -683,7 +769,11 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
"status": "ok", "status": "ok",
"id": agent_id, "id": agent_id,
"agent": agent_name, "agent": agent_name,
"message": format!("Agent '{agent_name}' spawned as '{agent_id}'. Use agent__check or agent__collect to get results."), "message": format!("Agent '{agent_name}' spawned as '{agent_id}' and is running in the background. CRITICAL: \
you MUST reclaim this agent before ending your turn — call `agent__collect` (blocks until \
done, returns output) or `agent__cancel` (if you no longer need it). Ending your turn with \
unreclaimed agents will be rejected and forces you to handle them. Do NOT assume the agent \
will 'report back' on its own."),
})) }))
} }
@@ -743,7 +833,7 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
.cloned() .cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?; .ok_or_else(|| anyhow!("No supervisor active"))?;
{ let target_abort = {
let sup = supervisor.read(); let sup = supervisor.read();
if sup.is_finished(id).is_none() { if sup.is_finished(id).is_none() {
return Ok(json!({ return Ok(json!({
@@ -751,7 +841,8 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
"message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.") "message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.")
})); }));
} }
} sup.abort_signal_for(id)
};
loop { loop {
let is_finished = { let is_finished = {
@@ -775,7 +866,27 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
})); }));
} }
time::sleep(Duration::from_millis(200)).await; match target_abort.as_ref() {
Some(abort) if abort.aborted() => {
let deadline = Instant::now() + Duration::from_secs(2);
while Instant::now() < deadline {
if supervisor.read().is_finished(id).unwrap_or(false) {
break;
}
time::sleep(Duration::from_millis(50)).await;
}
break;
}
Some(abort) => {
tokio::select! {
_ = time::sleep(Duration::from_millis(200)) => {}
_ = wait_abort_signal(abort) => {}
}
}
None => {
time::sleep(Duration::from_millis(200)).await;
}
}
} }
let handle = { let handle = {
@@ -792,6 +903,7 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
.map_err(|e| anyhow!("Agent failed: {e}"))?; .map_err(|e| anyhow!("Agent failed: {e}"))?;
let output = summarize_output(ctx, &result.agent_name, &result.output).await?; let output = summarize_output(ctx, &result.agent_name, &result.output).await?;
ctx.pending_agents_guardrail_count = 0;
Ok(json!({ Ok(json!({
"status": "completed", "status": "completed",
@@ -836,7 +948,7 @@ fn handle_list(ctx: &mut RequestContext) -> Result<Value> {
})) }))
} }
fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result<Value> { async fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
let id = args let id = args
.get("id") .get("id")
.and_then(Value::as_str) .and_then(Value::as_str)
@@ -847,14 +959,34 @@ fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
.as_ref() .as_ref()
.cloned() .cloned()
.ok_or_else(|| anyhow!("No supervisor active"))?; .ok_or_else(|| anyhow!("No supervisor active"))?;
let mut sup = supervisor.write();
match sup.take(id) { let handle = {
let mut sup = supervisor.write();
sup.take(id)
};
match handle {
Some(handle) => { Some(handle) => {
let agent_name = handle.agent_name.clone();
if let Some(child_sup) = handle.child_supervisor.as_ref() {
child_sup.read().cancel_recursive();
}
handle.abort_signal.set_ctrlc(); handle.abort_signal.set_ctrlc();
let cleanup = tokio::time::timeout(Duration::from_secs(5), handle.join_handle).await;
ctx.pending_agents_guardrail_count = 0;
let message = match cleanup {
Ok(_) => format!("Cancelled agent '{agent_name}' and waited for cleanup."),
Err(_) => format!(
"Cancelled agent '{agent_name}'; cleanup did not complete within 5s. Its descendants have been signalled and will tear down asynchronously."
),
};
Ok(json!({ Ok(json!({
"status": "ok", "status": "ok",
"message": format!("Cancelled agent '{}'", handle.agent_name), "message": message,
})) }))
} }
None => Ok(json!({ None => Ok(json!({
@@ -1283,6 +1415,7 @@ mod tests {
inbox: Arc::new(Inbox::new()), inbox: Arc::new(Inbox::new()),
abort_signal: create_abort_signal(), abort_signal: create_abort_signal(),
join_handle, join_handle,
child_supervisor: None,
}; };
ctx.supervisor ctx.supervisor
.as_ref() .as_ref()
@@ -1362,6 +1495,7 @@ mod tests {
inbox, inbox,
abort_signal: abort, abort_signal: abort,
join_handle, join_handle,
child_supervisor: None,
}; };
ctx.supervisor ctx.supervisor
.as_ref() .as_ref()
@@ -1381,7 +1515,7 @@ mod tests {
fn handle_cancel_registered_agent() { fn handle_cancel_registered_agent() {
let mut ctx = ctx_with_supervisor(4, 3); let mut ctx = ctx_with_supervisor(4, 3);
register_fake_agent(&mut ctx, "a1", "explore"); register_fake_agent(&mut ctx, "a1", "explore");
let result = handle_cancel(&mut ctx, &json!({"id": "a1"})).unwrap(); let result = run_async(handle_cancel(&mut ctx, &json!({"id": "a1"}))).unwrap();
assert_eq!(result["status"], "ok"); assert_eq!(result["status"], "ok");
assert_eq!(ctx.supervisor.as_ref().unwrap().read().active_count(), 0); assert_eq!(ctx.supervisor.as_ref().unwrap().read().active_count(), 0);
} }
@@ -1389,14 +1523,14 @@ mod tests {
#[test] #[test]
fn handle_cancel_unknown_agent() { fn handle_cancel_unknown_agent() {
let mut ctx = ctx_with_supervisor(4, 3); let mut ctx = ctx_with_supervisor(4, 3);
let result = handle_cancel(&mut ctx, &json!({"id": "missing"})).unwrap(); let result = run_async(handle_cancel(&mut ctx, &json!({"id": "missing"}))).unwrap();
assert_eq!(result["status"], "error"); assert_eq!(result["status"], "error");
} }
#[test] #[test]
fn handle_cancel_no_supervisor_errors() { fn handle_cancel_no_supervisor_errors() {
let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
let result = handle_cancel(&mut ctx, &json!({"id": "x"})); let result = run_async(handle_cancel(&mut ctx, &json!({"id": "x"})));
assert!(result.is_err()); assert!(result.is_err());
} }
+24 -1
View File
@@ -7,8 +7,10 @@ use crate::config::{
Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions, Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions,
}; };
use crate::function::skill::skill_function_declarations; use crate::function::skill::skill_function_declarations;
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
use crate::utils::create_abort_signal; use crate::utils::create_abort_signal;
use anyhow::{Context, Error, Result, anyhow, bail}; use anyhow::{Context, Error, Result, anyhow, bail};
use log::warn;
use serde_json::Value; use serde_json::Value;
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::Arc; use std::sync::Arc;
@@ -266,7 +268,28 @@ async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -
} }
if tool_results.is_empty() { if tool_results.is_empty() {
return Ok(accumulated); match check_pending_agents_guardrail(ctx) {
GuardrailAction::NoAction => return Ok(accumulated),
GuardrailAction::ForceTerminate(ids) => {
warn!(
"Pending-agent guardrail force-cancelled {} agent(s) after max reminders: {:?}",
ids.len(),
ids
);
return Ok(accumulated);
}
GuardrailAction::Inject(prompt) => {
if turn + 1 == node.max_iterations {
bail!(
"llm node hit max_iterations ({}) before LLM concluded",
node.max_iterations
);
}
let role = ctx.role.clone();
input = Input::from_str(ctx, &prompt, role)?;
continue;
}
}
} }
if turn + 1 == node.max_iterations { if turn + 1 == node.max_iterations {
+17 -1
View File
@@ -26,6 +26,7 @@ use crate::config::{
RequestContext, SHELL_ROLE, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists, RequestContext, SHELL_ROLE, TEMP_SESSION_NAME, WorkingMode, ensure_parent_exists,
install_builtins, list_agents, load_env_file, macro_execute, sync_models, install_builtins, list_agents, load_env_file, macro_execute, sync_models,
}; };
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
use crate::render::{prompt_theme, render_error}; use crate::render::{prompt_theme, render_error};
use crate::repl::Repl; use crate::repl::Repl;
use crate::utils::*; use crate::utils::*;
@@ -35,7 +36,7 @@ use clap::{CommandFactory, Parser};
use clap_complete::CompleteEnv; use clap_complete::CompleteEnv;
use client::ClientConfig; use client::ClientConfig;
use inquire::{Select, Text, set_global_render_config}; use inquire::{Select, Text, set_global_render_config};
use log::LevelFilter; use log::{LevelFilter, warn};
use log4rs::append::console::ConsoleAppender; use log4rs::append::console::ConsoleAppender;
use log4rs::append::file::FileAppender; use log4rs::append::file::FileAppender;
use log4rs::config::{Appender, Logger, Root}; use log4rs::config::{Appender, Logger, Root};
@@ -419,6 +420,21 @@ async fn start_directive(
abort_signal, abort_signal,
) )
.await?; .await?;
} else {
match check_pending_agents_guardrail(ctx) {
GuardrailAction::Inject(prompt) => {
let guardrail_input = Input::from_str(ctx, &prompt, None)?;
return start_directive(ctx, guardrail_input, code_mode, abort_signal).await;
}
GuardrailAction::ForceTerminate(ids) => {
warn!(
"Pending-agent guardrail force-cancelled {} agent(s) after max reminders: {:?}",
ids.len(),
ids
);
}
GuardrailAction::NoAction => {}
}
} }
ctx.exit_session()?; ctx.exit_session()?;
+24
View File
@@ -12,6 +12,7 @@ use crate::config::{
macro_execute, macro_execute,
}; };
use crate::config::{AssetCategory, paths}; use crate::config::{AssetCategory, paths};
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
use crate::render::render_error; use crate::render::render_error;
use crate::utils::{ use crate::utils::{
AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file, AbortSignal, abortable_run_with_spinner, create_abort_signal, dimmed_text, set_text, temp_file,
@@ -306,6 +307,9 @@ Type ".help" for additional help.
} }
Ok(Signal::CtrlC) => { Ok(Signal::CtrlC) => {
self.abort_signal.set_ctrlc(); self.abort_signal.set_ctrlc();
if let Some(supervisor) = self.ctx.read().supervisor.clone() {
supervisor.read().cancel_recursive();
}
println!("(To exit, press Ctrl+D or enter \".exit\")\n"); println!("(To exit, press Ctrl+D or enter \".exit\")\n");
} }
Ok(Signal::CtrlD) => { Ok(Signal::CtrlD) => {
@@ -315,6 +319,11 @@ Type ".help" for additional help.
_ => {} _ => {}
} }
} }
if let Some(supervisor) = self.ctx.read().supervisor.clone() {
supervisor.read().cancel_recursive();
}
self.ctx.write().exit_session()?; self.ctx.write().exit_session()?;
Ok(()) Ok(())
} }
@@ -435,6 +444,7 @@ pub async fn run_repl_command(
abort_signal: AbortSignal, abort_signal: AbortSignal,
mut line: &str, mut line: &str,
) -> Result<bool> { ) -> Result<bool> {
ctx.pending_agents_guardrail_count = 0;
if let Ok(Some(captures)) = MULTILINE_RE.captures(line) if let Ok(Some(captures)) = MULTILINE_RE.captures(line)
&& let Some(text_match) = captures.get(1) && let Some(text_match) = captures.get(1)
{ {
@@ -1011,6 +1021,20 @@ async fn ask(
) )
.await .await
} else { } else {
match check_pending_agents_guardrail(ctx) {
GuardrailAction::Inject(prompt) => {
let guardrail_input = Input::from_str(ctx, &prompt, None)?;
return ask(ctx, abort_signal, guardrail_input, false).await;
}
GuardrailAction::ForceTerminate(ids) => {
warn!(
"Pending-agent guardrail force-cancelled {} agent(s) after max reminders: {:?}",
ids.len(),
ids
);
}
GuardrailAction::NoAction => {}
}
let do_continue = should_continue(ctx); let do_continue = should_continue(ctx);
if do_continue { if do_continue {
+16
View File
@@ -5,6 +5,7 @@ pub mod taskqueue;
use crate::utils::AbortSignal; use crate::utils::AbortSignal;
use fmt::{Debug, Formatter}; use fmt::{Debug, Formatter};
use mailbox::Inbox; use mailbox::Inbox;
use parking_lot::RwLock;
use taskqueue::TaskQueue; use taskqueue::TaskQueue;
use anyhow::{Result, bail}; use anyhow::{Result, bail};
@@ -33,6 +34,7 @@ pub struct AgentHandle {
pub inbox: Arc<Inbox>, pub inbox: Arc<Inbox>,
pub abort_signal: AbortSignal, pub abort_signal: AbortSignal,
pub join_handle: JoinHandle<Result<AgentResult>>, pub join_handle: JoinHandle<Result<AgentResult>>,
pub child_supervisor: Option<Arc<RwLock<Supervisor>>>,
} }
pub struct Supervisor { pub struct Supervisor {
@@ -103,6 +105,10 @@ impl Supervisor {
self.handles.get(id).map(|h| &h.inbox) self.handles.get(id).map(|h| &h.inbox)
} }
pub fn abort_signal_for(&self, id: &str) -> Option<AbortSignal> {
self.handles.get(id).map(|h| h.abort_signal.clone())
}
pub fn list_agents(&self) -> Vec<(&str, &str)> { pub fn list_agents(&self) -> Vec<(&str, &str)> {
self.handles self.handles
.values() .values()
@@ -115,6 +121,15 @@ impl Supervisor {
handle.abort_signal.set_ctrlc(); handle.abort_signal.set_ctrlc();
} }
} }
pub fn cancel_recursive(&self) {
for handle in self.handles.values() {
handle.abort_signal.set_ctrlc();
if let Some(child_sup) = handle.child_supervisor.as_ref() {
child_sup.read().cancel_recursive();
}
}
}
} }
impl Debug for Supervisor { impl Debug for Supervisor {
@@ -152,6 +167,7 @@ mod tests {
inbox: Arc::new(Inbox::new()), inbox: Arc::new(Inbox::new()),
abort_signal: create_abort_signal(), abort_signal: create_abort_signal(),
join_handle, join_handle,
child_supervisor: None,
} }
} }