From 4cc89bb57d408a141f4367c5e42449d7b13c032a Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 16 Dec 2025 17:05:09 +0000 Subject: [PATCH] tool id --- codex-rs/core/src/agents/mod.rs | 4 +- codex-rs/core/src/codex.rs | 68 ++++++++++--------- codex-rs/core/src/state/collaboration.rs | 83 +++++++++++++++--------- 3 files changed, 93 insertions(+), 62 deletions(-) diff --git a/codex-rs/core/src/agents/mod.rs b/codex-rs/core/src/agents/mod.rs index c3e52e2502..2b92540178 100644 --- a/codex-rs/core/src/agents/mod.rs +++ b/codex-rs/core/src/agents/mod.rs @@ -124,8 +124,8 @@ impl AgentsConfig { } fn validate_agents(agents: &HashMap) -> Result<(), String> { - if !agents.contains_key("main") { - return Err("missing required agent: main".to_string()); + if !agents.contains_key("orchestrator") { + return Err("missing required agent: orchestrator".to_string()); } for agent in agents.values() { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index d7f40167e0..729d9129a2 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -398,7 +398,7 @@ impl TurnContext { } pub(crate) fn collaboration_agent(&self) -> AgentId { - self.collaboration_agent + self.collaboration_agent.clone() } } @@ -610,7 +610,7 @@ impl Session { model_family, self.conversation_id, sub_id, - agent.id, + agent.id.clone(), ); if let Some(agents_config) = self.agents_config() && let Some(agent_config) = agents_config.agent(agent.name.as_str()) @@ -1004,6 +1004,7 @@ impl Session { .models_manager .construct_model_family(session_configuration.model.as_str(), &per_turn_config) .await; + let root_id = AgentId::root(); let mut turn_context: TurnContext = Self::make_turn_context( Some(Arc::clone(&self.services.auth_manager)), &self.services.otel_manager, @@ -1013,7 +1014,7 @@ impl Session { model_family, self.conversation_id, sub_id, - AgentId(0), + root_id.clone(), ); if let Some(final_schema) = updates.final_output_json_schema { turn_context.final_output_json_schema = final_schema; @@ -1029,12 +1030,12 @@ impl Session { let mut collab = self.collaboration.lock().await; collab.ensure_root_agent(&session_configuration, &session_history); drop(collab); - self.register_sub_id(AgentId(0), turn_context.sub_id.clone()) + self.register_sub_id(&root_id, turn_context.sub_id.clone()) .await; Arc::new(turn_context) } - pub(crate) async fn register_sub_id(&self, agent: AgentId, sub_id: String) { + pub(crate) async fn register_sub_id(&self, agent: &AgentId, sub_id: String) { let mut collab = self.collaboration.lock().await; collab.register_sub_id(agent, sub_id); } @@ -1067,11 +1068,17 @@ impl Session { /// Persist the event to rollout and send it to clients. pub(crate) async fn send_event(&self, turn_context: &TurnContext, msg: EventMsg) { let agent = turn_context.collaboration_agent(); - if !Self::should_emit_event_for_agent(agent, &msg) { + let is_root = agent.is_root(); + if !Self::should_emit_event_for_agent(is_root, &msg) { return; } let legacy_source = msg.clone(); - let agent_idx = Some(turn_context.collaboration_agent().0); + let agent_idx = if is_root { + Some(0) + } else { + let collab = self.collaboration.lock().await; + collab.agent_index(&agent) + }; let event = Event { id: turn_context.sub_id.clone(), agent_idx, @@ -1090,8 +1097,8 @@ impl Session { } } - fn should_emit_event_for_agent(agent: AgentId, msg: &EventMsg) -> bool { - if agent == AgentId(0) { + fn should_emit_event_for_agent(is_root: bool, msg: &EventMsg) -> bool { + if is_root { return true; } !matches!( @@ -1115,7 +1122,7 @@ impl Session { pub(crate) async fn send_event_raw(&self, event: Event) { if let Some(agent_idx) = event.agent_idx && agent_idx != 0 - && !Self::should_emit_event_for_agent(AgentId(agent_idx), &event.msg) + && !Self::should_emit_event_for_agent(false, &event.msg) { return; } @@ -1362,7 +1369,7 @@ impl Session { turn_context: &TurnContext, ) { let agent = turn_context.collaboration_agent(); - if agent == AgentId(0) { + if agent.is_root() { let (history, token_info, config) = { let mut state = self.state.lock().await; state.record_items(items.iter(), turn_context.truncation_policy); @@ -1375,13 +1382,14 @@ impl Session { let mut collab = self.collaboration.lock().await; collab.ensure_root_agent(&config, &history); - if let Some(root) = collab.agent_mut(AgentId(0)) { + let root_id = AgentId::root(); + if let Some(root) = collab.agent_mut(&root_id) { root.history = history; root.history.set_token_info(token_info); } } else { let mut collab = self.collaboration.lock().await; - if let Some(agent_state) = collab.agent_mut(agent) { + if let Some(agent_state) = collab.agent_mut(&agent) { agent_state .history .record_items(items.iter(), turn_context.truncation_policy); @@ -1406,7 +1414,8 @@ impl Session { } pub(crate) async fn replace_history(&self, items: Vec) { - self.set_history_for_agent(AgentId(0), items, None).await; + let root_id = AgentId::root(); + self.set_history_for_agent(&root_id, items, None).await; } async fn persist_rollout_response_items(&self, items: &[ResponseItem]) { @@ -1473,11 +1482,12 @@ impl Session { } pub(crate) async fn clone_history(&self) -> ContextManager { - self.clone_history_for_agent(AgentId(0)).await + let root_id = AgentId::root(); + self.clone_history_for_agent(&root_id).await } - pub(crate) async fn clone_history_for_agent(&self, agent: AgentId) -> ContextManager { - if agent == AgentId(0) { + pub(crate) async fn clone_history_for_agent(&self, agent: &AgentId) -> ContextManager { + if agent.is_root() { let state = self.state.lock().await; return state.clone_history(); } @@ -1489,11 +1499,11 @@ impl Session { pub(crate) async fn set_history_for_agent( &self, - agent: AgentId, + agent: &AgentId, items: Vec, token_info: Option, ) { - if agent == AgentId(0) { + if agent.is_root() { let mut state = self.state.lock().await; state.replace_history(items.clone()); state.set_token_info(token_info.clone()); @@ -1509,7 +1519,7 @@ impl Session { token_usage: Option<&TokenUsage>, ) { let agent = turn_context.collaboration_agent(); - if agent == AgentId(0) { + if agent == AgentId::root() { let token_info = { let mut state = self.state.lock().await; if let Some(token_usage) = token_usage { @@ -1522,7 +1532,7 @@ impl Session { }; { let mut collab = self.collaboration.lock().await; - if let Some(root) = collab.agent_mut(AgentId(0)) { + if let Some(root) = collab.agent_mut(AgentId::root()) { root.history.set_token_info(token_info); } } @@ -1539,12 +1549,12 @@ impl Session { pub(crate) async fn recompute_token_usage(&self, turn_context: &TurnContext) { let agent = turn_context.collaboration_agent(); - let history = self.clone_history_for_agent(agent).await; + let history = self.clone_history_for_agent(&agent).await; let Some(estimated_total_tokens) = history.estimate_token_count(turn_context) else { return; }; - if agent == AgentId(0) { + if agent == AgentId::root() { let token_info = { let mut state = self.state.lock().await; let mut info = state.token_info().unwrap_or(TokenUsageInfo { @@ -1570,7 +1580,7 @@ impl Session { }; { let mut collab = self.collaboration.lock().await; - if let Some(root) = collab.agent_mut(AgentId(0)) { + if let Some(root) = collab.agent_mut(AgentId::root()) { root.history.set_token_info(Some(token_info)); } } @@ -1623,7 +1633,7 @@ impl Session { let agent = turn_context.collaboration_agent(); let context_window = turn_context.client.get_model_context_window(); if let Some(context_window) = context_window { - if agent == AgentId(0) { + if agent == AgentId::root() { let token_info = { let mut state = self.state.lock().await; state.set_token_usage_full(context_window); @@ -1631,7 +1641,7 @@ impl Session { }; { let mut collab = self.collaboration.lock().await; - if let Some(root) = collab.agent_mut(AgentId(0)) { + if let Some(root) = collab.agent_mut(AgentId::root()) { root.history.set_token_info(token_info); } } @@ -2403,7 +2413,7 @@ async fn spawn_review_thread( tool_call_gate: Arc::new(ReadinessFlag::new()), exec_policy: parent_turn_context.exec_policy.clone(), truncation_policy: TruncationPolicy::new(&per_turn_config, model_family.truncation_policy), - collaboration_agent: AgentId(0), + collaboration_agent: AgentId::root(), }; // Seed the child task with the review prompt as the initial user message. @@ -3439,7 +3449,7 @@ mod tests { model_family, conversation_id, "turn_id".to_string(), - AgentId(0), + AgentId::root(), ); let session = Session { @@ -3536,7 +3546,7 @@ mod tests { model_family, conversation_id, "turn_id".to_string(), - AgentId(0), + AgentId::root(), )); let session = Arc::new(Session { diff --git a/codex-rs/core/src/state/collaboration.rs b/codex-rs/core/src/state/collaboration.rs index b5661810e4..51e87b0987 100644 --- a/codex-rs/core/src/state/collaboration.rs +++ b/codex-rs/core/src/state/collaboration.rs @@ -1,6 +1,7 @@ //! Session-scoped collaboration state for multi-agent flows. use std::collections::HashMap; +use std::fmt; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; @@ -45,6 +46,16 @@ impl AgentId { pub fn random() -> Self { Self(uuid::Uuid::new_v4().to_string()) } + + pub fn is_root(&self) -> bool { + self.0 == "root" + } +} + +impl fmt::Display for AgentId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } } #[allow(dead_code)] @@ -150,6 +161,7 @@ pub(crate) struct CollaborationState { limits: CollaborationLimits, next_sub_id: i64, sub_ids: HashMap, + agent_indices: HashMap, } impl CollaborationState { @@ -160,6 +172,7 @@ impl CollaborationState { limits, next_sub_id: 0, sub_ids: HashMap::new(), + agent_indices: HashMap::new(), } } @@ -174,7 +187,7 @@ impl CollaborationState { ) -> AgentId { if self.agents.is_empty() { let root = AgentState::new_root( - "main".to_string(), + "orchestrator".to_string(), session_configuration.clone(), session_history.clone(), session_configuration @@ -182,6 +195,8 @@ impl CollaborationState { .or_else(|| session_configuration.user_instructions()), ); self.agents.push(root); + self.agent_indices + .insert(self.agents[0].id.clone(), 0); } else if let Some(root) = self.agents.get_mut(0) { root.config = session_configuration.clone(); root.history = session_history.clone(); @@ -190,6 +205,7 @@ impl CollaborationState { .developer_instructions() .or_else(|| session_configuration.user_instructions()); } + self.agent_indices.insert(root.id.clone(), 0); } AgentId::root() } @@ -198,34 +214,34 @@ impl CollaborationState { &self.agents } - pub(crate) fn agent(&self, id: AgentId) -> Option<&AgentState> { + pub(crate) fn agent(&self, id: &AgentId) -> Option<&AgentState> { self.index_for(id).and_then(|idx| self.agents.get(idx)) } - pub(crate) fn agent_mut(&mut self, id: AgentId) -> Option<&mut AgentState> { + pub(crate) fn agent_mut(&mut self, id: &AgentId) -> Option<&mut AgentState> { let index = self.index_for(id)?; self.agents.get_mut(index) } - pub(crate) fn clone_agent_history(&self, id: AgentId) -> Option { + pub(crate) fn clone_agent_history(&self, id: &AgentId) -> Option { self.agent(id).map(|agent| agent.history.clone()) } pub(crate) fn set_agent_history( &mut self, - id: AgentId, + id: &AgentId, items: Vec, token_info: Option, ) -> Result<(), String> { let agent = self .agent_mut(id) - .ok_or_else(|| format!("unknown agent {}", id.0))?; + .ok_or_else(|| format!("unknown agent {id}"))?; agent.history.replace(items); agent.history.set_token_info(token_info); Ok(()) } - pub(crate) fn record_message_for_agent(&mut self, id: AgentId, message: ResponseItem) { + pub(crate) fn record_message_for_agent(&mut self, id: &AgentId, message: ResponseItem) { let role = match &message { ResponseItem::Message { role, .. } => role.as_str(), _ => "other", @@ -233,7 +249,7 @@ impl CollaborationState { let content = content_for_log(&message); if let Some(agent) = self.agent_mut(id) { warn!( - agent_idx = id.0, + agent_idx = %id, agent_name = agent.name.as_str(), role, content, @@ -244,7 +260,7 @@ impl CollaborationState { .record_items([message].iter(), TruncationPolicy::Bytes(10_000)); } else { warn!( - agent_idx = id.0, + agent_idx = %id, agent_name = "", role, content, @@ -256,7 +272,7 @@ impl CollaborationState { #[allow(dead_code)] pub(crate) fn record_items_for_agent( &mut self, - id: AgentId, + id: &AgentId, items: &[ResponseItem], policy: TruncationPolicy, ) { @@ -273,20 +289,24 @@ impl CollaborationState { return Err("max collaboration depth reached".to_string()); } - let id = AgentId::random(); - agent.id = id.clone(); + let id = agent.id.clone(); + if self.agent_indices.contains_key(&id) { + return Err(format!("duplicate agent id {id}")); + } if let Some(parent) = agent.parent.as_ref() { self.children.entry(parent.clone()).or_default().push(id.clone()); } + let index = self.agents.len(); self.agents.push(agent); + self.agent_indices.insert(id.clone(), index); Ok(id) } - pub(crate) fn is_direct_child(&self, parent: AgentId, child: AgentId) -> bool { + pub(crate) fn is_direct_child(&self, parent: &AgentId, child: &AgentId) -> bool { self.children - .get(&parent) + .get(parent) .map(|kids| kids.contains(&child)) .unwrap_or(false) } @@ -295,39 +315,40 @@ impl CollaborationState { let mut result = Vec::new(); let mut stack: Vec = roots.to_vec(); while let Some(id) = stack.pop() { - result.push(id); if let Some(children) = self.children.get(&id) { for child in children { - stack.push(*child); + stack.push(child.clone()); } } + result.push(id); } result } - pub(crate) fn next_sub_id(&mut self, agent: AgentId) -> String { - let sub_id = format!("collab-agent-{}-{}", agent.0, self.next_sub_id); + pub(crate) fn next_agent_id(&self) -> AgentId { + AgentId::random() + } + + pub(crate) fn next_sub_id(&mut self, agent: &AgentId) -> String { + let next_sub_id = self.next_sub_id; + let sub_id = format!("collab-agent-{agent}-{next_sub_id}"); self.next_sub_id += 1; sub_id } - fn index_for(&self, id: AgentId) -> Option { - if id.0 < 0 { - return None; - } - let index = id.0 as usize; - if index < self.agents.len() { - Some(index) - } else { - None - } + pub(crate) fn agent_index(&self, id: &AgentId) -> Option { + self.index_for(id).and_then(|idx| i32::try_from(idx).ok()) } - pub(crate) fn register_sub_id(&mut self, agent: AgentId, sub_id: String) { - self.sub_ids.insert(sub_id, agent); + fn index_for(&self, id: &AgentId) -> Option { + self.agent_indices.get(id).copied() + } + + pub(crate) fn register_sub_id(&mut self, agent: &AgentId, sub_id: String) { + self.sub_ids.insert(sub_id, agent.clone()); } pub(crate) fn agent_for_sub_id(&self, sub_id: &str) -> Option { - self.sub_ids.get(sub_id).copied() + self.sub_ids.get(sub_id).cloned() } }