diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a7504939c1..89fcbc760c 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -31,9 +31,6 @@ use serde::Deserialize; use serde::Serialize; use serde_json; use tokio::sync::Mutex; -use tokio::sync::mpsc::UnboundedReceiver; -use tokio::sync::mpsc::UnboundedSender; -use tokio::sync::mpsc::unbounded_channel; use tokio::sync::oneshot; use tokio::task::AbortHandle; use tracing::debug; @@ -148,9 +145,8 @@ use self::compact::collect_user_messages; /// It operates as a queue pair where you send submissions and receive events. pub struct Codex { next_id: AtomicU64, - tx_sub: Sender, + tx_sub: Sender, rx_event: Receiver, - turn_readiness_tx: UnboundedSender>, } /// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`], @@ -171,6 +167,13 @@ pub(crate) const MODEL_FORMAT_HEAD_LINES: usize = MODEL_FORMAT_MAX_LINES / 2; pub(crate) const MODEL_FORMAT_TAIL_LINES: usize = MODEL_FORMAT_MAX_LINES - MODEL_FORMAT_HEAD_LINES; // 128 pub(crate) const MODEL_FORMAT_HEAD_BYTES: usize = MODEL_FORMAT_MAX_BYTES / 2; +type TurnReadinessTx = oneshot::Sender>; + +struct PendingSubmission { + submission: Submission, + readiness: Option, +} + impl Codex { /// Spawn a new [`Codex`] and initialize the session. pub async fn spawn( @@ -180,7 +183,6 @@ impl Codex { ) -> CodexResult { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); - let (turn_readiness_tx, turn_readiness_rx) = unbounded_channel(); let user_instructions = get_user_instructions(&config).await; @@ -215,18 +217,11 @@ impl Codex { let conversation_id = session.conversation_id; // This task will run until Op::Shutdown is received. - tokio::spawn(submission_loop( - session, - turn_context, - config, - rx_sub, - turn_readiness_rx, - )); + tokio::spawn(submission_loop(session, turn_context, config, rx_sub)); let codex = Codex { next_id: AtomicU64::new(0), tx_sub, rx_event, - turn_readiness_tx, }; Ok(CodexSpawnOk { @@ -235,10 +230,6 @@ impl Codex { }) } - pub(crate) fn turn_readiness_sender(&self) -> UnboundedSender> { - self.turn_readiness_tx.clone() - } - /// Submit the `op` wrapped in a `Submission` with a unique ID. pub async fn submit(&self, op: Op) -> CodexResult { let id = self @@ -246,18 +237,47 @@ impl Codex { .fetch_add(1, std::sync::atomic::Ordering::SeqCst) .to_string(); let sub = Submission { id: id.clone(), op }; - self.submit_with_id(sub).await?; + self.enqueue_submission(sub, None) + .await + .map_err(|_| CodexErr::InternalAgentDied)?; Ok(id) } /// Use sparingly: prefer `submit()` so Codex is responsible for generating /// unique IDs for each submission. pub async fn submit_with_id(&self, sub: Submission) -> CodexResult<()> { - self.tx_sub - .send(sub) + self.enqueue_submission(sub, None) + .await + .map_err(|_| CodexErr::InternalAgentDied) + } + + pub(crate) async fn submit_with_readiness( + &self, + op: Op, + readiness: Option, + ) -> CodexResult { + let id = self + .next_id + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + .to_string(); + let sub = Submission { id: id.clone(), op }; + self.enqueue_submission(sub, readiness) .await .map_err(|_| CodexErr::InternalAgentDied)?; - Ok(()) + Ok(id) + } + + async fn enqueue_submission( + &self, + submission: Submission, + readiness: Option, + ) -> Result<(), async_channel::SendError> { + self.tx_sub + .send(PendingSubmission { + submission, + readiness, + }) + .await } pub async fn next_event(&self) -> CodexResult { @@ -973,16 +993,6 @@ impl Session { [history, extra].concat() } - pub async fn queue_turn_readiness(&self, flag: Arc) { - let mut state = self.state.lock().await; - state.push_readiness(flag); - } - - pub async fn next_turn_readiness(&self) -> Option> { - let mut state = self.state.lock().await; - state.next_readiness().and_then(super::state::session::TurnReadinessGuard::take) - } - /// Returns the input if there was no task running to inject into pub async fn inject_input( &self, @@ -1019,7 +1029,6 @@ impl Session { { let mut state = self.state.lock().await; state.pending_approvals.clear(); - state.clear_readiness(); } let task = { @@ -1040,7 +1049,6 @@ impl Session { fn interrupt_task_sync(&self) { if let Ok(mut state) = self.state.try_lock() { state.pending_approvals.clear(); - state.clear_readiness(); } if let Ok(mut current_turn) = self.current_turn.try_lock() { current_turn.take(); @@ -1168,17 +1176,15 @@ async fn submission_loop( sess: Arc, turn_context: TurnContext, config: Arc, - rx_sub: Receiver, - mut turn_readiness_rx: UnboundedReceiver>, + rx_sub: Receiver, ) { // Wrap once to avoid cloning TurnContext for each task. let mut turn_context = Arc::new(turn_context); // To break out of this loop, send Op::Shutdown. - while let Ok(sub) = rx_sub.recv().await { - debug!(?sub, "Submission"); - while let Ok(flag) = turn_readiness_rx.try_recv() { - sess.queue_turn_readiness(flag).await; - } + while let Ok(pending) = rx_sub.recv().await { + debug!(?pending.submission, "Submission"); + let mut readiness = pending.readiness; + let sub = pending.submission; match sub.op { Op::Interrupt => { sess.interrupt_task().await; @@ -1272,19 +1278,13 @@ async fn submission_loop( } } Op::UserInput { items } => { - let readiness = match sess.next_turn_readiness().await { - Some(flag) => Some(flag), - None => { - warn!("missing readiness flag for user input"); - None - } - }; - if let Err((items, readiness)) = sess.inject_input(items, readiness.clone()).await { + let readiness_flag = prepare_turn_readiness(&mut readiness).await; + if let Err((items, readiness)) = sess.inject_input(items, readiness_flag).await { let turn_state = Arc::new(TurnState::new( sub.id.clone(), Arc::clone(&turn_context), items, - readiness.clone(), + readiness, )); let task = AgentTask::spawn(sess.clone(), Arc::clone(&turn_state)); sess.set_task(task, Some(turn_state)).await; @@ -1300,14 +1300,8 @@ async fn submission_loop( summary, final_output_json_schema, } => { - let readiness = match sess.next_turn_readiness().await { - Some(flag) => Some(flag), - None => { - warn!("missing readiness flag for user input"); - None - } - }; - if let Err((items, readiness)) = sess.inject_input(items, readiness.clone()).await { + let readiness_flag = prepare_turn_readiness(&mut readiness).await; + if let Err((items, readiness)) = sess.inject_input(items, readiness_flag).await { // Derive a fresh TurnContext for this turn using the provided overrides. let provider = turn_context.client.get_provider(); let auth_manager = turn_context.client.get_auth_manager(); @@ -1374,7 +1368,7 @@ async fn submission_loop( sub.id.clone(), Arc::clone(&turn_context), items, - readiness.clone(), + readiness, )); let task = AgentTask::spawn(sess.clone(), Arc::clone(&turn_state)); sess.set_task(task, Some(turn_state)).await; @@ -1557,10 +1551,40 @@ async fn submission_loop( // Ignore unknown ops; enum is non_exhaustive to allow extensions. } } + send_ready_flag(readiness).await; } debug!("Agent loop exited"); } +async fn prepare_turn_readiness( + readiness: &mut Option, +) -> Option> { + let sender = readiness.take()?; + let flag = Arc::new(ReadinessFlag::new()); + if sender.send(Arc::clone(&flag)).is_err() { + mark_flag_ready(&flag).await; + } + Some(flag) +} + +async fn send_ready_flag(readiness: Option) { + let Some(sender) = readiness else { + return; + }; + let flag = Arc::new(ReadinessFlag::new()); + mark_flag_ready(&flag).await; + let _ = sender.send(flag); +} + +async fn mark_flag_ready(flag: &Arc) { + if flag.is_ready() { + return; + } + if let Ok(token) = flag.subscribe().await { + let _ = flag.mark_ready(token).await; + } +} + /// Spawn a review thread using the given prompt. async fn spawn_review_thread( sess: Arc, diff --git a/codex-rs/core/src/codex_conversation.rs b/codex-rs/core/src/codex_conversation.rs index 2366520cbd..5320259529 100644 --- a/codex-rs/core/src/codex_conversation.rs +++ b/codex-rs/core/src/codex_conversation.rs @@ -5,12 +5,14 @@ use crate::protocol::Op; use crate::protocol::Submission; use codex_utils_readiness::ReadinessFlag; use std::sync::Arc; -use tokio::sync::mpsc::UnboundedSender; +use tokio::sync::oneshot; pub struct CodexConversation { codex: Codex, } +pub type TurnReadinessTx = oneshot::Sender>; + /// Conduit for the bidirectional stream of messages that compose a conversation /// in Codex. impl CodexConversation { @@ -31,7 +33,11 @@ impl CodexConversation { self.codex.next_event().await } - pub fn turn_readiness_sender(&self) -> UnboundedSender> { - self.codex.turn_readiness_sender() + pub async fn submit_with_readiness( + &self, + op: Op, + readiness: Option, + ) -> CodexResult { + self.codex.submit_with_readiness(op, readiness).await } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 20b7a56559..261dd11a78 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -16,6 +16,7 @@ mod codex_conversation; mod state; pub mod token_data; pub use codex_conversation::CodexConversation; +pub use codex_conversation::TurnReadinessTx; pub mod config; pub mod config_edit; pub mod config_profile; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 98310c9534..1ae027b55c 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -1,9 +1,6 @@ use std::collections::HashMap; use std::collections::HashSet; -use std::collections::VecDeque; -use std::sync::Arc; -use codex_utils_readiness::ReadinessFlag; use tokio::sync::oneshot; use crate::conversation_history::ConversationHistory; @@ -18,50 +15,4 @@ pub(crate) struct SessionState { pub(crate) history: ConversationHistory, pub(crate) token_info: Option, pub(crate) latest_rate_limits: Option, - readiness_queue: VecDeque>, -} - -impl SessionState { - pub(crate) fn push_readiness(&mut self, flag: Arc) { - self.readiness_queue.push_back(flag); - } - - pub(crate) fn next_readiness(&mut self) -> Option> { - if self.readiness_queue.is_empty() { - None - } else { - Some(TurnReadinessGuard::new(&mut self.readiness_queue)) - } - } - - pub(crate) fn clear_readiness(&mut self) { - self.readiness_queue.clear(); - } -} - -pub(crate) struct TurnReadinessGuard<'a> { - queue: &'a mut VecDeque>, - consumed: bool, -} - -impl<'a> TurnReadinessGuard<'a> { - fn new(queue: &'a mut VecDeque>) -> Self { - Self { - queue, - consumed: false, - } - } - - pub(crate) fn take(mut self) -> Option> { - self.consumed = true; - self.queue.pop_front() - } -} - -impl Drop for TurnReadinessGuard<'_> { - fn drop(&mut self) { - if !self.consumed { - let _ = self.queue.pop_front(); - } - } } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 149bc7b5af..13ca4931d7 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -112,6 +112,7 @@ use codex_git_tooling::create_ghost_commit; use codex_git_tooling::restore_ghost_commit; use codex_utils_readiness::Readiness; use codex_utils_readiness::ReadinessFlag; +use tokio::sync::oneshot; use tracing::warn; const MAX_TRACKED_GHOST_COMMITS: usize = 20; @@ -185,8 +186,7 @@ pub(crate) struct ChatWidgetInit { pub(crate) struct ChatWidget { app_event_tx: AppEventSender, - codex_op_tx: UnboundedSender, - turn_readiness: UnboundedSender>, + codex_op_tx: UnboundedSender, bottom_pane: BottomPane, active_exec_cell: Option, config: Config, @@ -784,7 +784,6 @@ impl ChatWidget { app_event_tx: app_event_tx.clone(), frame_requester: frame_requester.clone(), codex_op_tx: agent_channels.op_tx, - turn_readiness: agent_channels.turn_readiness, bottom_pane: BottomPane::new(BottomPaneParams { frame_requester, app_event_tx, @@ -846,7 +845,6 @@ impl ChatWidget { app_event_tx: app_event_tx.clone(), frame_requester: frame_requester.clone(), codex_op_tx: agent_channels.op_tx, - turn_readiness: agent_channels.turn_readiness, bottom_pane: BottomPane::new(BottomPaneParams { frame_requester, app_event_tx, @@ -1129,14 +1127,15 @@ impl ChatWidget { return; } - let readiness_flag = Arc::new(ReadinessFlag::new()); - agent::send_turn_readiness(&self.turn_readiness, Arc::clone(&readiness_flag)); - let readiness_to_mark = Arc::clone(&readiness_flag); + let (readiness_tx, readiness_rx) = oneshot::channel::>(); let capture_snapshot = !self.ghost_snapshots_disabled; let repo_path = self.config.cwd.clone(); let app_event_tx = self.app_event_tx.clone(); tokio::spawn(async move { - let readiness_token = readiness_to_mark.subscribe().await.ok(); + let Ok(flag) = readiness_rx.await else { + return; + }; + let readiness_token = flag.subscribe().await.ok(); if capture_snapshot { let event = match create_ghost_commit(&CreateGhostCommitOptions::new( repo_path.as_path(), @@ -1167,7 +1166,7 @@ impl ChatWidget { app_event_tx.send(event); } if let Some(token) = readiness_token { - let _ = readiness_to_mark.mark_ready(token).await; + let _ = flag.mark_ready(token).await; } }); @@ -1182,7 +1181,10 @@ impl ChatWidget { } self.codex_op_tx - .send(Op::UserInput { items }) + .send(agent::OutgoingOp { + op: Op::UserInput { items }, + readiness: Some(readiness_tx), + }) .unwrap_or_else(|e| { tracing::error!("failed to send message: {e}"); }); @@ -1190,7 +1192,10 @@ impl ChatWidget { // Persist the text to cross-session message history. if !text.is_empty() { self.codex_op_tx - .send(Op::AddToHistory { text: text.clone() }) + .send(agent::OutgoingOp { + op: Op::AddToHistory { text: text.clone() }, + readiness: None, + }) .unwrap_or_else(|e| { tracing::error!("failed to send AddHistory op: {e}"); }); @@ -1679,7 +1684,10 @@ impl ChatWidget { pub(crate) fn submit_op(&self, op: Op) { // Record outbound operation for session replay fidelity. crate::session_log::log_outbound_op(&op); - if let Err(e) = self.codex_op_tx.send(op) { + if let Err(e) = self.codex_op_tx.send(agent::OutgoingOp { + op, + readiness: None, + }) { tracing::error!("failed to submit op: {e}"); } } diff --git a/codex-rs/tui/src/chatwidget/agent.rs b/codex-rs/tui/src/chatwidget/agent.rs index eeb99ec9f3..4c3a0cb886 100644 --- a/codex-rs/tui/src/chatwidget/agent.rs +++ b/codex-rs/tui/src/chatwidget/agent.rs @@ -3,11 +3,9 @@ use std::sync::Arc; use codex_core::CodexConversation; use codex_core::ConversationManager; use codex_core::NewConversation; +use codex_core::TurnReadinessTx; use codex_core::config::Config; use codex_core::protocol::Op; -use codex_utils_readiness::Readiness; -use codex_utils_readiness::ReadinessFlag; -use tokio::sync::mpsc::UnboundedReceiver; use tokio::sync::mpsc::UnboundedSender; use tokio::sync::mpsc::unbounded_channel; @@ -15,37 +13,12 @@ use crate::app_event::AppEvent; use crate::app_event_sender::AppEventSender; pub(crate) struct AgentChannels { - pub(crate) op_tx: UnboundedSender, - pub(crate) turn_readiness: UnboundedSender>, -} -fn mark_ready(flag: Arc) { - tokio::spawn(async move { - if let Ok(token) = flag.subscribe().await { - let _ = flag.mark_ready(token).await; - } - }); + pub(crate) op_tx: UnboundedSender, } -fn spawn_readiness_forwarder( - mut rx: UnboundedReceiver>, - sender: UnboundedSender>, -) { - tokio::spawn(async move { - while let Some(flag) = rx.recv().await { - if sender.send(Arc::clone(&flag)).is_err() { - mark_ready(flag); - } - } - }); -} - -pub(crate) fn send_turn_readiness( - sender: &UnboundedSender>, - flag: Arc, -) { - if sender.send(Arc::clone(&flag)).is_err() { - mark_ready(flag); - } +pub(crate) struct OutgoingOp { + pub(crate) op: Op, + pub(crate) readiness: Option, } /// Spawn the agent bootstrapper and op forwarding loop, returning the @@ -55,8 +28,7 @@ pub(crate) fn spawn_agent( app_event_tx: AppEventSender, server: Arc, ) -> AgentChannels { - let (codex_op_tx, mut codex_op_rx) = unbounded_channel::(); - let (turn_readiness_tx, turn_readiness_rx) = unbounded_channel::>(); + let (codex_op_tx, mut codex_op_rx) = unbounded_channel::(); let app_event_tx_clone = app_event_tx; tokio::spawn(async move { @@ -73,9 +45,6 @@ pub(crate) fn spawn_agent( } }; - let readiness_sender = conversation.turn_readiness_sender(); - spawn_readiness_forwarder(turn_readiness_rx, readiness_sender); - // Forward the captured `SessionConfigured` event so it can be rendered in the UI. let ev = codex_core::protocol::Event { // The `id` does not matter for rendering, so we can use a fake value. @@ -86,8 +55,10 @@ pub(crate) fn spawn_agent( let conversation_clone = conversation.clone(); tokio::spawn(async move { - while let Some(op) = codex_op_rx.recv().await { - let id = conversation_clone.submit(op).await; + while let Some(outgoing) = codex_op_rx.recv().await { + let id = conversation_clone + .submit_with_readiness(outgoing.op, outgoing.readiness) + .await; if let Err(e) = id { tracing::error!("failed to submit op: {e}"); } @@ -99,10 +70,7 @@ pub(crate) fn spawn_agent( } }); - AgentChannels { - op_tx: codex_op_tx, - turn_readiness: turn_readiness_tx, - } + AgentChannels { op_tx: codex_op_tx } } /// Spawn agent loops for an existing conversation (e.g., a forked conversation). @@ -113,9 +81,7 @@ pub(crate) fn spawn_agent_from_existing( session_configured: codex_core::protocol::SessionConfiguredEvent, app_event_tx: AppEventSender, ) -> AgentChannels { - let (codex_op_tx, mut codex_op_rx) = unbounded_channel::(); - let (turn_readiness_tx, turn_readiness_rx) = unbounded_channel::>(); - spawn_readiness_forwarder(turn_readiness_rx, conversation.turn_readiness_sender()); + let (codex_op_tx, mut codex_op_rx) = unbounded_channel::(); let app_event_tx_clone = app_event_tx; tokio::spawn(async move { @@ -128,8 +94,10 @@ pub(crate) fn spawn_agent_from_existing( let conversation_clone = conversation.clone(); tokio::spawn(async move { - while let Some(op) = codex_op_rx.recv().await { - let id = conversation_clone.submit(op).await; + while let Some(outgoing) = codex_op_rx.recv().await { + let id = conversation_clone + .submit_with_readiness(outgoing.op, outgoing.readiness) + .await; if let Err(e) = id { tracing::error!("failed to submit op: {e}"); } @@ -141,8 +109,5 @@ pub(crate) fn spawn_agent_from_existing( } }); - AgentChannels { - op_tx: codex_op_tx, - turn_readiness: turn_readiness_tx, - } + AgentChannels { op_tx: codex_op_tx } } diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index fc1f7f4a97..d72e7866f9 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::app_event::AppEvent; use crate::app_event_sender::AppEventSender; +use crate::chatwidget::agent; use codex_core::AuthManager; use codex_core::CodexAuth; use codex_core::config::Config; @@ -34,7 +35,6 @@ use codex_core::protocol::StreamErrorEvent; use codex_core::protocol::TaskCompleteEvent; use codex_core::protocol::TaskStartedEvent; use codex_protocol::mcp_protocol::ConversationId; -use codex_utils_readiness::ReadinessFlag; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyModifiers; @@ -297,11 +297,11 @@ async fn helpers_are_available_and_do_not_panic() { fn make_chatwidget_manual() -> ( ChatWidget, tokio::sync::mpsc::UnboundedReceiver, - tokio::sync::mpsc::UnboundedReceiver, + tokio::sync::mpsc::UnboundedReceiver, ) { let (tx_raw, rx) = unbounded_channel::(); let app_event_tx = AppEventSender::new(tx_raw); - let (op_tx, op_rx) = unbounded_channel::(); + let (op_tx, op_rx) = unbounded_channel::(); let cfg = test_config(); let bottom = BottomPane::new(BottomPaneParams { app_event_tx: app_event_tx.clone(), @@ -312,11 +312,9 @@ fn make_chatwidget_manual() -> ( disable_paste_burst: false, }); let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test")); - let (turn_readiness, _rx) = unbounded_channel::>(); let widget = ChatWidget { app_event_tx, codex_op_tx: op_tx, - turn_readiness, bottom_pane: bottom, active_exec_cell: None, config: cfg.clone(), @@ -349,7 +347,7 @@ pub(crate) fn make_chatwidget_manual_with_sender() -> ( ChatWidget, AppEventSender, tokio::sync::mpsc::UnboundedReceiver, - tokio::sync::mpsc::UnboundedReceiver, + tokio::sync::mpsc::UnboundedReceiver, ) { let (widget, rx, op_rx) = make_chatwidget_manual(); let app_event_tx = widget.app_event_tx.clone(); @@ -1676,7 +1674,8 @@ fn apply_patch_full_flow_integration_like() { let forwarded = op_rx .try_recv() .expect("expected op forwarded to codex channel"); - match forwarded { + assert!(forwarded.readiness.is_none()); + match forwarded.op { Op::PatchApproval { id, decision } => { assert_eq!(id, "sub-xyz"); assert!(matches!(