openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-14 11:57:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

NV 1

Browse Source
This commit is contained in:
jimmyfraiture
2025-09-24 15:08:34 +01:00
parent 40d17f68e1
commit e05540ea7f
7 changed files with 136 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
codex-rs/core/src/codex.rs
View File

@@ -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<Submission>,
tx_sub: Sender<PendingSubmission>,
rx_event: Receiver<Event>,
turn_readiness_tx: UnboundedSender<Arc<ReadinessFlag>>,
}
/// 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<Arc<ReadinessFlag>>;
struct PendingSubmission {
submission: Submission,
readiness: Option<TurnReadinessTx>,
}
impl Codex {
/// Spawn a new [`Codex`] and initialize the session.
pub async fn spawn(
@@ -180,7 +183,6 @@ impl Codex {
) -> CodexResult<CodexSpawnOk> {
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<Arc<ReadinessFlag>> {
self.turn_readiness_tx.clone()
}
/// Submit the `op` wrapped in a `Submission` with a unique ID.
pub async fn submit(&self, op: Op) -> CodexResult<String> {
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<TurnReadinessTx>,
) -> CodexResult<String> {
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<TurnReadinessTx>,
) -> Result<(), async_channel::SendError<PendingSubmission>> {
self.tx_sub
.send(PendingSubmission {
submission,
readiness,
})
.await
}
pub async fn next_event(&self) -> CodexResult<Event> {
@@ -973,16 +993,6 @@ impl Session {
[history, extra].concat()
}
pub async fn queue_turn_readiness(&self, flag: Arc<ReadinessFlag>) {
let mut state = self.state.lock().await;
state.push_readiness(flag);
}
pub async fn next_turn_readiness(&self) -> Option<Arc<ReadinessFlag>> {
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<Session>,
turn_context: TurnContext,
config: Arc<Config>,
rx_sub: Receiver<Submission>,
mut turn_readiness_rx: UnboundedReceiver<Arc<ReadinessFlag>>,
rx_sub: Receiver<PendingSubmission>,
) {
// 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<TurnReadinessTx>,
) -> Option<Arc<ReadinessFlag>> {
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<TurnReadinessTx>) {
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<ReadinessFlag>) {
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<Session>,

12
codex-rs/core/src/codex_conversation.rs
View File

@@ -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<Arc<ReadinessFlag>>;
/// 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<Arc<ReadinessFlag>> {
self.codex.turn_readiness_sender()
pub async fn submit_with_readiness(
&self,
op: Op,
readiness: Option<TurnReadinessTx>,
) -> CodexResult<String> {
self.codex.submit_with_readiness(op, readiness).await
}
}

1
codex-rs/core/src/lib.rs
View File

@@ -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;

49
codex-rs/core/src/state/session.rs
View File

@@ -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<TokenUsageInfo>,
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
readiness_queue: VecDeque<Arc<ReadinessFlag>>,
}
impl SessionState {
pub(crate) fn push_readiness(&mut self, flag: Arc<ReadinessFlag>) {
self.readiness_queue.push_back(flag);
}
pub(crate) fn next_readiness(&mut self) -> Option<TurnReadinessGuard<'_>> {
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<Arc<ReadinessFlag>>,
consumed: bool,
}
impl<'a> TurnReadinessGuard<'a> {
fn new(queue: &'a mut VecDeque<Arc<ReadinessFlag>>) -> Self {
Self {
queue,
consumed: false,
}
}
pub(crate) fn take(mut self) -> Option<Arc<ReadinessFlag>> {
self.consumed = true;
self.queue.pop_front()
}
}
impl Drop for TurnReadinessGuard<'_> {
fn drop(&mut self) {
if !self.consumed {
let _ = self.queue.pop_front();
}
}
}

32
codex-rs/tui/src/chatwidget.rs
View File

@@ -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<Op>,
turn_readiness: UnboundedSender<Arc<ReadinessFlag>>,
codex_op_tx: UnboundedSender<agent::OutgoingOp>,
bottom_pane: BottomPane,
active_exec_cell: Option<ExecCell>,
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::<Arc<ReadinessFlag>>();
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}");
}
}

69
codex-rs/tui/src/chatwidget/agent.rs
View File

@@ -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<Op>,
pub(crate) turn_readiness: UnboundedSender<Arc<ReadinessFlag>>,
}
fn mark_ready(flag: Arc<ReadinessFlag>) {
tokio::spawn(async move {
if let Ok(token) = flag.subscribe().await {
let _ = flag.mark_ready(token).await;
}
});
pub(crate) op_tx: UnboundedSender<OutgoingOp>,
}
fn spawn_readiness_forwarder(
mut rx: UnboundedReceiver<Arc<ReadinessFlag>>,
sender: UnboundedSender<Arc<ReadinessFlag>>,
) {
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<Arc<ReadinessFlag>>,
flag: Arc<ReadinessFlag>,
) {
if sender.send(Arc::clone(&flag)).is_err() {
mark_ready(flag);
}
pub(crate) struct OutgoingOp {
pub(crate) op: Op,
pub(crate) readiness: Option<TurnReadinessTx>,
}
/// 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<ConversationManager>,
) -> AgentChannels {
let (codex_op_tx, mut codex_op_rx) = unbounded_channel::<Op>();
let (turn_readiness_tx, turn_readiness_rx) = unbounded_channel::<Arc<ReadinessFlag>>();
let (codex_op_tx, mut codex_op_rx) = unbounded_channel::<OutgoingOp>();
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::<Op>();
let (turn_readiness_tx, turn_readiness_rx) = unbounded_channel::<Arc<ReadinessFlag>>();
spawn_readiness_forwarder(turn_readiness_rx, conversation.turn_readiness_sender());
let (codex_op_tx, mut codex_op_rx) = unbounded_channel::<OutgoingOp>();
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 }
}

13
codex-rs/tui/src/chatwidget/tests.rs
View File

@@ -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<AppEvent>,
tokio::sync::mpsc::UnboundedReceiver<Op>,
tokio::sync::mpsc::UnboundedReceiver<agent::OutgoingOp>,
) {
let (tx_raw, rx) = unbounded_channel::<AppEvent>();
let app_event_tx = AppEventSender::new(tx_raw);
let (op_tx, op_rx) = unbounded_channel::<Op>();
let (op_tx, op_rx) = unbounded_channel::<agent::OutgoingOp>();
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::<Arc<ReadinessFlag>>();
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<AppEvent>,
tokio::sync::mpsc::UnboundedReceiver<Op>,
tokio::sync::mpsc::UnboundedReceiver<agent::OutgoingOp>,
) {
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!(