diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 87ce5afec9..93866cf089 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::AuthManager; use crate::CodexAuth; use crate::ThreadManager; +use crate::agent::WatchdogRegistration; use crate::built_in_model_providers; use crate::codex::make_session_and_context; use crate::config::DEFAULT_AGENT_MAX_DEPTH; @@ -30,6 +31,7 @@ use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHand use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2; use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2; use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2; +use crate::tools::handlers::multi_agents_v2::WatchdogSelfCloseHandlerV2; use crate::turn_diff_tracker::TurnDiffTracker; use codex_features::Feature; use codex_protocol::AgentPath; @@ -89,6 +91,54 @@ fn thread_manager() -> ThreadManager { ) } +async fn attach_watchdog_helper_for_tests( + manager: &ThreadManager, + agent_control: &crate::agent::AgentControl, + config: &crate::config::Config, +) -> ThreadId { + let owner = manager + .start_thread(config.clone()) + .await + .expect("owner thread should start"); + let target = agent_control + .spawn_agent( + config.clone(), + Op::UserInput { + items: Vec::new(), + final_output_json_schema: None, + }, + None, + ) + .await + .expect("watchdog target thread should start"); + let helper = manager + .start_thread(config.clone()) + .await + .expect("helper thread should start"); + agent_control + .register_watchdog(WatchdogRegistration { + owner_thread_id: owner.thread_id, + target_thread_id: target, + child_depth: 1, + interval_s: 30, + prompt: "check in".to_string(), + config: config.clone(), + }) + .await + .expect("watchdog registration should succeed"); + agent_control + .set_watchdog_active_helper_for_tests(target, helper.thread_id) + .await; + assert_eq!( + agent_control + .watchdog_owner_for_active_helper(helper.thread_id) + .await, + Some(owner.thread_id), + "watchdog helper should be registered for owner" + ); + helper.thread_id +} + fn history_contains_inter_agent_communication( history_items: &[ResponseItem], expected: &InterAgentCommunication, @@ -175,6 +225,11 @@ struct ListedAgentResult { last_task_message: Option, } +#[derive(Debug, Deserialize, PartialEq)] +struct WatchdogSelfCloseResult { + previous_status: AgentStatus, +} + #[tokio::test] async fn handler_rejects_non_function_payloads() { let (session, turn) = make_session_and_context().await; @@ -3001,6 +3056,116 @@ async fn close_agent_submits_shutdown_and_returns_previous_status() { assert_eq!(status_after, AgentStatus::NotFound); } +#[tokio::test] +async fn watchdog_self_close_rejects_non_watchdog_thread() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + let agent_control = manager.agent_control(); + let thread = manager + .start_thread(turn.config.as_ref().clone()) + .await + .expect("thread should start"); + session.services.agent_control = agent_control.clone(); + session.conversation_id = thread.thread_id; + + let err = WatchdogSelfCloseHandler + .handle(invocation( + Arc::new(session), + Arc::new(turn), + "watchdog_self_close", + function_payload(json!({})), + )) + .await + .expect_err("non-watchdog threads should be rejected"); + + assert_eq!( + err, + FunctionCallError::RespondToModel( + "watchdog_self_close is only available in watchdog check-in threads.".to_string(), + ) + ); +} + +#[tokio::test] +async fn watchdog_self_close_closes_watchdog_helper_and_returns_previous_status() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + let agent_control = manager.agent_control(); + session.services.agent_control = agent_control.clone(); + let helper_thread_id = + attach_watchdog_helper_for_tests(&manager, &agent_control, turn.config.as_ref()).await; + session.conversation_id = helper_thread_id; + let status_before = agent_control.get_status(helper_thread_id).await; + + let output = WatchdogSelfCloseHandler + .handle(invocation( + Arc::new(session), + Arc::new(turn), + "watchdog_self_close", + function_payload(json!({})), + )) + .await + .expect("watchdog helper should be allowed to self-close"); + let (content, success) = expect_text_output(output); + let result: WatchdogSelfCloseResult = + serde_json::from_str(&content).expect("watchdog self-close result should be json"); + + assert_eq!( + result, + WatchdogSelfCloseResult { + previous_status: status_before, + } + ); + assert_eq!(success, Some(true)); + assert_eq!( + agent_control.get_status(helper_thread_id).await, + AgentStatus::NotFound + ); +} + +#[tokio::test] +async fn multi_agent_v2_watchdog_self_close_closes_watchdog_helper_and_returns_previous_status() { + let (mut session, mut turn) = make_session_and_context().await; + let manager = thread_manager(); + let agent_control = manager.agent_control(); + session.services.agent_control = agent_control.clone(); + let mut config = turn.config.as_ref().clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + let helper_thread_id = + attach_watchdog_helper_for_tests(&manager, &agent_control, &config).await; + session.conversation_id = helper_thread_id; + turn.config = Arc::new(config); + let status_before = agent_control.get_status(helper_thread_id).await; + + let output = WatchdogSelfCloseHandlerV2 + .handle(invocation( + Arc::new(session), + Arc::new(turn), + "watchdog_self_close", + function_payload(json!({})), + )) + .await + .expect("watchdog helper should be allowed to self-close"); + let (content, success) = expect_text_output(output); + let result: WatchdogSelfCloseResult = + serde_json::from_str(&content).expect("watchdog self-close result should be json"); + + assert_eq!( + result, + WatchdogSelfCloseResult { + previous_status: status_before, + } + ); + assert_eq!(success, Some(true)); + assert_eq!( + agent_control.get_status(helper_thread_id).await, + AgentStatus::NotFound + ); +} + #[tokio::test] async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed() { let (_session, turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index 4c9fb1bf7c..d0b7dec8d9 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs @@ -34,6 +34,7 @@ pub(crate) use list_agents::Handler as ListAgentsHandler; pub(crate) use send_message::Handler as SendMessageHandler; pub(crate) use spawn::Handler as SpawnAgentHandler; pub(crate) use wait::Handler as WaitAgentHandler; +pub(crate) use watchdog_self_close::Handler as WatchdogSelfCloseHandlerV2; mod assign_task; mod close_agent; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 4cd4952611..80a8c5f440 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -37,29 +37,25 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_tools::CommandToolOptions; +use codex_tools::DiscoverablePluginInfo; use codex_tools::DiscoverableTool; +use codex_tools::DiscoverableToolAction; use codex_tools::DiscoverableToolType; +use codex_tools::FreeformTool; +use codex_tools::FreeformToolFormat; use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ResponsesApiTool; use codex_tools::ShellToolOptions; use codex_tools::SpawnAgentToolOptions; -use codex_tools::ToolSearchAppInfo; -use codex_tools::ToolSuggestEntry; use codex_tools::ViewImageToolOptions; use codex_tools::WaitAgentTimeoutOptions; use codex_tools::augment_tool_spec_for_code_mode; use codex_tools::create_assign_task_tool; use codex_tools::create_close_agent_tool_v1; use codex_tools::create_close_agent_tool_v2; -use codex_tools::create_code_mode_tool; use codex_tools::create_exec_command_tool; -use codex_tools::create_js_repl_reset_tool; -use codex_tools::create_js_repl_tool; -use codex_tools::create_list_agents_tool; -use codex_tools::create_list_dir_tool; -use codex_tools::create_list_mcp_resource_templates_tool; -use codex_tools::create_list_mcp_resources_tool; -use codex_tools::create_read_mcp_resource_tool; +use codex_tools::create_list_agents_tool as create_list_agents_tool_v2; use codex_tools::create_report_agent_job_result_tool; use codex_tools::create_request_permissions_tool; use codex_tools::create_request_user_input_tool; @@ -71,33 +67,47 @@ use codex_tools::create_shell_tool; use codex_tools::create_spawn_agent_tool_v1; use codex_tools::create_spawn_agent_tool_v2; use codex_tools::create_spawn_agents_on_csv_tool; -use codex_tools::create_test_sync_tool; -use codex_tools::create_tool_search_tool; -use codex_tools::create_tool_suggest_tool; use codex_tools::create_view_image_tool; use codex_tools::create_wait_agent_tool_v1; use codex_tools::create_wait_agent_tool_v2; -use codex_tools::create_wait_tool; +use codex_tools::create_watchdog_self_close_tool; use codex_tools::create_write_stdin_tool; use codex_tools::dynamic_tool_to_responses_api_tool; use codex_tools::mcp_tool_to_responses_api_tool; use codex_tools::tool_spec_to_code_mode_tool_definition; use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_template::Template; use serde::Deserialize; use serde::Serialize; use std::collections::BTreeMap; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::LazyLock; pub type JsonSchema = codex_tools::JsonSchema; #[cfg(test)] pub(crate) use codex_tools::mcp_call_tool_result_output_schema; +const TOOL_SEARCH_DESCRIPTION_TEMPLATE_SOURCE: &str = + include_str!("../../templates/search_tool/tool_description.md"); +const TOOL_SEARCH_DESCRIPTION_TEMPLATE_KEY: &str = "app_descriptions"; +const WATCHDOG_TOOLS_NAMESPACE: &str = "watchdog"; +const WATCHDOG_TOOLS_NAMESPACE_DESCRIPTION: &str = + "Watchdog-only tools for parent-thread recovery and watchdog check-in lifecycle control."; +static TOOL_SEARCH_DESCRIPTION_TEMPLATE: LazyLock