diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 8401b15018..88817b655e 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -530,20 +530,28 @@ impl Session { } } - pub(crate) fn make_collaboration_turn_context( + pub(crate) async fn make_collaboration_turn_context( &self, agent: &AgentState, sub_id: String, - ) -> TurnContext { - Self::make_turn_context( + ) -> Arc { + let per_turn_config = Self::build_per_turn_config(&agent.config); + let model_family = self + .services + .models_manager + .construct_model_family(agent.config.model.as_str(), &per_turn_config) + .await; + Arc::new(Self::make_turn_context( Some(Arc::clone(&self.services.auth_manager)), &self.services.otel_event_manager, agent.config.provider.clone(), &agent.config, + per_turn_config, + model_family, self.conversation_id, sub_id, agent.id, - ) + )) } #[allow(clippy::too_many_arguments)] @@ -1954,6 +1962,7 @@ mod handlers { let warning = EventMsg::Warning(WarningEvent { message }); sess.send_event_raw(Event { id: id.clone(), + agent_idx: Some(0), msg: warning, }) .await; @@ -2537,15 +2546,19 @@ pub(crate) async fn run_collaboration_turn( turn_diff_tracker: SharedTurnDiffTracker, input: Vec, cancellation_token: CancellationToken, -) -> CodexResult> { - run_turn( +) -> CodexResult<(bool, Option)> { + let TurnRunResult { + needs_follow_up, + last_agent_message, + } = run_turn( sess, turn_context, turn_diff_tracker, input, cancellation_token, ) - .await + .await?; + Ok((needs_follow_up, last_agent_message)) } /// When the model is prompted, it returns a stream of events. Some of these diff --git a/codex-rs/core/src/tasks/collaboration.rs b/codex-rs/core/src/tasks/collaboration.rs index cf7b03944f..21ea574d93 100644 --- a/codex-rs/core/src/tasks/collaboration.rs +++ b/codex-rs/core/src/tasks/collaboration.rs @@ -9,7 +9,6 @@ use tokio_util::sync::CancellationToken; use crate::codex::Session; use crate::codex::TurnContext; use crate::codex::run_collaboration_turn; -use crate::response_processing::process_items; use crate::state::AgentId; use crate::state::AgentLifecycleState; use crate::state::TaskKind; @@ -246,8 +245,9 @@ async fn run_agent_turns( collab.next_sub_id(target) }; - let turn_context = - Arc::new(session.make_collaboration_turn_context(&agent_snapshot, sub_id.clone())); + let turn_context = session + .make_collaboration_turn_context(&agent_snapshot, sub_id.clone()) + .await; session.register_sub_id(target, sub_id.clone()).await; let tracker: SharedTurnDiffTracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); @@ -259,25 +259,22 @@ async fn run_agent_turns( Arc::clone(&session), Arc::clone(&turn_context), tracker, - agent_history.get_history(), + agent_history.get_history_for_prompt(), CancellationToken::new(), ) .await; let (delta_tokens, continue_running) = match run_result { - Ok(processed_items) => { - let (responses, items_to_record) = - process_items(processed_items, session.as_ref(), turn_context.as_ref()).await; + Ok((needs_follow_up, last)) => { let new_history = session.clone_history_for_agent(target).await; let after_tokens = new_history.get_total_token_usage(); let delta_tokens = after_tokens .saturating_sub(before_tokens) .clamp(0, i32::MAX as i64) as i32; - let last = crate::codex::get_last_assistant_message_from_turn(&items_to_record); { let mut collab = session.collaboration_state().lock().await; if let Some(agent) = collab.agent_mut(target) { - if !responses.is_empty() { + if needs_follow_up { agent_status = AgentLifecycleState::Running; } else { agent_status = AgentLifecycleState::Idle { @@ -289,7 +286,7 @@ async fn run_agent_turns( } } last_message = last; - (delta_tokens, !responses.is_empty()) + (delta_tokens, needs_follow_up) } Err(err) => { { diff --git a/codex-rs/core/src/tools/handlers/collaboration.rs b/codex-rs/core/src/tools/handlers/collaboration.rs index b39aafd915..0b3b6aa52e 100644 --- a/codex-rs/core/src/tools/handlers/collaboration.rs +++ b/codex-rs/core/src/tools/handlers/collaboration.rs @@ -746,9 +746,9 @@ async fn handle_wait( targets .iter() .filter(|id| { - collab.agent(**id).is_some_and(|agent| { - matches!(agent.status, AgentLifecycleState::Running) - }) + collab + .agent(**id) + .is_some_and(|agent| matches!(agent.status, AgentLifecycleState::Running)) }) .map(|id| id.0) .collect()