agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: split agent control modules (#26610)

Browse Source 
## Summary

Mechanically splits `AgentControl` into focused modules so later agent
runtime changes are easier to review. The shared lookup, messaging, and
completion logic remains in `control.rs`, while spawn-specific code and
V1 legacy close/resume behavior move into dedicated files.

## Changes

- Extract spawn-agent code into `agent/control/spawn.rs`.
- Extract V1-only legacy close/resume behavior into
`agent/control/legacy.rs`.
- Keep shared control-plane behavior in `agent/control.rs`.
- Preserve existing behavior; this PR is intended to be mechanical.

## Stack

1. This PR - Mechanical `AgentControl` split: extracts spawn and V1
legacy code without behavior changes.
2. #26614 - Execution slot accounting: separates logical agents from
active execution slots.
3. #26611 - Residency and reload runtime: adds resident-agent LRU,
eviction/reload, durable lookup, and V2 delivery through reload.
4. #26612 - V2 tool semantics: narrows `close_agent` to interrupt-only
and updates V2 tool coverage.
This commit is contained in:
jif
2026-06-05 16:24:22 +02:00
committed by GitHub
Unverified
parent 66232220e2
commit 0b1512c2c8
3 changed files with 693 additions and 683 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+4 -683
View File
@@ -42,9 +42,11 @@ use std::sync::Weak;
use tokio::sync::watch;
use tracing::warn;
const AGENT_NAMES: &str = include_str!("agent_names.txt");
const ROOT_LAST_TASK_MESSAGE: &str = "Main thread";
mod legacy;
mod spawn;
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum SpawnAgentForkMode {
FullHistory,
@@ -59,11 +61,6 @@ pub(crate) struct SpawnAgentOptions {
pub(crate) environments: Option<Vec<TurnEnvironmentSelection>>,
}
struct SpawnAgentThreadInheritance {
shell_snapshot: Option<Arc<ShellSnapshot>>,
exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
}
#[derive(Clone, Debug)]
pub(crate) struct LiveAgent {
pub(crate) thread_id: ThreadId,
@@ -78,77 +75,6 @@ pub(crate) struct ListedAgent {
pub(crate) last_task_message: Option<String>,
}
fn default_agent_nickname_list() -> Vec<&'static str> {
AGENT_NAMES
.lines()
.map(str::trim)
.filter(|name| !name.is_empty())
.collect()
}
fn agent_nickname_candidates(config: &Config, role_name: Option<&str>) -> Vec<String> {
let role_name = role_name.unwrap_or(DEFAULT_ROLE_NAME);
if let Some(candidates) =
resolve_role_config(config, role_name).and_then(|role| role.nickname_candidates.clone())
{
return candidates;
}
default_agent_nickname_list()
.into_iter()
.map(ToOwned::to_owned)
.collect()
}
fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: bool) -> bool {
match item {
RolloutItem::ResponseItem(ResponseItem::Message { role, phase, .. }) => match role.as_str()
{
"system" | "developer" | "user" => true,
"assistant" => *phase == Some(MessagePhase::FinalAnswer),
_ => false,
},
RolloutItem::ResponseItem(
ResponseItem::AgentMessage { .. }
| ResponseItem::Reasoning { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::ToolSearchCall { .. }
| ResponseItem::FunctionCallOutput { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::CustomToolCallOutput { .. }
| ResponseItem::ToolSearchOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::ImageGenerationCall { .. }
| ResponseItem::Compaction { .. }
| ResponseItem::CompactionTrigger
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other,
) => false,
// Full-history forks preserve the cached prompt prefix and can keep diffing
// from the parent's durable baseline. Truncated forks drop part of that prompt,
// so they must rebuild context on their first child turn.
RolloutItem::TurnContext(_) => preserve_reference_context_item,
RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true,
}
}
fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &[String]) -> bool {
let ResponseItem::Message { role, content, .. } = item else {
return false;
};
if role != "developer" {
return false;
}
let [ContentItem::InputText { text }] = content.as_slice() else {
return false;
};
usage_hint_texts
.iter()
.any(|usage_hint_text| usage_hint_text == text)
}
/// Control-plane handle for multi-agent operations.
/// `AgentControl` is held by each session (via `SessionServices`). It provides capability to
/// spawn new agents and the inter-agent communication layer.
@@ -185,531 +111,6 @@ impl AgentControl {
self.session_id
}
/// Spawn a new agent thread and submit the initial prompt.
#[cfg(test)]
pub(crate) async fn spawn_agent(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
) -> CodexResult<ThreadId> {
let spawned_agent = Box::pin(self.spawn_agent_internal(
config,
initial_operation,
session_source,
SpawnAgentOptions::default(),
))
.await?;
Ok(spawned_agent.thread_id)
}
/// Spawn an agent thread with some metadata.
pub(crate) async fn spawn_agent_with_metadata(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
options: SpawnAgentOptions, // TODO(jif) drop with new fork.
) -> CodexResult<LiveAgent> {
Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options))
.await
}
async fn spawn_agent_internal(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
options: SpawnAgentOptions,
) -> CodexResult<LiveAgent> {
let state = self.upgrade()?;
let multi_agent_version = state
.effective_multi_agent_version_for_spawn(
&InitialHistory::New,
session_source.as_ref(),
options.parent_thread_id,
/*forked_from_thread_id*/ None,
&config,
)
.await;
let agent_max_threads = config.effective_agent_max_threads(multi_agent_version);
let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?;
let inheritance = SpawnAgentThreadInheritance {
shell_snapshot: self
.inherited_shell_snapshot_for_source(&state, session_source.as_ref())
.await,
exec_policy: self
.inherited_exec_policy_for_source(&state, session_source.as_ref(), &config)
.await,
};
let (session_source, mut agent_metadata) = match session_source {
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth,
agent_path,
agent_role,
..
})) => {
let (session_source, agent_metadata) = self.prepare_thread_spawn(
&mut reservation,
&config,
parent_thread_id,
depth,
agent_path,
agent_role,
/*preferred_agent_nickname*/ None,
)?;
(Some(session_source), agent_metadata)
}
other => (other, AgentMetadata::default()),
};
let notification_source = session_source.clone();
// The same `AgentControl` is sent to spawn the thread.
let new_thread = match (session_source, options.fork_mode.as_ref(), inheritance) {
(Some(session_source), Some(_), inheritance) => {
Box::pin(self.spawn_forked_thread(
&state,
config,
session_source,
&options,
inheritance,
multi_agent_version,
))
.await?
}
(Some(session_source), None, inheritance) => {
Box::pin(state.spawn_new_thread_with_source(
config.clone(),
self.clone(),
session_source,
options.parent_thread_id,
/*forked_from_thread_id*/ None,
/*thread_source*/ Some(ThreadSource::Subagent),
/*metrics_service_name*/ None,
inheritance.shell_snapshot,
inheritance.exec_policy,
options.environments.clone(),
))
.await?
}
(None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?,
};
agent_metadata.agent_id = Some(new_thread.thread_id);
reservation.commit(agent_metadata.clone());
if let Some(SessionSource::SubAgent(
subagent_source @ SubAgentSource::ThreadSpawn {
parent_thread_id, ..
},
)) = notification_source.as_ref()
{
let client_metadata = match state.get_thread(*parent_thread_id).await {
Ok(parent_thread) => {
parent_thread
.codex
.session
.app_server_client_metadata()
.await
}
Err(error) => {
tracing::warn!(
error = %error,
parent_thread_id = %parent_thread_id,
"skipping subagent thread analytics: failed to load parent thread metadata"
);
crate::session::session::AppServerClientMetadata {
client_name: None,
client_version: None,
}
}
};
let thread_config = new_thread.thread.codex.thread_config_snapshot().await;
let parent_thread_id = thread_config.parent_thread_id;
emit_subagent_session_started(
&new_thread
.thread
.codex
.session
.services
.analytics_events_client,
client_metadata,
new_thread.thread.codex.session.session_id(),
new_thread.thread_id,
parent_thread_id,
thread_config,
subagent_source.clone(),
);
}
// Notify a new thread has been created. This notification will be processed by clients
// to subscribe or drain this newly created thread.
// TODO(jif) add helper for drain
state.notify_thread_created(new_thread.thread_id);
self.persist_thread_spawn_edge_for_source(
new_thread.thread.as_ref(),
new_thread.thread_id,
notification_source.as_ref(),
)
.await;
self.send_input(new_thread.thread_id, initial_operation)
.await?;
if multi_agent_version != MultiAgentVersion::V2 {
let child_reference = agent_metadata
.agent_path
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| new_thread.thread_id.to_string());
self.maybe_start_completion_watcher(
new_thread.thread_id,
notification_source,
child_reference,
agent_metadata.agent_path.clone(),
);
}
Ok(LiveAgent {
thread_id: new_thread.thread_id,
metadata: agent_metadata,
status: self.get_status(new_thread.thread_id).await,
})
}
async fn spawn_forked_thread(
&self,
state: &Arc<ThreadManagerState>,
config: Config,
session_source: SessionSource,
options: &SpawnAgentOptions,
inheritance: SpawnAgentThreadInheritance,
multi_agent_version: MultiAgentVersion,
) -> CodexResult<crate::thread_manager::NewThread> {
let SpawnAgentThreadInheritance {
shell_snapshot: inherited_shell_snapshot,
exec_policy: inherited_exec_policy,
} = inheritance;
if options.fork_parent_spawn_call_id.is_none() {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a parent spawn call id".to_string(),
));
}
let Some(fork_mode) = options.fork_mode.as_ref() else {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a fork mode".to_string(),
));
};
let SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id, ..
}) = &session_source
else {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a thread-spawn session source".to_string(),
));
};
let parent_thread_id = *parent_thread_id;
let parent_thread = state.get_thread(parent_thread_id).await.ok();
if let Some(parent_thread) = parent_thread.as_ref() {
// `record_conversation_items` only queues persistence writes asynchronously.
// Flush before snapshotting store history for a fork.
parent_thread.ensure_rollout_materialized().await;
parent_thread.flush_rollout().await?;
}
let parent_history = state
.read_stored_thread(ReadThreadParams {
thread_id: parent_thread_id,
include_archived: true,
include_history: true,
})
.await?
.history
.ok_or_else(|| {
CodexErr::Fatal(format!(
"parent thread history unavailable for fork: {parent_thread_id}"
))
})?;
let mut forked_rollout_items = parent_history.items;
if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode {
forked_rollout_items =
truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns);
}
let multi_agent_v2_usage_hint_texts_to_filter: Vec<String> =
if let Some(parent_thread) = parent_thread.as_ref() {
if multi_agent_version == MultiAgentVersion::V2 {
let parent_config = parent_thread.codex.session.get_config().await;
[
parent_config
.multi_agent_v2
.root_agent_usage_hint_text
.clone(),
parent_config
.multi_agent_v2
.subagent_usage_hint_text
.clone(),
]
.into_iter()
.flatten()
.collect()
} else {
Vec::new()
}
} else if multi_agent_version == MultiAgentVersion::V2 {
[
config.multi_agent_v2.root_agent_usage_hint_text.clone(),
config.multi_agent_v2.subagent_usage_hint_text.clone(),
]
.into_iter()
.flatten()
.collect()
} else {
Vec::new()
};
let preserve_reference_context_item = matches!(fork_mode, SpawnAgentForkMode::FullHistory);
forked_rollout_items.retain(|item| {
keep_forked_rollout_item(item, preserve_reference_context_item)
&& !matches!(
item,
RolloutItem::ResponseItem(response_item)
if is_multi_agent_v2_usage_hint_message(
response_item,
&multi_agent_v2_usage_hint_texts_to_filter,
)
)
});
for item in &mut forked_rollout_items {
if let RolloutItem::Compacted(compacted) = item
&& let Some(replacement_history) = compacted.replacement_history.as_mut()
{
replacement_history.retain(|response_item| {
!is_multi_agent_v2_usage_hint_message(
response_item,
&multi_agent_v2_usage_hint_texts_to_filter,
)
});
}
}
if preserve_reference_context_item
&& multi_agent_version == MultiAgentVersion::V2
&& config.multi_agent_v2.usage_hint_enabled
&& let Some(subagent_usage_hint_text) =
config.multi_agent_v2.subagent_usage_hint_text.clone()
&& let Some(subagent_usage_hint_message) =
crate::context_manager::updates::build_developer_update_item(vec![
subagent_usage_hint_text,
])
{
forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message));
}
state
.fork_thread_with_source(
config.clone(),
InitialHistory::Forked(forked_rollout_items),
self.clone(),
session_source,
/*thread_source*/ Some(ThreadSource::Subagent),
/*parent_thread_id*/ Some(parent_thread_id),
/*forked_from_thread_id*/ Some(parent_thread_id),
inherited_shell_snapshot,
inherited_exec_policy,
options.environments.clone(),
)
.await
}
/// Resume an existing agent thread from a recorded rollout file.
pub(crate) async fn resume_agent_from_rollout(
&self,
config: Config,
thread_id: ThreadId,
session_source: SessionSource,
) -> CodexResult<ThreadId> {
let root_depth = thread_spawn_depth(&session_source).unwrap_or(0);
let resumed_thread_id = Box::pin(self.resume_single_agent_from_rollout(
config.clone(),
thread_id,
session_source,
))
.await?;
let state = self.upgrade()?;
let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else {
return Ok(resumed_thread_id);
};
let Some(state_db_ctx) = resumed_thread.state_db() else {
return Ok(resumed_thread_id);
};
let mut resume_queue = VecDeque::from([(thread_id, root_depth)]);
while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() {
let child_ids = match state_db_ctx
.list_thread_spawn_children_with_status(
parent_thread_id,
DirectionalThreadSpawnEdgeStatus::Open,
)
.await
{
Ok(child_ids) => child_ids,
Err(err) => {
warn!(
"failed to load persisted thread-spawn children for {parent_thread_id}: {err}"
);
continue;
}
};
for child_thread_id in child_ids {
let child_depth = parent_depth + 1;
let child_resumed = if state.get_thread(child_thread_id).await.is_ok() {
true
} else {
let child_session_source =
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: child_depth,
agent_path: None,
agent_nickname: None,
agent_role: None,
});
match Box::pin(self.resume_single_agent_from_rollout(
config.clone(),
child_thread_id,
child_session_source,
))
.await
{
Ok(_) => true,
Err(err) => {
warn!("failed to resume descendant thread {child_thread_id}: {err}");
false
}
}
};
if child_resumed {
resume_queue.push_back((child_thread_id, child_depth));
}
}
}
Ok(resumed_thread_id)
}
async fn resume_single_agent_from_rollout(
&self,
config: Config,
thread_id: ThreadId,
session_source: SessionSource,
) -> CodexResult<ThreadId> {
let state = self.upgrade()?;
let state_db_ctx = state.state_db();
let stored_thread = state
.read_stored_thread(ReadThreadParams {
thread_id,
include_archived: true,
include_history: true,
})
.await?;
let history = stored_thread
.history
.ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?
.items;
let initial_history = InitialHistory::Resumed(ResumedHistory {
conversation_id: thread_id,
history,
rollout_path: stored_thread.rollout_path,
});
let parent_thread_id = stored_thread.parent_thread_id;
let multi_agent_version = state
.effective_multi_agent_version_for_spawn(
&initial_history,
Some(&session_source),
parent_thread_id,
/*forked_from_thread_id*/ None,
&config,
)
.await;
let agent_max_threads = config.effective_agent_max_threads(multi_agent_version);
let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?;
let (session_source, agent_metadata) = match session_source {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth,
agent_path,
agent_role: _,
agent_nickname: _,
}) => {
let (resumed_agent_nickname, resumed_agent_role) =
if let Some(state_db_ctx) = state_db_ctx.as_ref() {
match state_db_ctx.get_thread(thread_id).await {
Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role),
Ok(None) | Err(_) => (None, None),
}
} else {
(None, None)
};
self.prepare_thread_spawn(
&mut reservation,
&config,
parent_thread_id,
depth,
agent_path,
resumed_agent_role,
resumed_agent_nickname,
)?
}
other => (other, AgentMetadata::default()),
};
let notification_source = session_source.clone();
let inherited_shell_snapshot = self
.inherited_shell_snapshot_for_source(&state, Some(&session_source))
.await;
let inherited_exec_policy = self
.inherited_exec_policy_for_source(&state, Some(&session_source), &config)
.await;
let resumed_thread = state
.resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions {
config: config.clone(),
initial_history,
agent_control: self.clone(),
session_source,
parent_thread_id,
inherited_shell_snapshot,
inherited_exec_policy,
})
.await?;
let mut agent_metadata = agent_metadata;
agent_metadata.agent_id = Some(resumed_thread.thread_id);
reservation.commit(agent_metadata.clone());
// Resumed threads are re-registered in-memory and need the same listener
// attachment path as freshly spawned threads.
state.notify_thread_created(resumed_thread.thread_id);
if multi_agent_version != MultiAgentVersion::V2 {
let child_reference = agent_metadata
.agent_path
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| resumed_thread.thread_id.to_string());
self.maybe_start_completion_watcher(
resumed_thread.thread_id,
Some(notification_source.clone()),
child_reference,
agent_metadata.agent_path.clone(),
);
}
self.persist_thread_spawn_edge_for_source(
resumed_thread.thread.as_ref(),
resumed_thread.thread_id,
Some(&notification_source),
)
.await;
Ok(resumed_thread.thread_id)
}
/// Send rich user input items to an existing agent thread.
pub(crate) async fn send_input(
&self,
@@ -787,86 +188,6 @@ impl AgentControl {
result
}
/// Submit a shutdown request for a live agent without marking it explicitly closed in
/// persisted spawn-edge state.
pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
let state = self.upgrade()?;
let result = if let Ok(thread) = state.get_thread(agent_id).await {
thread.codex.session.ensure_rollout_materialized().await;
thread.codex.session.flush_rollout().await?;
let result = if matches!(thread.agent_status().await, AgentStatus::Shutdown) {
Ok(String::new())
} else {
state.send_op(agent_id, Op::Shutdown {}).await
};
thread.wait_until_terminated().await;
result
} else {
state.send_op(agent_id, Op::Shutdown {}).await
};
let _ = state.remove_thread(&agent_id).await;
self.state.release_spawned_thread(agent_id);
result
}
/// 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> {
let state = self.upgrade()?;
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,
}
}
/// Shut down `agent_id` and any live descendants reachable from the in-memory spawn tree.
async fn shutdown_agent_tree(&self, agent_id: ThreadId) -> CodexResult<String> {
let descendant_ids = self.live_thread_spawn_descendants(agent_id).await?;
let result = self.shutdown_live_agent(agent_id).await;
for descendant_id in descendant_ids {
match self.shutdown_live_agent(descendant_id).await {
Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => {}
Err(err) => return Err(err),
}
}
result
}
/// Fetch the last known status for `agent_id`, returning `NotFound` when unavailable.
pub(crate) async fn get_status(&self, agent_id: ThreadId) -> AgentStatus {
let Ok(state) = self.upgrade() else {
@@ -1140,7 +461,7 @@ impl AgentControl {
if let Some(agent_path) = agent_path.as_ref() {
reservation.reserve_agent_path(agent_path)?;
}
let candidate_names = agent_nickname_candidates(config, agent_role.as_deref());
let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref());
let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect();
let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference(
&candidate_name_refs,
codex-rs/core/src/agent/control/legacy.rs
+83
View File
@@ -0,0 +1,83 @@
use super::*;
impl AgentControl {
/// Submit a shutdown request for a live agent without marking it explicitly closed in
/// persisted spawn-edge state.
pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
let state = self.upgrade()?;
let result = if let Ok(thread) = state.get_thread(agent_id).await {
thread.codex.session.ensure_rollout_materialized().await;
thread.codex.session.flush_rollout().await?;
let result = if matches!(thread.agent_status().await, AgentStatus::Shutdown) {
Ok(String::new())
} else {
state.send_op(agent_id, Op::Shutdown {}).await
};
thread.wait_until_terminated().await;
result
} else {
state.send_op(agent_id, Op::Shutdown {}).await
};
let _ = state.remove_thread(&agent_id).await;
self.state.release_spawned_thread(agent_id);
result
}
/// 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> {
let state = self.upgrade()?;
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,
}
}
/// Shut down `agent_id` and any live descendants reachable from the in-memory spawn tree.
pub(crate) async fn shutdown_agent_tree(&self, agent_id: ThreadId) -> CodexResult<String> {
let descendant_ids = self.live_thread_spawn_descendants(agent_id).await?;
let result = self.shutdown_live_agent(agent_id).await;
for descendant_id in descendant_ids {
match self.shutdown_live_agent(descendant_id).await {
Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => {}
Err(err) => return Err(err),
}
}
result
}
}
codex-rs/core/src/agent/control/spawn.rs
+606
View File
@@ -0,0 +1,606 @@
use super::*;
const AGENT_NAMES: &str = include_str!("../agent_names.txt");
struct SpawnAgentThreadInheritance {
shell_snapshot: Option<Arc<ShellSnapshot>>,
exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
}
fn default_agent_nickname_list() -> Vec<&'static str> {
AGENT_NAMES
.lines()
.map(str::trim)
.filter(|name| !name.is_empty())
.collect()
}
pub(super) fn agent_nickname_candidates(config: &Config, role_name: Option<&str>) -> Vec<String> {
let role_name = role_name.unwrap_or(DEFAULT_ROLE_NAME);
if let Some(candidates) =
resolve_role_config(config, role_name).and_then(|role| role.nickname_candidates.clone())
{
return candidates;
}
default_agent_nickname_list()
.into_iter()
.map(ToOwned::to_owned)
.collect()
}
fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: bool) -> bool {
match item {
RolloutItem::ResponseItem(ResponseItem::Message { role, phase, .. }) => match role.as_str()
{
"system" | "developer" | "user" => true,
"assistant" => *phase == Some(MessagePhase::FinalAnswer),
_ => false,
},
RolloutItem::ResponseItem(
ResponseItem::AgentMessage { .. }
| ResponseItem::Reasoning { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::ToolSearchCall { .. }
| ResponseItem::FunctionCallOutput { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::CustomToolCallOutput { .. }
| ResponseItem::ToolSearchOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::ImageGenerationCall { .. }
| ResponseItem::Compaction { .. }
| ResponseItem::CompactionTrigger
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other,
) => false,
// Full-history forks preserve the cached prompt prefix and can keep diffing
// from the parent's durable baseline. Truncated forks drop part of that prompt,
// so they must rebuild context on their first child turn.
RolloutItem::TurnContext(_) => preserve_reference_context_item,
RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true,
}
}
fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &[String]) -> bool {
let ResponseItem::Message { role, content, .. } = item else {
return false;
};
if role != "developer" {
return false;
}
let [ContentItem::InputText { text }] = content.as_slice() else {
return false;
};
usage_hint_texts
.iter()
.any(|usage_hint_text| usage_hint_text == text)
}
impl AgentControl {
/// Spawn a new agent thread and submit the initial prompt.
#[cfg(test)]
pub(crate) async fn spawn_agent(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
) -> CodexResult<ThreadId> {
let spawned_agent = Box::pin(self.spawn_agent_internal(
config,
initial_operation,
session_source,
SpawnAgentOptions::default(),
))
.await?;
Ok(spawned_agent.thread_id)
}
/// Spawn an agent thread with some metadata.
pub(crate) async fn spawn_agent_with_metadata(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
options: SpawnAgentOptions, // TODO(jif) drop with new fork.
) -> CodexResult<LiveAgent> {
Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options))
.await
}
async fn spawn_agent_internal(
&self,
config: Config,
initial_operation: Op,
session_source: Option<SessionSource>,
options: SpawnAgentOptions,
) -> CodexResult<LiveAgent> {
let state = self.upgrade()?;
let multi_agent_version = state
.effective_multi_agent_version_for_spawn(
&InitialHistory::New,
session_source.as_ref(),
options.parent_thread_id,
/*forked_from_thread_id*/ None,
&config,
)
.await;
let agent_max_threads = config.effective_agent_max_threads(multi_agent_version);
let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?;
let inheritance = SpawnAgentThreadInheritance {
shell_snapshot: self
.inherited_shell_snapshot_for_source(&state, session_source.as_ref())
.await,
exec_policy: self
.inherited_exec_policy_for_source(&state, session_source.as_ref(), &config)
.await,
};
let (session_source, mut agent_metadata) = match session_source {
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth,
agent_path,
agent_role,
..
})) => {
let (session_source, agent_metadata) = self.prepare_thread_spawn(
&mut reservation,
&config,
parent_thread_id,
depth,
agent_path,
agent_role,
/*preferred_agent_nickname*/ None,
)?;
(Some(session_source), agent_metadata)
}
other => (other, AgentMetadata::default()),
};
let notification_source = session_source.clone();
// The same `AgentControl` is sent to spawn the thread.
let new_thread = match (session_source, options.fork_mode.as_ref(), inheritance) {
(Some(session_source), Some(_), inheritance) => {
Box::pin(self.spawn_forked_thread(
&state,
config,
session_source,
&options,
inheritance,
multi_agent_version,
))
.await?
}
(Some(session_source), None, inheritance) => {
Box::pin(state.spawn_new_thread_with_source(
config.clone(),
self.clone(),
session_source,
options.parent_thread_id,
/*forked_from_thread_id*/ None,
/*thread_source*/ Some(ThreadSource::Subagent),
/*metrics_service_name*/ None,
inheritance.shell_snapshot,
inheritance.exec_policy,
options.environments.clone(),
))
.await?
}
(None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?,
};
agent_metadata.agent_id = Some(new_thread.thread_id);
reservation.commit(agent_metadata.clone());
if let Some(SessionSource::SubAgent(
subagent_source @ SubAgentSource::ThreadSpawn {
parent_thread_id, ..
},
)) = notification_source.as_ref()
{
let client_metadata = match state.get_thread(*parent_thread_id).await {
Ok(parent_thread) => {
parent_thread
.codex
.session
.app_server_client_metadata()
.await
}
Err(error) => {
tracing::warn!(
error = %error,
parent_thread_id = %parent_thread_id,
"skipping subagent thread analytics: failed to load parent thread metadata"
);
crate::session::session::AppServerClientMetadata {
client_name: None,
client_version: None,
}
}
};
let thread_config = new_thread.thread.codex.thread_config_snapshot().await;
let parent_thread_id = thread_config.parent_thread_id;
emit_subagent_session_started(
&new_thread
.thread
.codex
.session
.services
.analytics_events_client,
client_metadata,
new_thread.thread.codex.session.session_id(),
new_thread.thread_id,
parent_thread_id,
thread_config,
subagent_source.clone(),
);
}
// Notify a new thread has been created. This notification will be processed by clients
// to subscribe or drain this newly created thread.
// TODO(jif) add helper for drain
state.notify_thread_created(new_thread.thread_id);
self.persist_thread_spawn_edge_for_source(
new_thread.thread.as_ref(),
new_thread.thread_id,
notification_source.as_ref(),
)
.await;
self.send_input(new_thread.thread_id, initial_operation)
.await?;
if multi_agent_version != MultiAgentVersion::V2 {
let child_reference = agent_metadata
.agent_path
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| new_thread.thread_id.to_string());
self.maybe_start_completion_watcher(
new_thread.thread_id,
notification_source,
child_reference,
agent_metadata.agent_path.clone(),
);
}
Ok(LiveAgent {
thread_id: new_thread.thread_id,
metadata: agent_metadata,
status: self.get_status(new_thread.thread_id).await,
})
}
async fn spawn_forked_thread(
&self,
state: &Arc<ThreadManagerState>,
config: Config,
session_source: SessionSource,
options: &SpawnAgentOptions,
inheritance: SpawnAgentThreadInheritance,
multi_agent_version: MultiAgentVersion,
) -> CodexResult<crate::thread_manager::NewThread> {
let SpawnAgentThreadInheritance {
shell_snapshot: inherited_shell_snapshot,
exec_policy: inherited_exec_policy,
} = inheritance;
if options.fork_parent_spawn_call_id.is_none() {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a parent spawn call id".to_string(),
));
}
let Some(fork_mode) = options.fork_mode.as_ref() else {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a fork mode".to_string(),
));
};
let SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id, ..
}) = &session_source
else {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a thread-spawn session source".to_string(),
));
};
let parent_thread_id = *parent_thread_id;
let parent_thread = state.get_thread(parent_thread_id).await.ok();
if let Some(parent_thread) = parent_thread.as_ref() {
// `record_conversation_items` only queues persistence writes asynchronously.
// Flush before snapshotting store history for a fork.
parent_thread.ensure_rollout_materialized().await;
parent_thread.flush_rollout().await?;
}
let parent_history = state
.read_stored_thread(ReadThreadParams {
thread_id: parent_thread_id,
include_archived: true,
include_history: true,
})
.await?
.history
.ok_or_else(|| {
CodexErr::Fatal(format!(
"parent thread history unavailable for fork: {parent_thread_id}"
))
})?;
let mut forked_rollout_items = parent_history.items;
if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode {
forked_rollout_items =
truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns);
}
let multi_agent_v2_usage_hint_texts_to_filter: Vec<String> =
if let Some(parent_thread) = parent_thread.as_ref() {
if multi_agent_version == MultiAgentVersion::V2 {
let parent_config = parent_thread.codex.session.get_config().await;
[
parent_config
.multi_agent_v2
.root_agent_usage_hint_text
.clone(),
parent_config
.multi_agent_v2
.subagent_usage_hint_text
.clone(),
]
.into_iter()
.flatten()
.collect()
} else {
Vec::new()
}
} else if multi_agent_version == MultiAgentVersion::V2 {
[
config.multi_agent_v2.root_agent_usage_hint_text.clone(),
config.multi_agent_v2.subagent_usage_hint_text.clone(),
]
.into_iter()
.flatten()
.collect()
} else {
Vec::new()
};
let preserve_reference_context_item = matches!(fork_mode, SpawnAgentForkMode::FullHistory);
forked_rollout_items.retain(|item| {
keep_forked_rollout_item(item, preserve_reference_context_item)
&& !matches!(
item,
RolloutItem::ResponseItem(response_item)
if is_multi_agent_v2_usage_hint_message(
response_item,
&multi_agent_v2_usage_hint_texts_to_filter,
)
)
});
for item in &mut forked_rollout_items {
if let RolloutItem::Compacted(compacted) = item
&& let Some(replacement_history) = compacted.replacement_history.as_mut()
{
replacement_history.retain(|response_item| {
!is_multi_agent_v2_usage_hint_message(
response_item,
&multi_agent_v2_usage_hint_texts_to_filter,
)
});
}
}
if preserve_reference_context_item
&& multi_agent_version == MultiAgentVersion::V2
&& config.multi_agent_v2.usage_hint_enabled
&& let Some(subagent_usage_hint_text) =
config.multi_agent_v2.subagent_usage_hint_text.clone()
&& let Some(subagent_usage_hint_message) =
crate::context_manager::updates::build_developer_update_item(vec![
subagent_usage_hint_text,
])
{
forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message));
}
state
.fork_thread_with_source(
config.clone(),
InitialHistory::Forked(forked_rollout_items),
self.clone(),
session_source,
/*thread_source*/ Some(ThreadSource::Subagent),
/*parent_thread_id*/ Some(parent_thread_id),
/*forked_from_thread_id*/ Some(parent_thread_id),
inherited_shell_snapshot,
inherited_exec_policy,
options.environments.clone(),
)
.await
}
/// Resume an existing agent thread from a recorded rollout file.
pub(crate) async fn resume_agent_from_rollout(
&self,
config: Config,
thread_id: ThreadId,
session_source: SessionSource,
) -> CodexResult<ThreadId> {
let root_depth = thread_spawn_depth(&session_source).unwrap_or(0);
let resumed_thread_id = Box::pin(self.resume_single_agent_from_rollout(
config.clone(),
thread_id,
session_source,
))
.await?;
let state = self.upgrade()?;
let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else {
return Ok(resumed_thread_id);
};
let Some(state_db_ctx) = resumed_thread.state_db() else {
return Ok(resumed_thread_id);
};
let mut resume_queue = VecDeque::from([(thread_id, root_depth)]);
while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() {
let child_ids = match state_db_ctx
.list_thread_spawn_children_with_status(
parent_thread_id,
DirectionalThreadSpawnEdgeStatus::Open,
)
.await
{
Ok(child_ids) => child_ids,
Err(err) => {
warn!(
"failed to load persisted thread-spawn children for {parent_thread_id}: {err}"
);
continue;
}
};
for child_thread_id in child_ids {
let child_depth = parent_depth + 1;
let child_resumed = if state.get_thread(child_thread_id).await.is_ok() {
true
} else {
let child_session_source =
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: child_depth,
agent_path: None,
agent_nickname: None,
agent_role: None,
});
match Box::pin(self.resume_single_agent_from_rollout(
config.clone(),
child_thread_id,
child_session_source,
))
.await
{
Ok(_) => true,
Err(err) => {
warn!("failed to resume descendant thread {child_thread_id}: {err}");
false
}
}
};
if child_resumed {
resume_queue.push_back((child_thread_id, child_depth));
}
}
}
Ok(resumed_thread_id)
}
async fn resume_single_agent_from_rollout(
&self,
config: Config,
thread_id: ThreadId,
session_source: SessionSource,
) -> CodexResult<ThreadId> {
let state = self.upgrade()?;
let state_db_ctx = state.state_db();
let stored_thread = state
.read_stored_thread(ReadThreadParams {
thread_id,
include_archived: true,
include_history: true,
})
.await?;
let history = stored_thread
.history
.ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?
.items;
let initial_history = InitialHistory::Resumed(ResumedHistory {
conversation_id: thread_id,
history,
rollout_path: stored_thread.rollout_path,
});
let parent_thread_id = stored_thread.parent_thread_id;
let multi_agent_version = state
.effective_multi_agent_version_for_spawn(
&initial_history,
Some(&session_source),
parent_thread_id,
/*forked_from_thread_id*/ None,
&config,
)
.await;
let agent_max_threads = config.effective_agent_max_threads(multi_agent_version);
let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?;
let (session_source, agent_metadata) = match session_source {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth,
agent_path,
agent_role: _,
agent_nickname: _,
}) => {
let (resumed_agent_nickname, resumed_agent_role) =
if let Some(state_db_ctx) = state_db_ctx.as_ref() {
match state_db_ctx.get_thread(thread_id).await {
Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role),
Ok(None) | Err(_) => (None, None),
}
} else {
(None, None)
};
self.prepare_thread_spawn(
&mut reservation,
&config,
parent_thread_id,
depth,
agent_path,
resumed_agent_role,
resumed_agent_nickname,
)?
}
other => (other, AgentMetadata::default()),
};
let notification_source = session_source.clone();
let inherited_shell_snapshot = self
.inherited_shell_snapshot_for_source(&state, Some(&session_source))
.await;
let inherited_exec_policy = self
.inherited_exec_policy_for_source(&state, Some(&session_source), &config)
.await;
let resumed_thread = state
.resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions {
config: config.clone(),
initial_history,
agent_control: self.clone(),
session_source,
parent_thread_id,
inherited_shell_snapshot,
inherited_exec_policy,
})
.await?;
let mut agent_metadata = agent_metadata;
agent_metadata.agent_id = Some(resumed_thread.thread_id);
reservation.commit(agent_metadata.clone());
// Resumed threads are re-registered in-memory and need the same listener
// attachment path as freshly spawned threads.
state.notify_thread_created(resumed_thread.thread_id);
if multi_agent_version != MultiAgentVersion::V2 {
let child_reference = agent_metadata
.agent_path
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| resumed_thread.thread_id.to_string());
self.maybe_start_completion_watcher(
resumed_thread.thread_id,
Some(notification_source.clone()),
child_reference,
agent_metadata.agent_path.clone(),
);
}
self.persist_thread_spawn_edge_for_source(
resumed_thread.thread.as_ref(),
resumed_thread.thread_id,
Some(&notification_source),
)
.await;
Ok(resumed_thread.thread_id)
}
}