mirror of
https://github.com/openai/codex.git
synced 2026-11-03 19:53:25 +00:00
rka-oai
32a696dbac
## Stack Depends on #28746. This PR implements shared rollout-budget accounting and model-visible reminders using the configuration defined in #28746. # Description / Main changes to Core: `AgentControl` will now be the area where "rollout level" features & accounting will have to live. It is incorrectly named for this responsibility, but I think it can hold all the necessary shared state & features (rollout token budget, mutliple thread interruption responsibilitym etc) In this PR, we have one "token ledger" that each thread will subtract from when sampling. The "charge" will occur when response.completed() is done and the calculation will be done on the responses api usage carrier. The calculation will weigh sampling and pre-fill tokens as specified. Every time the budget crosses the configured reminder threshold, a developer message is appended before the thread's next request This remaining budget will _always_ be restated/reminded after a compaction event. Expiration and fan-out interruption will be in the stacked follow-up (and also live in Agent Control). ## Reminders "You have weighted {session_tokens_left} tokens left in the shared session token budget." The first request in each thread context receives the current remainder. Later reminders are emitted after aggregate weighted usage crosses a configured interval. If several intervals are crossed before a thread sends another request, Core inserts one reminder with the latest remainder. Compaction response usage is charged before the next context starts. The next reminder is appended after the compaction summary, leaving the initial context content stable. ## Tests Integration coverage verifies: - weighted output and non-cached input accounting - initial and periodic reminders - shared accounting between a root and sub-agent - post-compaction remainder and message placement Local checks: - `just fmt` - `just test -p codex-core rollout_budget` - `git diff --check` The full workspace test suite was not run locally.
747 lines
26 KiB
Rust
747 lines
26 KiB
Rust
use crate::agent::AgentStatus;
|
|
use crate::agent::registry::AgentMetadata;
|
|
use crate::agent::registry::AgentRegistry;
|
|
use crate::agent::role::DEFAULT_ROLE_NAME;
|
|
use crate::agent::role::resolve_role_config;
|
|
use crate::agent::status::is_final;
|
|
use crate::codex_thread::ThreadConfigSnapshot;
|
|
use crate::config::Config;
|
|
use crate::config::RolloutBudgetConfig;
|
|
use crate::environment_selection::TurnEnvironmentSnapshot;
|
|
use crate::rollout_budget::RolloutBudget;
|
|
use crate::session::emit_subagent_session_started;
|
|
use crate::session_prefix::format_inter_agent_completion_message;
|
|
use crate::session_prefix::format_subagent_context_line;
|
|
use crate::session_prefix::format_subagent_notification_message;
|
|
use crate::thread_manager::ResumeThreadWithHistoryOptions;
|
|
use crate::thread_manager::ThreadManagerState;
|
|
use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns;
|
|
use codex_protocol::AgentPath;
|
|
use codex_protocol::SessionId;
|
|
use codex_protocol::ThreadId;
|
|
use codex_protocol::error::CodexErr;
|
|
use codex_protocol::error::Result as CodexResult;
|
|
use codex_protocol::models::ContentItem;
|
|
use codex_protocol::models::MessagePhase;
|
|
use codex_protocol::models::ResponseItem;
|
|
use codex_protocol::protocol::InitialHistory;
|
|
use codex_protocol::protocol::InterAgentCommunication;
|
|
use codex_protocol::protocol::MultiAgentVersion;
|
|
use codex_protocol::protocol::Op;
|
|
use codex_protocol::protocol::ResumedHistory;
|
|
use codex_protocol::protocol::RolloutItem;
|
|
use codex_protocol::protocol::SessionSource;
|
|
use codex_protocol::protocol::SubAgentSource;
|
|
use codex_protocol::protocol::ThreadSource;
|
|
use codex_protocol::protocol::TurnEnvironmentSelection;
|
|
use codex_protocol::user_input::UserInput;
|
|
use codex_state::DirectionalThreadSpawnEdgeStatus;
|
|
use codex_thread_store::ReadThreadParams;
|
|
use serde::Serialize;
|
|
use std::collections::HashMap;
|
|
use std::collections::VecDeque;
|
|
use std::sync::Arc;
|
|
use std::sync::Weak;
|
|
use tokio::sync::watch;
|
|
use tracing::warn;
|
|
|
|
pub(crate) use self::execution::AgentExecutionGuard;
|
|
use self::execution::AgentExecutionLimiter;
|
|
use self::residency::V2Residency;
|
|
|
|
const ROOT_LAST_TASK_MESSAGE: &str = "Main thread";
|
|
|
|
mod execution;
|
|
mod legacy;
|
|
mod residency;
|
|
mod spawn;
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub(crate) enum SpawnAgentForkMode {
|
|
FullHistory,
|
|
LastNTurns(usize),
|
|
}
|
|
|
|
#[derive(Clone, Debug, Default)]
|
|
pub(crate) struct SpawnAgentOptions {
|
|
pub(crate) fork_parent_spawn_call_id: Option<String>,
|
|
pub(crate) fork_mode: Option<SpawnAgentForkMode>,
|
|
pub(crate) parent_thread_id: Option<ThreadId>,
|
|
pub(crate) environments: Option<Vec<TurnEnvironmentSelection>>,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(crate) struct LiveAgent {
|
|
pub(crate) thread_id: ThreadId,
|
|
pub(crate) metadata: AgentMetadata,
|
|
pub(crate) status: AgentStatus,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
|
|
pub(crate) struct ListedAgent {
|
|
pub(crate) agent_name: String,
|
|
pub(crate) agent_status: AgentStatus,
|
|
pub(crate) last_task_message: Option<String>,
|
|
}
|
|
|
|
/// 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.
|
|
/// An `AgentControl` instance is intended to be created at most once per root thread/session
|
|
/// tree. That same `AgentControl` is then shared with every sub-agent spawned from that root,
|
|
/// which keeps the registry scoped to that root thread rather than the entire `ThreadManager`.
|
|
#[derive(Clone, Default)]
|
|
pub(crate) struct AgentControl {
|
|
/// ID shared by the whole agent control session. This means every sub-agents from a common
|
|
/// root share the same session ID.
|
|
session_id: SessionId,
|
|
/// Weak handle back to the global thread registry/state.
|
|
/// This is `Weak` to avoid reference cycles and shadow persistence of the form
|
|
/// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`.
|
|
manager: Weak<ThreadManagerState>,
|
|
state: Arc<AgentRegistry>,
|
|
v2_residency: Arc<V2Residency>,
|
|
agent_execution_limiter: Arc<AgentExecutionLimiter>,
|
|
/// Session-scoped state shared by the root thread and every cloned sub-agent control handle.
|
|
rollout_budget: Arc<RolloutBudget>,
|
|
}
|
|
|
|
impl AgentControl {
|
|
/// Construct a new `AgentControl` that can spawn/message agents via the given manager state.
|
|
pub(crate) fn new(
|
|
manager: Weak<ThreadManagerState>,
|
|
rollout_budget: Option<RolloutBudgetConfig>,
|
|
) -> Self {
|
|
let control = Self {
|
|
manager,
|
|
..Default::default()
|
|
};
|
|
if let Some(rollout_budget) = rollout_budget {
|
|
control.rollout_budget.configure(rollout_budget);
|
|
}
|
|
control
|
|
}
|
|
|
|
pub(crate) fn with_session_id(mut self, session_id: SessionId, max_threads: usize) -> Self {
|
|
self.session_id = session_id;
|
|
self.agent_execution_limiter.initialize(max_threads);
|
|
self
|
|
}
|
|
|
|
pub(crate) fn session_id(&self) -> SessionId {
|
|
self.session_id
|
|
}
|
|
|
|
pub(crate) fn rollout_budget(&self) -> &RolloutBudget {
|
|
self.rollout_budget.as_ref()
|
|
}
|
|
|
|
/// Send rich user input items to an existing agent thread.
|
|
pub(crate) async fn send_input(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
initial_operation: Op,
|
|
) -> CodexResult<String> {
|
|
let state = self.upgrade()?;
|
|
self.ensure_execution_capacity_for_op(agent_id, &initial_operation)
|
|
.await?;
|
|
self.send_input_after_capacity_check(agent_id, &state, initial_operation)
|
|
.await
|
|
}
|
|
|
|
async fn send_input_after_capacity_check(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
state: &Arc<ThreadManagerState>,
|
|
initial_operation: Op,
|
|
) -> CodexResult<String> {
|
|
let last_task_message = match &initial_operation {
|
|
Op::InterAgentCommunication { communication } => {
|
|
last_task_message_from_communication(communication)
|
|
}
|
|
_ => non_empty_task_message(render_input_preview(&initial_operation)),
|
|
};
|
|
let result = self
|
|
.handle_thread_request_result(
|
|
agent_id,
|
|
state,
|
|
state.send_op(agent_id, initial_operation).await,
|
|
)
|
|
.await;
|
|
if result.is_ok() {
|
|
match last_task_message {
|
|
Some(last_task_message) => self
|
|
.state
|
|
.update_last_task_message(agent_id, last_task_message),
|
|
None => self.state.clear_last_task_message(agent_id),
|
|
}
|
|
}
|
|
result
|
|
}
|
|
|
|
pub(crate) async fn send_inter_agent_communication(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
communication: InterAgentCommunication,
|
|
) -> CodexResult<String> {
|
|
let last_task_message = last_task_message_from_communication(&communication);
|
|
let state = self.upgrade()?;
|
|
let op = Op::InterAgentCommunication { communication };
|
|
self.ensure_execution_capacity_for_op(agent_id, &op).await?;
|
|
let result = self
|
|
.handle_thread_request_result(agent_id, &state, state.send_op(agent_id, op).await)
|
|
.await;
|
|
if result.is_ok() {
|
|
match last_task_message {
|
|
Some(last_task_message) => self
|
|
.state
|
|
.update_last_task_message(agent_id, last_task_message),
|
|
None => self.state.clear_last_task_message(agent_id),
|
|
}
|
|
}
|
|
result
|
|
}
|
|
|
|
/// Interrupt the current task for an existing agent thread.
|
|
pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
|
|
let state = self.upgrade()?;
|
|
self.handle_thread_request_result(
|
|
agent_id,
|
|
&state,
|
|
state.send_op(agent_id, Op::Interrupt).await,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn handle_thread_request_result(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
state: &Arc<ThreadManagerState>,
|
|
result: CodexResult<String>,
|
|
) -> CodexResult<String> {
|
|
if matches!(result, Err(CodexErr::InternalAgentDied)) {
|
|
let _ = state.remove_thread(&agent_id).await;
|
|
self.forget_v2_residency(agent_id);
|
|
self.state.release_spawned_thread(agent_id);
|
|
}
|
|
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 {
|
|
// No agent available if upgrade fails.
|
|
return AgentStatus::NotFound;
|
|
};
|
|
let Ok(thread) = state.get_thread(agent_id).await else {
|
|
return AgentStatus::NotFound;
|
|
};
|
|
thread.agent_status().await
|
|
}
|
|
|
|
pub(crate) fn register_session_root(
|
|
&self,
|
|
current_thread_id: ThreadId,
|
|
current_parent_thread_id: Option<ThreadId>,
|
|
) {
|
|
if current_parent_thread_id.is_none() {
|
|
self.state.register_root_thread(current_thread_id);
|
|
}
|
|
}
|
|
|
|
pub(crate) fn get_agent_metadata(&self, agent_id: ThreadId) -> Option<AgentMetadata> {
|
|
self.state.agent_metadata_for_thread(agent_id)
|
|
}
|
|
|
|
pub(crate) fn ensure_agent_known(&self, agent_id: ThreadId) -> CodexResult<AgentMetadata> {
|
|
self.state
|
|
.agent_metadata_for_thread(agent_id)
|
|
.ok_or(CodexErr::ThreadNotFound(agent_id))
|
|
}
|
|
|
|
pub(crate) async fn list_live_agent_subtree_thread_ids(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
) -> CodexResult<Vec<ThreadId>> {
|
|
let mut thread_ids = vec![agent_id];
|
|
thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?);
|
|
Ok(thread_ids)
|
|
}
|
|
|
|
pub(crate) async fn get_agent_config_snapshot(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
) -> Option<ThreadConfigSnapshot> {
|
|
let Ok(state) = self.upgrade() else {
|
|
return None;
|
|
};
|
|
let Ok(thread) = state.get_thread(agent_id).await else {
|
|
return None;
|
|
};
|
|
Some(thread.config_snapshot().await)
|
|
}
|
|
|
|
pub(crate) async fn resolve_agent_reference(
|
|
&self,
|
|
_current_thread_id: ThreadId,
|
|
current_session_source: &SessionSource,
|
|
agent_reference: &str,
|
|
) -> CodexResult<ThreadId> {
|
|
let current_agent_path = current_session_source
|
|
.get_agent_path()
|
|
.unwrap_or_else(AgentPath::root);
|
|
let agent_path = current_agent_path
|
|
.resolve(agent_reference)
|
|
.map_err(CodexErr::UnsupportedOperation)?;
|
|
if let Some(thread_id) = self.state.agent_id_for_path(&agent_path) {
|
|
return Ok(thread_id);
|
|
}
|
|
Err(CodexErr::UnsupportedOperation(format!(
|
|
"live agent path `{}` not found",
|
|
agent_path.as_str()
|
|
)))
|
|
}
|
|
|
|
/// Subscribe to status updates for `agent_id`, yielding the latest value and changes.
|
|
pub(crate) async fn subscribe_status(
|
|
&self,
|
|
agent_id: ThreadId,
|
|
) -> CodexResult<watch::Receiver<AgentStatus>> {
|
|
let state = self.upgrade()?;
|
|
let thread = state.get_thread(agent_id).await?;
|
|
Ok(thread.subscribe_status())
|
|
}
|
|
|
|
pub(crate) async fn format_environment_context_subagents(
|
|
&self,
|
|
parent_thread_id: ThreadId,
|
|
) -> String {
|
|
let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else {
|
|
return String::new();
|
|
};
|
|
|
|
agents
|
|
.into_iter()
|
|
.map(|(thread_id, metadata)| {
|
|
let reference = metadata
|
|
.agent_path
|
|
.as_ref()
|
|
.map(|agent_path| agent_path.name().to_string())
|
|
.unwrap_or_else(|| thread_id.to_string());
|
|
format_subagent_context_line(reference.as_str(), metadata.agent_nickname.as_deref())
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n")
|
|
}
|
|
|
|
pub(crate) async fn list_agents(
|
|
&self,
|
|
current_session_source: &SessionSource,
|
|
path_prefix: Option<&str>,
|
|
) -> CodexResult<Vec<ListedAgent>> {
|
|
let state = self.upgrade()?;
|
|
let resolved_prefix = path_prefix
|
|
.map(|prefix| {
|
|
current_session_source
|
|
.get_agent_path()
|
|
.unwrap_or_else(AgentPath::root)
|
|
.resolve(prefix)
|
|
.map_err(CodexErr::UnsupportedOperation)
|
|
})
|
|
.transpose()?;
|
|
|
|
let mut live_agents = self.state.live_agents();
|
|
live_agents.sort_by(|left, right| {
|
|
left.agent_path
|
|
.as_deref()
|
|
.unwrap_or_default()
|
|
.cmp(right.agent_path.as_deref().unwrap_or_default())
|
|
.then_with(|| {
|
|
left.agent_id
|
|
.map(|id| id.to_string())
|
|
.unwrap_or_default()
|
|
.cmp(&right.agent_id.map(|id| id.to_string()).unwrap_or_default())
|
|
})
|
|
});
|
|
|
|
let root_path = AgentPath::root();
|
|
let mut agents = Vec::with_capacity(live_agents.len().saturating_add(1));
|
|
if resolved_prefix
|
|
.as_ref()
|
|
.is_none_or(|prefix| agent_matches_prefix(Some(&root_path), prefix))
|
|
&& let Some(root_thread_id) = self.state.agent_id_for_path(&root_path)
|
|
&& let Ok(root_thread) = state.get_thread(root_thread_id).await
|
|
{
|
|
agents.push(ListedAgent {
|
|
agent_name: root_path.to_string(),
|
|
agent_status: root_thread.agent_status().await,
|
|
last_task_message: Some(ROOT_LAST_TASK_MESSAGE.to_string()),
|
|
});
|
|
}
|
|
|
|
for metadata in live_agents {
|
|
let Some(thread_id) = metadata.agent_id else {
|
|
continue;
|
|
};
|
|
if resolved_prefix
|
|
.as_ref()
|
|
.is_some_and(|prefix| !agent_matches_prefix(metadata.agent_path.as_ref(), prefix))
|
|
{
|
|
continue;
|
|
}
|
|
|
|
let Ok(thread) = state.get_thread(thread_id).await else {
|
|
continue;
|
|
};
|
|
let agent_name = metadata
|
|
.agent_path
|
|
.as_ref()
|
|
.map(ToString::to_string)
|
|
.unwrap_or_else(|| thread_id.to_string());
|
|
let last_task_message = metadata.last_task_message.clone();
|
|
agents.push(ListedAgent {
|
|
agent_name,
|
|
agent_status: thread.agent_status().await,
|
|
last_task_message,
|
|
});
|
|
}
|
|
|
|
Ok(agents)
|
|
}
|
|
|
|
/// Starts a detached watcher for sub-agents spawned from another thread.
|
|
///
|
|
/// This is only enabled for `SubAgentSource::ThreadSpawn`, where a parent thread exists and
|
|
/// can receive completion notifications.
|
|
fn maybe_start_completion_watcher(
|
|
&self,
|
|
child_thread_id: ThreadId,
|
|
session_source: Option<SessionSource>,
|
|
child_reference: String,
|
|
child_agent_path: Option<AgentPath>,
|
|
) {
|
|
let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
|
parent_thread_id, ..
|
|
})) = session_source
|
|
else {
|
|
return;
|
|
};
|
|
let control = self.clone();
|
|
tokio::spawn(async move {
|
|
let status = match control.subscribe_status(child_thread_id).await {
|
|
Ok(mut status_rx) => {
|
|
let mut status = status_rx.borrow().clone();
|
|
while !is_final(&status) {
|
|
if status_rx.changed().await.is_err() {
|
|
status = control.get_status(child_thread_id).await;
|
|
break;
|
|
}
|
|
status = status_rx.borrow().clone();
|
|
}
|
|
status
|
|
}
|
|
Err(_) => control.get_status(child_thread_id).await,
|
|
};
|
|
if !is_final(&status) {
|
|
return;
|
|
}
|
|
|
|
let Ok(state) = control.upgrade() else {
|
|
return;
|
|
};
|
|
let child_thread = state.get_thread(child_thread_id).await.ok();
|
|
let child_uses_multi_agent_v2 = match child_thread.as_ref() {
|
|
Some(child_thread) => {
|
|
child_thread.multi_agent_version() == Some(MultiAgentVersion::V2)
|
|
}
|
|
None => true,
|
|
};
|
|
if child_agent_path.is_some() && child_uses_multi_agent_v2 {
|
|
let Some(child_agent_path) = child_agent_path.clone() else {
|
|
return;
|
|
};
|
|
let Some(parent_agent_path) = child_agent_path
|
|
.as_str()
|
|
.rsplit_once('/')
|
|
.and_then(|(parent, _)| AgentPath::try_from(parent).ok())
|
|
else {
|
|
return;
|
|
};
|
|
let Some(message) = format_inter_agent_completion_message(
|
|
parent_agent_path.clone(),
|
|
child_agent_path.clone(),
|
|
&status,
|
|
) else {
|
|
return;
|
|
};
|
|
let communication = InterAgentCommunication::new(
|
|
child_agent_path,
|
|
parent_agent_path,
|
|
Vec::new(),
|
|
message,
|
|
/*trigger_turn*/ false,
|
|
);
|
|
let _ = control
|
|
.send_inter_agent_communication(parent_thread_id, communication)
|
|
.await;
|
|
return;
|
|
}
|
|
let message = format_subagent_notification_message(child_reference.as_str(), &status);
|
|
let Ok(parent_thread) = state.get_thread(parent_thread_id).await else {
|
|
return;
|
|
};
|
|
parent_thread
|
|
.inject_user_message_without_turn(message)
|
|
.await;
|
|
});
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn prepare_thread_spawn(
|
|
&self,
|
|
reservation: &mut crate::agent::registry::SpawnReservation,
|
|
config: &Config,
|
|
parent_thread_id: ThreadId,
|
|
depth: i32,
|
|
agent_path: Option<AgentPath>,
|
|
agent_role: Option<String>,
|
|
preferred_agent_nickname: Option<String>,
|
|
) -> CodexResult<(SessionSource, AgentMetadata)> {
|
|
if depth == 1 {
|
|
self.state.register_root_thread(parent_thread_id);
|
|
}
|
|
if let Some(agent_path) = agent_path.as_ref() {
|
|
reservation.reserve_agent_path(agent_path)?;
|
|
}
|
|
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,
|
|
preferred_agent_nickname.as_deref(),
|
|
)?);
|
|
let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
|
parent_thread_id,
|
|
depth,
|
|
agent_path: agent_path.clone(),
|
|
agent_nickname: agent_nickname.clone(),
|
|
agent_role: agent_role.clone(),
|
|
});
|
|
let agent_metadata = AgentMetadata {
|
|
agent_id: None,
|
|
agent_path,
|
|
agent_nickname,
|
|
agent_role,
|
|
last_task_message: None,
|
|
};
|
|
Ok((session_source, agent_metadata))
|
|
}
|
|
|
|
fn upgrade(&self) -> CodexResult<Arc<ThreadManagerState>> {
|
|
self.manager
|
|
.upgrade()
|
|
.ok_or_else(|| CodexErr::UnsupportedOperation("thread manager dropped".to_string()))
|
|
}
|
|
|
|
async fn inherited_environments_for_source(
|
|
&self,
|
|
state: &Arc<ThreadManagerState>,
|
|
session_source: Option<&SessionSource>,
|
|
) -> Option<TurnEnvironmentSnapshot> {
|
|
let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
|
parent_thread_id, ..
|
|
})) = session_source
|
|
else {
|
|
return None;
|
|
};
|
|
|
|
let parent_thread = state.get_thread(*parent_thread_id).await.ok()?;
|
|
Some(
|
|
parent_thread
|
|
.codex
|
|
.session
|
|
.services
|
|
.turn_environments
|
|
.snapshot()
|
|
.await,
|
|
)
|
|
}
|
|
|
|
async fn inherited_exec_policy_for_source(
|
|
&self,
|
|
state: &Arc<ThreadManagerState>,
|
|
session_source: Option<&SessionSource>,
|
|
child_config: &Config,
|
|
) -> Option<Arc<crate::exec_policy::ExecPolicyManager>> {
|
|
let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
|
parent_thread_id, ..
|
|
})) = session_source
|
|
else {
|
|
return None;
|
|
};
|
|
|
|
let parent_thread = state.get_thread(*parent_thread_id).await.ok()?;
|
|
let parent_config = parent_thread.codex.session.get_config().await;
|
|
if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) {
|
|
return None;
|
|
}
|
|
|
|
Some(Arc::clone(
|
|
&parent_thread.codex.session.services.exec_policy,
|
|
))
|
|
}
|
|
|
|
async fn open_thread_spawn_children(
|
|
&self,
|
|
parent_thread_id: ThreadId,
|
|
) -> CodexResult<Vec<(ThreadId, AgentMetadata)>> {
|
|
let mut children_by_parent = self.live_thread_spawn_children().await?;
|
|
Ok(children_by_parent
|
|
.remove(&parent_thread_id)
|
|
.unwrap_or_default())
|
|
}
|
|
|
|
async fn live_thread_spawn_children(
|
|
&self,
|
|
) -> CodexResult<HashMap<ThreadId, Vec<(ThreadId, AgentMetadata)>>> {
|
|
let state = self.upgrade()?;
|
|
let mut children_by_parent = HashMap::<ThreadId, Vec<(ThreadId, AgentMetadata)>>::new();
|
|
|
|
for (parent_thread_id, child_thread_id) in state.list_live_thread_spawn_edges().await {
|
|
children_by_parent
|
|
.entry(parent_thread_id)
|
|
.or_default()
|
|
.push((
|
|
child_thread_id,
|
|
self.state
|
|
.agent_metadata_for_thread(child_thread_id)
|
|
.unwrap_or(AgentMetadata {
|
|
agent_id: Some(child_thread_id),
|
|
..Default::default()
|
|
}),
|
|
));
|
|
}
|
|
|
|
for children in children_by_parent.values_mut() {
|
|
children.sort_by(|left, right| {
|
|
left.1
|
|
.agent_path
|
|
.as_deref()
|
|
.unwrap_or_default()
|
|
.cmp(right.1.agent_path.as_deref().unwrap_or_default())
|
|
.then_with(|| left.0.to_string().cmp(&right.0.to_string()))
|
|
});
|
|
}
|
|
|
|
Ok(children_by_parent)
|
|
}
|
|
|
|
async fn persist_thread_spawn_edge_for_source(
|
|
&self,
|
|
thread: &crate::CodexThread,
|
|
child_thread_id: ThreadId,
|
|
session_source: Option<&SessionSource>,
|
|
) {
|
|
let Some(parent_thread_id) = session_source.and_then(SessionSource::parent_thread_id)
|
|
else {
|
|
return;
|
|
};
|
|
let Some(state_db_ctx) = thread.state_db() else {
|
|
return;
|
|
};
|
|
if let Err(err) = state_db_ctx
|
|
.upsert_thread_spawn_edge(
|
|
parent_thread_id,
|
|
child_thread_id,
|
|
DirectionalThreadSpawnEdgeStatus::Open,
|
|
)
|
|
.await
|
|
{
|
|
warn!("failed to persist thread-spawn edge: {err}");
|
|
}
|
|
}
|
|
|
|
async fn live_thread_spawn_descendants(
|
|
&self,
|
|
root_thread_id: ThreadId,
|
|
) -> CodexResult<Vec<ThreadId>> {
|
|
let mut children_by_parent = self.live_thread_spawn_children().await?;
|
|
let mut descendants = Vec::new();
|
|
let mut stack = children_by_parent
|
|
.remove(&root_thread_id)
|
|
.unwrap_or_default()
|
|
.into_iter()
|
|
.map(|(child_thread_id, _)| child_thread_id)
|
|
.rev()
|
|
.collect::<Vec<_>>();
|
|
|
|
while let Some(thread_id) = stack.pop() {
|
|
descendants.push(thread_id);
|
|
if let Some(children) = children_by_parent.remove(&thread_id) {
|
|
for (child_thread_id, _) in children.into_iter().rev() {
|
|
stack.push(child_thread_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(descendants)
|
|
}
|
|
}
|
|
|
|
fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> bool {
|
|
if prefix.is_root() {
|
|
return true;
|
|
}
|
|
|
|
agent_path.is_some_and(|agent_path| {
|
|
agent_path == prefix
|
|
|| agent_path
|
|
.as_str()
|
|
.strip_prefix(prefix.as_str())
|
|
.is_some_and(|suffix| suffix.starts_with('/'))
|
|
})
|
|
}
|
|
|
|
pub(crate) fn render_input_preview(initial_operation: &Op) -> String {
|
|
match initial_operation {
|
|
Op::UserInput { items, .. } => items
|
|
.iter()
|
|
.map(|item| match item {
|
|
UserInput::Text { text, .. } => text.clone(),
|
|
UserInput::Image { .. } => "[image]".to_string(),
|
|
UserInput::LocalImage { path, .. } => {
|
|
format!("[local_image:{}]", path.display())
|
|
}
|
|
UserInput::Skill { name, path, .. } => {
|
|
format!("[skill:${name}]({})", path.display())
|
|
}
|
|
UserInput::Mention { name, path, .. } => format!("[mention:${name}]({path})"),
|
|
_ => "[input]".to_string(),
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n"),
|
|
Op::InterAgentCommunication { communication } => communication.content.clone(),
|
|
_ => String::new(),
|
|
}
|
|
}
|
|
|
|
fn last_task_message_from_communication(communication: &InterAgentCommunication) -> Option<String> {
|
|
if communication.encrypted_content.is_some() {
|
|
return None;
|
|
}
|
|
non_empty_task_message(communication.content.clone())
|
|
}
|
|
|
|
fn non_empty_task_message(message: String) -> Option<String> {
|
|
(!message.is_empty()).then_some(message)
|
|
}
|
|
|
|
fn thread_spawn_depth(session_source: &SessionSource) -> Option<i32> {
|
|
match session_source {
|
|
SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) => Some(*depth),
|
|
_ => None,
|
|
}
|
|
}
|
|
#[cfg(test)]
|
|
#[path = "control_tests.rs"]
|
|
mod tests;
|