From bfc3b7bfeaaecbeaff7647fef09bd0d4174ceeb5 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Tue, 25 Aug 2026 13:55:15 -0600 Subject: [PATCH] test: pin current tool-eval, guardrail, and truncation behavior ahead of background-jobs work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit T0 characterization safety net per plans/background-jobs-design.md §9.3: 43 tests pinning handle_collect/check/cancel/spawn, the pending-agents guardrail (incl. ForceTerminate + counter resets), eval_tool_calls partition/re-sort/soft-fail/ loop-alert/truncation, truncate_if_needed's UTF-8 boundary edge, ToolCall::eval prefix routing, merge_tool_results shape, and cancel_recursive recursion. Known-buggy behaviors deliberately pinned for visible later diffs: handle_check consumes finished handles, guardrail ignores finished-but-uncollected agents, mid-char truncation returns the full string with marker prepended. Not covered (findings): empty-after-dedup bail is unreachable from non-empty input; run_child_agent needs a mock LLM client (none exists) — manual case; over-threshold summarization pinned via deterministic unknown-model failure. --- src/config/input.rs | 68 +++++++ src/function/mod.rs | 313 +++++++++++++++++++++++++++- src/function/supervisor.rs | 408 ++++++++++++++++++++++++++++++++++++- src/supervisor/mod.rs | 21 ++ 4 files changed, 808 insertions(+), 2 deletions(-) diff --git a/src/config/input.rs b/src/config/input.rs index 71f82d5..5d35064 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -593,6 +593,8 @@ mod tests { use super::*; use crate::config::request_context::RequestContext; use crate::config::{AppState, WorkingMode}; + use crate::function::ToolCall; + use serde_json::json; use std::fs; use std::sync::Arc; use std::time::SystemTime; @@ -973,4 +975,70 @@ mod tests { )); assert!(result.is_err()); } + + fn tool_result(id: &str, output: &str) -> ToolResult { + ToolResult::new( + ToolCall::new("t".into(), json!({}), Some(id.to_string())), + json!(output), + ) + } + + #[test] + fn merge_tool_results_first_merge_creates_container() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None).unwrap(); + + let input = + input.merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]); + + let tool_calls = input.tool_calls().as_ref().unwrap(); + assert_eq!(tool_calls.text, "assistant text"); + assert!(!tool_calls.sequence); + assert_eq!(tool_calls.tool_results.len(), 1); + assert!(tool_calls.tool_results[0].text.is_none()); + } + + #[test] + fn merge_tool_results_second_merge_marks_sequence_and_tags_text() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None) + .unwrap() + .merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]); + + let input = + input.merge_tool_results("second text".into(), vec![tool_result("id-2", "ok2")]); + + let tool_calls = input.tool_calls().as_ref().unwrap(); + assert!(tool_calls.sequence); + assert_eq!(tool_calls.tool_results.len(), 2); + assert_eq!(tool_calls.text, "assistant text"); + assert!(tool_calls.tool_results[0].text.is_none()); + assert_eq!( + tool_calls.tool_results[1].text, + Some("second text".to_string()) + ); + } + + #[test] + fn build_messages_wraps_tool_results_in_single_assistant_message() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None) + .unwrap() + .merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]) + .merge_tool_results("second text".into(), vec![tool_result("id-2", "ok2")]); + + let messages = input.build_messages().unwrap(); + + let tool_call_messages: Vec<_> = messages + .iter() + .filter(|m| matches!(m.content, MessageContent::ToolCalls(_))) + .collect(); + assert_eq!(tool_call_messages.len(), 1); + let message = tool_call_messages[0]; + assert!(matches!(message.role, MessageRole::Assistant)); + let MessageContent::ToolCalls(tool_calls) = &message.content else { + unreachable!(); + }; + assert_eq!(tool_calls.tool_results.len(), 2); + } } diff --git a/src/function/mod.rs b/src/function/mod.rs index d32a24c..f552854 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -2474,7 +2474,7 @@ mod tests { FIXTURE_ANNOTATED_TEXT, FIXTURE_ANNOTATED_URI, FIXTURE_BLOB_BYTES, FIXTURE_BLOB_URI, FIXTURE_LOG_TEXT, FIXTURE_LOG_URI, FixtureServer, fixture_runtime, }; - use crate::config::{AppState, WorkingMode}; + use crate::config::{Agent, AgentConfig, AppConfig, AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use base64::Engine; use base64::engine::general_purpose::STANDARD; @@ -4004,4 +4004,315 @@ mod tests { assert!(dir.is_dir()); fs::remove_dir_all(&dir).unwrap(); } + + #[test] + fn eval_tool_calls_partitions_mcp_and_sequential_then_resorts() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![ + call("unknown_first", Some("id-1")), + ToolCall::new( + "mcp_search_foo".into(), + json!({"query": "q"}), + Some("id-2".into()), + ), + call("unknown_last", Some("id-3")), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 3); + assert_eq!(results[0].call.name, "unknown_first"); + assert_eq!(results[1].call.name, "mcp_search_foo"); + assert_eq!(results[2].call.name, "unknown_last"); + + for sequential in [&results[0], &results[2]] { + let err = sequential.output["tool_call_error"].as_str().unwrap(); + assert!( + err.contains("use only tools listed in your catalog"), + "{err}" + ); + } + let mcp_err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(mcp_err.starts_with("MCP search failed"), "{mcp_err}"); + assert!(!mcp_err.contains("use only tools listed in your catalog")); + } + + #[test] + fn eval_tool_calls_isolates_failures_within_a_batch() { + let app = AppState { + config: Arc::new(AppConfig { + auto_continue: true, + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + ctx.tool_scope.functions.append_todo_functions(); + let calls = vec![ + ToolCall::new( + "todo__init".into(), + json!({"goal": "ship it"}), + Some("id-1".into()), + ), + call("unknown_tool", Some("id-2")), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].call.name, "todo__init"); + assert_eq!(results[0].output["status"], "ok"); + assert!(results[0].output.get("tool_call_error").is_none()); + let err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_tool_calls_reports_loop_alert_without_executing() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let looped = call_with_args("looped_tool", json!({"a": 1})); + ctx.tool_scope.tool_tracker.record_call(looped.clone()); + ctx.tool_scope.tool_tracker.record_call(looped.clone()); + let calls = vec![ + looped, + ToolCall::new("other_tool".into(), json!({}), Some("id-2".into())), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 2); + let alert = results[0].output.as_str().unwrap(); + assert!(alert.starts_with("{\"tool_call_loop_alert\":"), "{alert}"); + let err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_tool_calls_truncates_with_global_max_chars() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(50), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = results[0].output.as_str().unwrap(); + assert!( + out.starts_with("[truncated: tool output exceeded 50 chars]\n"), + "{out}" + ); + } + + #[test] + fn eval_tool_calls_agent_max_chars_overrides_global() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(5000), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig { + max_tool_result_chars: Some(30), + ..Default::default() + })); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = results[0].output.as_str().unwrap(); + assert!( + out.starts_with("[truncated: tool output exceeded 30 chars]\n"), + "{out}" + ); + } + + #[test] + fn eval_tool_calls_zero_max_chars_disables_truncation() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(0), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert!(results[0].output["tool_call_error"].is_string()); + assert!(!results[0].output.to_string().contains("[truncated")); + } + + #[test] + fn eval_tool_calls_no_max_chars_configured_never_truncates() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert!(results[0].output["tool_call_error"].is_string()); + assert!(!results[0].output.to_string().contains("[truncated")); + } + + /// Pins current behavior: when the char cap lands inside a multi-byte + /// UTF-8 character of the serialized output, no prefix can be taken, so + /// the truncation marker is prepended to the FULL original output and the + /// "truncated" result is longer than the input. + #[test] + fn truncate_if_needed_utf8_boundary_returns_full_output_with_marker() { + let serialized = json!("aé").to_string(); + let result = ToolResult::new(call("t", Some("id-1")), json!("aé")); + + let truncated = result.truncate_if_needed(3); + + let out = truncated.output.as_str().unwrap(); + assert_eq!( + out, + format!("[truncated: tool output exceeded 3 chars]\n{serialized}") + ); + assert!(out.len() > serialized.len()); + } + + #[test] + fn eval_routes_agent_prefix_to_supervisor_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_supervisor_functions(); + + let out = + run_async(call_with_args("agent__check", json!({"id": "x"})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Supervisor tool failed"), "{err}"); + assert!(err.contains("No supervisor active"), "{err}"); + } + + #[test] + fn eval_routes_todo_prefix_to_todo_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_todo_functions(); + + let out = run_async(call_with_args("todo__list", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Todo tool failed"), "{err}"); + assert!(err.contains("Auto-continue is not enabled"), "{err}"); + } + + #[test] + fn eval_routes_memory_prefix_to_memory_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_memory_functions(); + + let out = run_async(call_with_args("memory__read", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Memory tool failed"), "{err}"); + assert!(err.contains("name is required"), "{err}"); + } + + #[test] + fn eval_routes_skill_prefix_to_skill_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_skill_functions(); + + let out = run_async(call_with_args("skill__load", json!({})).eval(&mut ctx)).unwrap(); + + assert_eq!(out["error"], "name is required"); + } + + #[test] + fn eval_routes_user_prefix_to_user_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_user_interaction_functions(); + + let out = run_async(call_with_args("user__confirm", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("User interaction failed"), "{err}"); + assert!(err.contains("'question' is required"), "{err}"); + } + + #[test] + fn eval_routes_rag_prefix_to_rag_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_rag_query_functions(); + + let out = + run_async(call_with_args("rag__query", json!({"query": "x"})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("RAG query failed"), "{err}"); + assert!(err.contains("No RAG is attached"), "{err}"); + } + + #[test] + fn eval_unknown_name_errors_with_unexpected_call() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + + let err = run_async(call_with_args("nope", json!({})).eval(&mut ctx)).unwrap_err(); + + assert!(err.to_string().contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_mcp_empty_runtime_returns_distinct_error_per_prefix() { + let ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let cases = [ + ( + "mcp_invoke_ghost", + json!({"tool": "t"}), + "MCP tool invocation failed", + ), + ( + "mcp_search_ghost", + json!({"query": "q"}), + "MCP search failed", + ), + ( + "mcp_describe_ghost", + json!({"tool": "t"}), + "MCP describe failed", + ), + ( + "mcp_read_ghost", + json!({"uri": "file:///x"}), + "MCP read failed", + ), + ( + "mcp_prompt_ghost", + json!({"prompt": "p"}), + "MCP prompt failed", + ), + ]; + + for (name, args, expected) in cases { + let out = run_async(call_with_args(name, args).eval_mcp(&ctx)).unwrap(); + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with(expected), "{name}: {err}"); + } + } } diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 32c5d48..87dfbe5 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -1473,14 +1473,24 @@ mod tests { } fn register_fake_agent(ctx: &mut RequestContext, id: &str, name: &str) { + register_fake_agent_with_output(ctx, id, name, "fake output"); + } + + fn register_fake_agent_with_output( + ctx: &mut RequestContext, + id: &str, + name: &str, + output: &str, + ) { let rt = tokio::runtime::Runtime::new().unwrap(); let id_owned = id.to_string(); let name_owned = name.to_string(); + let output_owned = output.to_string(); let join_handle = rt.spawn(async move { Ok(AgentResult { id: id_owned, agent_name: name_owned, - output: "fake output".into(), + output: output_owned, exit_status: AgentExitStatus::Completed, }) }); @@ -1511,6 +1521,48 @@ mod tests { .block_on(f) } + fn register_running_agent(ctx: &mut RequestContext, id: &str, name: &str) -> AbortSignal { + let abort = create_abort_signal(); + let id_owned = id.to_string(); + let name_owned = name.to_string(); + let join_handle = tokio::spawn(async move { + time::sleep(Duration::from_secs(60)).await; + Ok(AgentResult { + id: id_owned, + agent_name: name_owned, + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: id.to_string(), + agent_name: name.to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: abort.clone(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + abort + } + + fn wait_until_finished(ctx: &RequestContext, id: &str) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while ctx.supervisor.as_ref().unwrap().read().is_finished(id) != Some(true) { + assert!( + std::time::Instant::now() < deadline, + "agent '{id}' never finished" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + #[tokio::test] async fn sync_agent_functions_gates_meta_functions_on_live_capabilities() { let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); @@ -2127,4 +2179,358 @@ mod tests { } }); } + + #[test] + fn handle_collect_finished_agent_returns_output_and_consumes_handle() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + ctx.pending_agents_guardrail_count = 2; + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["id"], "a1"); + assert_eq!(result["agent"], "explore"); + assert_eq!(result["exit_status"], "Completed"); + assert_eq!(result["output"], "fake output"); + assert_eq!(ctx.pending_agents_guardrail_count, 0); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + } + + #[test] + fn handle_collect_pending_escalations_early_out_keeps_handle() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let _abort = register_running_agent(&mut ctx, "slow", "test"); + let queue = ctx.ensure_root_escalation_queue(); + let (tx, _rx) = tokio::sync::oneshot::channel(); + queue.submit(EscalationRequest { + id: "esc_1".into(), + from_agent_id: "a1".into(), + from_agent_name: "explore".into(), + question: "What do?".into(), + options: None, + reply_tx: tx, + }); + + let result = handle_collect(&mut ctx, &json!({"id": "slow"})) + .await + .unwrap(); + + assert_eq!(result["status"], "pending"); + assert!(result["pending_escalations"].is_array()); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("slow"), + Some(false) + ); + }); + } + + #[test] + fn handle_collect_unknown_agent_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + let result = run_async(handle_collect(&mut ctx, &json!({"id": "missing"}))).unwrap(); + assert_eq!(result["status"], "error"); + assert!(result["message"].as_str().unwrap().contains("not found")); + } + + #[test] + fn handle_collect_without_agent_passes_long_output_through_verbatim() { + let mut ctx = ctx_with_supervisor(4, 3); + let long_output = "x".repeat(10_000); + register_fake_agent_with_output(&mut ctx, "a1", "explore", &long_output); + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["output"], long_output); + } + + #[test] + fn handle_collect_output_below_agent_threshold_passes_through() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + summarization_threshold: 1_000_000, + ..Default::default() + })); + register_fake_agent(&mut ctx, "a1", "explore"); + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["output"], "fake output"); + } + + #[test] + fn handle_collect_over_threshold_with_unknown_summarization_model_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + summarization_threshold: 1, + summarization_model: Some("nonexistent_client:model".into()), + ..Default::default() + })); + register_fake_agent(&mut ctx, "a1", "explore"); + + let err = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap_err(); + + assert!(err.to_string().contains("nonexistent_client")); + } + + #[test] + fn guardrail_no_supervisor_is_no_action_and_resets_counter() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.pending_agents_guardrail_count = 2; + + assert!(matches!( + check_pending_agents_guardrail(&mut ctx), + GuardrailAction::NoAction + )); + assert_eq!(ctx.pending_agents_guardrail_count, 0); + } + + /// Pins current behavior: a finished-but-uncollected agent is not counted + /// as pending (only still-running agents are), so the turn-end guardrail + /// takes no action and the finished agent's result can be silently dropped. + #[test] + fn guardrail_ignores_finished_but_uncollected_agents() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + ctx.pending_agents_guardrail_count = 2; + + assert!(matches!( + check_pending_agents_guardrail(&mut ctx), + GuardrailAction::NoAction + )); + assert_eq!(ctx.pending_agents_guardrail_count, 0); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + Some(true) + ); + } + + #[test] + fn guardrail_force_terminates_at_max_and_cancels_agents() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let abort = register_running_agent(&mut ctx, "slow", "test"); + ctx.pending_agents_guardrail_count = PENDING_AGENTS_GUARDRAIL_MAX; + + match check_pending_agents_guardrail(&mut ctx) { + GuardrailAction::ForceTerminate(ids) => { + assert_eq!(ids, vec!["slow".to_string()]); + } + _ => panic!("expected ForceTerminate action"), + } + assert_eq!(ctx.pending_agents_guardrail_count, 0); + assert!(abort.aborted()); + }); + } + + #[test] + fn guardrail_injects_prompt_and_increments_counter_below_max() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let _abort = register_running_agent(&mut ctx, "slow", "test"); + ctx.pending_agents_guardrail_count = 1; + + match check_pending_agents_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("slow")); + assert!(prompt.contains("agent__collect")); + } + _ => panic!("expected Inject action"), + } + assert_eq!(ctx.pending_agents_guardrail_count, 2); + }); + } + + #[test] + fn handle_cancel_resets_guardrail_counter() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + ctx.pending_agents_guardrail_count = 2; + + let result = run_async(handle_cancel(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "ok"); + assert_eq!(ctx.pending_agents_guardrail_count, 0); + } + + #[test] + fn handle_spawn_missing_agent_arg_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + let err = run_async(handle_spawn(&mut ctx, &json!({}))).unwrap_err(); + assert!(err.to_string().contains("'agent' is required")); + } + + #[test] + fn handle_spawn_missing_prompt_arg_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + let err = run_async(handle_spawn(&mut ctx, &json!({"agent": "explore"}))).unwrap_err(); + assert!(err.to_string().contains("'prompt' is required")); + } + + #[test] + fn handle_spawn_rejects_agent_outside_whitelist() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + spawnable_agents: Some(vec!["allowed".into()]), + ..Default::default() + })); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "notallowed", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("spawnable_agents") + ); + } + + #[test] + fn handle_spawn_at_capacity_errors() { + let mut ctx = ctx_with_supervisor(1, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert_eq!( + result["message"], + "At capacity: 1/1 agents running. Wait for one to finish or cancel one." + ); + } + + #[test] + fn handle_spawn_exceeding_depth_errors() { + let mut ctx = ctx_with_supervisor(4, 0); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("Max agent depth exceeded") + ); + } + + #[test] + fn handle_spawn_no_supervisor_errors() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + let err = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap_err(); + assert!(err.to_string().contains("No supervisor active")); + } + + /// Pins current behavior: checking a finished agent does not report a + /// "finished, ready to collect" status; it silently delegates to collect, + /// returning the full result and consuming the handle. + #[test] + fn handle_check_finished_agent_delegates_to_collect_and_consumes_handle() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + + let result = run_async(handle_check(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["output"], "fake output"); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + } + + #[test] + fn handle_cancel_running_agent_aborts_and_waits_for_cleanup() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let sig = create_abort_signal(); + let sig2 = sig.clone(); + let join_handle = tokio::spawn(async move { + loop { + if sig2.aborted() { + return Ok(AgentResult { + id: "a1".into(), + agent_name: "explore".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }); + } + time::sleep(Duration::from_millis(10)).await; + } + }); + let handle = AgentHandle { + id: "a1".into(), + agent_name: "explore".into(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: sig.clone(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + ctx.pending_agents_guardrail_count = 2; + + let result = handle_cancel(&mut ctx, &json!({"id": "a1"})).await.unwrap(); + + assert_eq!(result["status"], "ok"); + let message = result["message"].as_str().unwrap(); + assert!(message.contains("Cancelled agent 'explore'")); + assert!(message.contains("waited for cleanup")); + assert!(sig.aborted()); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + assert_eq!(ctx.pending_agents_guardrail_count, 0); + }); + } } diff --git a/src/supervisor/mod.rs b/src/supervisor/mod.rs index dc55d21..0739404 100644 --- a/src/supervisor/mod.rs +++ b/src/supervisor/mod.rs @@ -294,4 +294,25 @@ mod tests { AgentExitStatus::Failed("x".into()) ); } + + #[test] + fn cancel_recursive_aborts_nested_supervisors() { + let child_sig = create_abort_signal(); + let mut child_handle = make_handle("c1", "worker", 2); + child_handle.abort_signal = child_sig.clone(); + let mut child_sup = Supervisor::new(4, 3); + child_sup.register(child_handle).unwrap(); + + let parent_sig = create_abort_signal(); + let mut parent_handle = make_handle("a1", "explore", 1); + parent_handle.abort_signal = parent_sig.clone(); + parent_handle.child_supervisor = Some(Arc::new(RwLock::new(child_sup))); + let mut sup = Supervisor::new(4, 3); + sup.register(parent_handle).unwrap(); + + sup.cancel_recursive(); + + assert!(parent_sig.aborted()); + assert!(child_sig.aborted()); + } }