diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 8f25074d76f0..f72f34cd4e6b 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -788,15 +788,45 @@ impl AgentControl { /// agent and any live descendants reached from the in-memory tree. pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; - if let Ok(thread) = state.get_thread(agent_id).await - && let Some(state_db_ctx) = thread.state_db() - && let Err(err) = state_db_ctx - .set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed) - .await - { - warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); + let known_agent = self.state.agent_metadata_for_thread(agent_id).is_some(); + match state.get_thread(agent_id).await { + Ok(thread) => { + if let Some(state_db_ctx) = thread.state_db() + && let Err(err) = state_db_ctx + .set_thread_spawn_edge_status( + agent_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + { + warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); + } + } + Err(CodexErr::ThreadNotFound(_)) if known_agent => { + if let Some(state_db_ctx) = state.state_db() + && let Err(err) = state_db_ctx + .set_thread_spawn_edge_status( + agent_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + { + return Err(CodexErr::Fatal(format!( + "failed to persist stale thread-spawn edge status for {agent_id}: {err}" + ))); + } + } + Err(CodexErr::ThreadNotFound(_)) => {} + Err(err) => { + warn!("failed to inspect agent before close {agent_id}: {err}"); + } + } + match Box::pin(self.shutdown_agent_tree(agent_id)).await { + Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) if known_agent => { + Ok(String::new()) + } + result => result, } - Box::pin(self.shutdown_agent_tree(agent_id)).await } /// Shut down `agent_id` and any live descendants reachable from the in-memory spawn tree. diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index de069854f106..bcfa5231b9d2 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -1,6 +1,7 @@ use super::*; use crate::tools::handlers::multi_agents_spec::create_close_agent_tool_v1; use crate::turn_timing::now_unix_timestamp_ms; +use codex_protocol::error::CodexErr; use codex_tools::ToolSpec; pub(crate) struct Handler; @@ -36,11 +37,9 @@ async fn handle_close_agent( let arguments = function_arguments(payload)?; let args: CloseAgentArgs = parse_arguments(&arguments)?; let agent_id = parse_agent_id_target(&args.target)?; - let receiver_agent = session - .services - .agent_control - .get_agent_metadata(agent_id) - .unwrap_or_default(); + let receiver_agent = session.services.agent_control.get_agent_metadata(agent_id); + let known_agent = receiver_agent.is_some(); + let receiver_agent = receiver_agent.unwrap_or_default(); session .send_event( &turn, @@ -60,6 +59,9 @@ async fn handle_close_agent( .await { Ok(mut status_rx) => status_rx.borrow_and_update().clone(), + Err(CodexErr::ThreadNotFound(_)) if known_agent => { + session.services.agent_control.get_status(agent_id).await + } Err(err) => { let status = session.services.agent_control.get_status(agent_id).await; session 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 68b19adea9e0..164d5410e0b0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -52,6 +52,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::user_input::UserInput; +use codex_state::DirectionalThreadSpawnEdgeStatus; use core_test_support::TempDirExt; use pretty_assertions::assert_eq; use serde::Deserialize; @@ -3751,6 +3752,126 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() { ); } +#[tokio::test] +async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() { + let (mut session, mut turn) = make_session_and_context().await; + let mut config = (*turn.config).clone(); + config.agent_max_threads = Some(1); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow sqlite"); + let state_db = init_state_db(&config) + .await + .expect("sqlite state db should initialize"); + let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + Some(state_db.clone()), + ); + let root = manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + session.services.agent_control = manager.agent_control(); + session.conversation_id = root.thread_id; + turn.config = Arc::new(config.clone()); + + let session = Arc::new(session); + let turn = Arc::new(turn); + SpawnAgentHandlerV2::default() + .handle(invocation( + session.clone(), + turn.clone(), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo", + "task_name": "worker" + })), + )) + .await + .expect("spawn_agent should succeed"); + + let agent_id = session + .services + .agent_control + .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker") + .await + .expect("worker path should resolve"); + let stale_thread = manager + .remove_thread(&agent_id) + .await + .expect("worker thread should be loaded before removal"); + stale_thread + .submit(Op::Shutdown {}) + .await + .expect("removed worker thread should still accept shutdown"); + stale_thread.wait_until_terminated().await; + + let output = CloseAgentHandlerV2 + .handle(invocation( + session.clone(), + turn.clone(), + "close_agent", + function_payload(json!({"target": "worker"})), + )) + .await + .expect("close_agent should reap stale v2 task names"); + let (content, success) = expect_text_output(output); + let result: close_agent::CloseAgentResult = + serde_json::from_str(&content).expect("close_agent result should be json"); + assert_eq!(result.previous_status, AgentStatus::NotFound); + assert_eq!(success, Some(true)); + + let open_children = state_db + .list_thread_spawn_children_with_status( + root.thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open children should load"); + assert_eq!(open_children, Vec::::new()); + let closed_children = state_db + .list_thread_spawn_children_with_status( + root.thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed children should load"); + assert_eq!(closed_children, vec![agent_id]); + + SpawnAgentHandlerV2::default() + .handle(invocation( + session.clone(), + turn.clone(), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo again", + "task_name": "replacement" + })), + )) + .await + .expect("spawn_agent should succeed after stale close releases the slot"); + let replacement_id = session + .services + .agent_control + .resolve_agent_reference(session.conversation_id, &turn.session_source, "replacement") + .await + .expect("replacement path should resolve"); + let _ = session + .services + .agent_control + .shutdown_live_agent(replacement_id) + .await + .expect("replacement should shut down"); +} + #[tokio::test] async fn multi_agent_v2_close_agent_rejects_root_target_and_id() { let (mut session, mut turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs index 088e5310623a..15502572a509 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs @@ -1,6 +1,7 @@ use super::*; use crate::tools::handlers::multi_agents_spec::create_close_agent_tool_v2; use crate::turn_timing::now_unix_timestamp_ms; +use codex_protocol::error::CodexErr; use codex_tools::ToolSpec; pub(crate) struct Handler; @@ -36,11 +37,9 @@ async fn handle_close_agent( let arguments = function_arguments(payload)?; let args: CloseAgentArgs = parse_arguments(&arguments)?; let agent_id = resolve_agent_target(&session, &turn, &args.target).await?; - let receiver_agent = session - .services - .agent_control - .get_agent_metadata(agent_id) - .unwrap_or_default(); + let receiver_agent = session.services.agent_control.get_agent_metadata(agent_id); + let known_agent = receiver_agent.is_some(); + let receiver_agent = receiver_agent.unwrap_or_default(); if receiver_agent .agent_path .as_ref() @@ -69,6 +68,9 @@ async fn handle_close_agent( .await { Ok(mut status_rx) => status_rx.borrow_and_update().clone(), + Err(CodexErr::ThreadNotFound(_)) if known_agent => { + session.services.agent_control.get_status(agent_id).await + } Err(err) => { let status = session.services.agent_control.get_status(agent_id).await; session