feat: make v2 close_agent interrupt only

This commit is contained in:
jif-oai
2026-06-05 15:33:17 +02:00
parent 0ae308399d
commit 83bbc0d8db
3 changed files with 155 additions and 65 deletions

View File

@@ -308,13 +308,13 @@ pub fn create_close_agent_tool_v2() -> ToolSpec {
let properties = BTreeMap::from([(
"target".to_string(),
JsonSchema::string(Some(
"Agent id or canonical task name to close (from spawn_agent).".to_string(),
"Agent id or canonical task name to interrupt (from spawn_agent).".to_string(),
)),
)]);
ToolSpec::Function(ResponsesApiTool {
name: "close_agent".to_string(),
description: "Close an agent and any open descendants when they are no longer needed, and return the target agent's previous status before shutdown was requested. Idle agents do not count toward the concurrency limit.".to_string(),
description: "Interrupt an agent's current turn, if any, and return its previous status. The agent remains available for messages and follow-up tasks.".to_string(),
strict: false,
defer_loading: None,
parameters: JsonSchema::object(properties, Some(vec!["target".to_string()]), Some(false.into())),
@@ -485,7 +485,7 @@ fn close_agent_output_schema() -> Value {
"type": "object",
"properties": {
"previous_status": {
"description": "The agent status observed before shutdown was requested.",
"description": "The agent status observed before the request was handled.",
"allOf": [agent_status_output_schema()]
}
},

View File

@@ -1680,6 +1680,82 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() {
assert_eq!(result.agents[0].last_task_message.as_deref(), Some("build"));
}
#[tokio::test]
async fn multi_agent_v2_list_agents_keeps_interrupted_resident_agents() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let root = manager
.start_thread((*turn.config).clone())
.await
.expect("root thread should start");
session.services.agent_control = manager.agent_control();
session.thread_id = root.thread_id;
let mut config = (*turn.config).clone();
let _ = config.features.enable(Feature::MultiAgentV2);
set_turn_config(&mut turn, config);
let session = Arc::new(session);
let turn = Arc::new(turn);
let spawn_output = 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 _ = expect_text_output(spawn_output);
let agent_id = session
.services
.agent_control
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
.await
.expect("worker path should resolve");
let agent_path = session
.services
.agent_control
.get_agent_metadata(agent_id)
.expect("worker metadata should exist")
.agent_path
.expect("worker path should exist");
let close_output = CloseAgentHandlerV2
.handle(invocation(
session.clone(),
turn.clone(),
"close_agent",
function_payload(json!({"target": "worker"})),
))
.await
.expect("close_agent should succeed");
let _ = expect_text_output(close_output);
let output = ListAgentsHandlerV2
.handle(invocation(
session,
turn,
"list_agents",
function_payload(json!({})),
))
.await
.expect("list_agents should succeed");
let (content, _) = expect_text_output(output);
let result: ListAgentsResult =
serde_json::from_str(&content).expect("list_agents result should be json");
assert_eq!(result.agents.len(), 2);
assert_eq!(result.agents[0].agent_name, "/root");
assert_eq!(
result.agents[0].last_task_message.as_deref(),
Some("Main thread")
);
assert_eq!(result.agents[1].agent_name, agent_path.as_str());
}
#[tokio::test]
async fn multi_agent_v2_send_message_rejects_legacy_items_field() {
let (mut session, mut turn) = make_session_and_context().await;
@@ -3672,11 +3748,34 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() {
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
.await
.expect("worker path should resolve");
let worker_thread = manager
.get_thread(agent_id)
.await
.expect("worker thread should be resident");
let worker_session = worker_thread.codex.session.clone();
SpawnAgentHandlerV2::default()
.handle(invocation(
worker_session.clone(),
worker_session.new_default_turn().await,
"spawn_agent",
function_payload(json!({
"message": "inspect a child task",
"task_name": "child"
})),
))
.await
.expect("child spawn should succeed");
let child_id = session
.services
.agent_control
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker/child")
.await
.expect("child path should resolve");
let output = CloseAgentHandlerV2
.handle(invocation(
session,
turn,
session.clone(),
turn.clone(),
"close_agent",
function_payload(json!({"target": "worker"})),
))
@@ -3688,13 +3787,38 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() {
assert_ne!(result.previous_status, AgentStatus::NotFound);
assert_eq!(success, Some(true));
assert_eq!(
manager.agent_control().get_status(agent_id).await,
AgentStatus::NotFound
session
.services
.agent_control
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
.await
.expect("worker path should remain resolvable"),
agent_id
);
manager
.get_thread(agent_id)
.await
.expect("worker should remain resident");
manager
.get_thread(child_id)
.await
.expect("child should remain resident");
let ops = manager.captured_ops();
assert!(
ops.iter()
.any(|(thread_id, op)| *thread_id == agent_id && matches!(op, Op::Interrupt))
);
assert!(!ops.iter().any(|(thread_id, op)| {
(*thread_id == agent_id || *thread_id == child_id) && matches!(op, Op::Shutdown)
}));
assert!(
!ops.iter()
.any(|(thread_id, op)| *thread_id == child_id && matches!(op, Op::Interrupt))
);
}
#[tokio::test]
async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
async fn multi_agent_v2_close_agent_accepts_unloaded_task_name_target() {
let (mut session, mut turn) = make_session_and_context().await;
let mut config = (*turn.config).clone();
config.multi_agent_v2.max_concurrent_threads_per_session = 2;
@@ -3763,7 +3887,7 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
function_payload(json!({"target": "worker"})),
))
.await
.expect("close_agent should reap stale v2 task names");
.expect("close_agent should accept unloaded 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");
@@ -3777,7 +3901,7 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
)
.await
.expect("open children should load");
assert_eq!(open_children, Vec::<ThreadId>::new());
assert_eq!(open_children, vec![agent_id]);
let closed_children = state_db
.list_thread_spawn_children_with_status(
root.thread_id,
@@ -3785,32 +3909,22 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
)
.await
.expect("closed children should load");
assert_eq!(closed_children, vec![agent_id]);
assert_eq!(closed_children, Vec::<ThreadId>::new());
SpawnAgentHandlerV2::default()
let output = ListAgentsHandlerV2
.handle(invocation(
session.clone(),
turn.clone(),
"spawn_agent",
function_payload(json!({
"message": "inspect this repo again",
"task_name": "replacement"
})),
"list_agents",
function_payload(json!({})),
))
.await
.expect("spawn_agent should succeed after stale close releases the slot");
let replacement_id = session
.services
.agent_control
.resolve_agent_reference(session.thread_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");
.expect("list_agents should succeed");
let (content, _) = expect_text_output(output);
let result: ListAgentsResult =
serde_json::from_str(&content).expect("list_agents result should be json");
assert_eq!(result.agents.len(), 1);
assert_eq!(result.agents[0].agent_name, "/root");
}
#[tokio::test]

View File

@@ -37,9 +37,12 @@ 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);
let known_agent = receiver_agent.is_some();
let receiver_agent = receiver_agent.unwrap_or_default();
let receiver_agent = session
.services
.agent_control
.ensure_agent_known(agent_id)
.await
.map_err(|err| collab_agent_error(agent_id, err))?;
if receiver_agent
.agent_path
.as_ref()
@@ -67,43 +70,16 @@ async fn handle_close_agent(
.into(),
)
.await;
let status = match session
let status = session.services.agent_control.get_status(agent_id).await;
let result = match session
.services
.agent_control
.subscribe_status(agent_id)
.interrupt_agent(agent_id)
.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
.send_event(
&turn,
CollabCloseEndEvent {
call_id: call_id.clone(),
completed_at_ms: now_unix_timestamp_ms(),
sender_thread_id: session.thread_id,
receiver_thread_id: agent_id,
receiver_agent_nickname: receiver_agent.agent_nickname.clone(),
receiver_agent_role: receiver_agent.agent_role.clone(),
status,
}
.into(),
)
.await;
return Err(collab_agent_error(agent_id, err));
}
Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => Ok(()),
Err(err) => Err(collab_agent_error(agent_id, err)),
};
let result = session
.services
.agent_control
.close_agent(agent_id)
.await
.map_err(|err| collab_agent_error(agent_id, err))
.map(|_| ());
session
.send_event(
&turn,