Return agent snapshots from close_agent (#47252)

## What changed

Return an `AgentInfo` snapshot captured before shutdown, including metadata for known unloaded agents. Use the snapshot's status and metadata in the tool's completion event and its status in `previous_status`, replacing the separate status subscription.

Preserve the pre-close status for error reporting when shutdown removes the target before a descendant fails to close. Propagate unexpected thread lookup errors before attempting shutdown.

GitOrigin-RevId: bdbd877aaf0389e830d6a7e84a5e14bb7d47277a
This commit is contained in:
jif
2026-09-22 10:58:32 +00:00
committed by copyberry
parent c117207a6f
commit fdd89e78ac
2 changed files with 28 additions and 55 deletions
+18 -11
View File
@@ -1,4 +1,5 @@
use super::*;
use crate::agent::api::AgentInfo;
use codex_protocol::error::CodexErrorDetails;
use codex_thread_store::PersistContext;
@@ -45,12 +46,19 @@ impl LocalAgentControl {
/// Mark `agent_id` as explicitly closed in persisted spawn-edge state, then shut down the
/// agent and any live descendants reached from the in-memory tree.
pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult<AgentInfo> {
let state = self.upgrade()?;
let known_agent = self.state.agent_metadata_for_thread(agent_id).is_some();
match state.get_thread(agent_id).await {
let metadata = self.get_agent_metadata(agent_id);
let known_agent = metadata.is_some();
let snapshot = match state.get_thread(agent_id).await {
Ok(thread) => {
if !thread.config_snapshot().await.ephemeral
let agent = LiveAgent {
thread_id: agent_id,
metadata: metadata.unwrap_or_default(),
status: thread.agent_status().await,
};
let config = Box::new(thread.config_snapshot().await);
if !config.ephemeral
&& let Some(agent_graph_store) = state.agent_graph_store()
&& let Err(err) = agent_graph_store
.set_thread_spawn_edge_status(
@@ -61,6 +69,7 @@ impl LocalAgentControl {
{
warn!("failed to persist thread-spawn edge status for {agent_id}: {err}");
}
AgentInfo::Loaded { agent, config }
}
Err(err)
if known_agent && matches!(err.details(), CodexErrorDetails::ThreadNotFound(_)) =>
@@ -77,12 +86,10 @@ impl LocalAgentControl {
"failed to persist stale thread-spawn edge status for {agent_id}: {err}"
)));
}
AgentInfo::Unloaded(metadata.unwrap_or_default())
}
Err(err) if matches!(err.details(), CodexErrorDetails::ThreadNotFound(_)) => {}
Err(err) => {
warn!("failed to inspect agent before close {agent_id}: {err}");
}
}
Err(err) => return Err(err),
};
match Box::pin(self.shutdown_agent_tree(agent_id)).await {
Err(err)
if known_agent
@@ -91,9 +98,9 @@ impl LocalAgentControl {
CodexErrorDetails::ThreadNotFound(_) | CodexErrorDetails::InternalAgentDied
) =>
{
Ok(String::new())
Ok(snapshot)
}
result => result,
result => result.map(|_| snapshot),
}
}
@@ -1,6 +1,5 @@
use super::*;
use crate::tools::handlers::multi_agents_spec::create_close_agent_tool_v1;
use codex_protocol::error::CodexErrorDetails;
use codex_tools::ToolSpec;
pub(crate) struct Handler;
@@ -43,7 +42,6 @@ async fn handle_close_agent(
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);
let known_agent = receiver_agent.is_some();
let receiver_agent = receiver_agent.unwrap_or_default();
session
.emit_turn_item_started(
@@ -62,48 +60,16 @@ async fn handle_close_agent(
}),
)
.await;
let status = match session
.services
.agent_control
.subscribe_status(agent_id)
.await
{
Ok(mut status_rx) => status_rx.borrow_and_update().clone(),
Err(err)
if known_agent && matches!(err.details(), CodexErrorDetails::ThreadNotFound(_)) =>
{
session.services.agent_control.get_status(agent_id).await
}
Err(err) => {
let status = session.services.agent_control.get_status(agent_id).await;
session
.emit_turn_item_completed(
&turn,
TurnItem::CollabAgentToolCall(CollabAgentToolCallItem {
id: call_id.clone(),
tool: CollabAgentTool::CloseAgent,
status: collab_tool_call_status(&status, Some(agent_id)),
sender_thread_id: session.thread_id(),
receiver_thread_ids: vec![agent_id],
receiver_agents: vec![CollabAgentRef {
thread_id: agent_id,
agent_nickname: receiver_agent.agent_nickname.clone(),
agent_role: receiver_agent.agent_role.clone(),
}],
prompt: None,
model: None,
reasoning_effort: None,
agents_states: [(agent_id, status)].into_iter().collect(),
}),
)
.await;
return Err(collab_agent_error(agent_id, err));
}
// Shutdown may remove the target before a descendant fails to close.
let previous_status = session.services.agent_control.get_status(agent_id).await;
let result = Box::pin(session.services.agent_control.close_agent(agent_id)).await;
let (status, receiver_agent) = match &result {
Ok(snapshot) => (
snapshot.status().cloned().unwrap_or(AgentStatus::NotFound),
snapshot.metadata().clone(),
),
Err(_) => (previous_status, receiver_agent),
};
let result = Box::pin(session.services.agent_control.close_agent(agent_id))
.await
.map_err(|err| collab_agent_error(agent_id, err))
.map(|_| ());
session
.emit_turn_item_completed(
&turn,
@@ -125,7 +91,7 @@ async fn handle_close_agent(
}),
)
.await;
result?;
result.map_err(|err| collab_agent_error(agent_id, err))?;
Ok(CloseAgentResult {
previous_status: status,