openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-01 19:33:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'etraut/next-turn-state-remove-user-input-with-context' into etraut/next-turn-state-input-op-consolidation

Browse Source 
# Conflicts:
#	codex-rs/core/src/session/handlers.rs
#	codex-rs/core/tests/common/test_codex.rs
#	codex-rs/core/tests/suite/request_permissions.rs
#	codex-rs/core/tests/suite/unified_exec.rs
#	codex-rs/protocol/src/protocol.rs
This commit is contained in:
Eric Traut
2026-05-18 19:22:12 -07:00
parent 4939645c44 6927359c0c
commit 45fe8ff8cc
285 changed files with 11487 additions and 3986 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/handlers.rs
+15 -9
View File
@@ -263,7 +263,9 @@ pub async fn inter_agent_communication(
communication: InterAgentCommunication,
) {
let trigger_turn = communication.trigger_turn;
sess.enqueue_mailbox_communication(communication);
sess.input_queue
.enqueue_mailbox_communication(communication)
.await;
if trigger_turn {
sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id)
.await;
@@ -582,12 +584,14 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
sess.guardian_review_session.shutdown().await;
}
fn emit_thread_stop_lifecycle(sess: &Session) {
async fn emit_thread_stop_lifecycle(sess: &Session) {
for contributor in sess.services.extensions.thread_lifecycle_contributors() {
contributor.on_thread_stop(codex_extension_api::ThreadStopInput {
session_store: &sess.services.session_extension_data,
thread_store: &sess.services.thread_extension_data,
});
contributor
.on_thread_stop(codex_extension_api::ThreadStopInput {
session_store: &sess.services.session_extension_data,
thread_store: &sess.services.thread_extension_data,
})
.await;
}
}
@@ -606,7 +610,7 @@ pub async fn shutdown(sess: &Arc<Session>, sub_id: String) -> bool {
&[],
);
emit_thread_stop_lifecycle(sess.as_ref());
emit_thread_stop_lifecycle(sess.as_ref()).await;
// Gracefully flush and shutdown thread persistence on session end so tests
// that inspect durable state do not race with the background writer.
@@ -859,7 +863,7 @@ pub(super) async fn submission_loop(
// explicit shutdown op, still run session teardown.
if !shutdown_received {
shutdown_session_runtime(&sess).await;
emit_thread_stop_lifecycle(sess.as_ref());
emit_thread_stop_lifecycle(sess.as_ref()).await;
}
debug!("Agent loop exited");
}
@@ -901,7 +905,9 @@ Approved action:
}];
if let Err(items) = sess.inject_response_items(items).await {
sess.queue_response_items_for_next_turn(items).await;
sess.input_queue
.queue_response_items_for_next_turn(items)
.await;
}
}
codex-rs/core/src/session/input_queue.rs
+392
View File
@@ -0,0 +1,392 @@
use crate::state::ActiveTurn;
use crate::state::MailboxDeliveryPhase;
use crate::state::TurnState;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::protocol::InterAgentCommunication;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::watch;
/// Turn-local pending input storage owned by the input queue flow.
#[derive(Default)]
pub(crate) struct TurnInputQueue {
items: Vec<ResponseInputItem>,
}
/// Session-scoped pending input storage and active-turn mailbox delivery coordination.
pub(crate) struct InputQueue {
mailbox_tx: watch::Sender<()>,
mailbox_pending_mails: Mutex<VecDeque<InterAgentCommunication>>,
idle_pending_input: Mutex<Vec<ResponseInputItem>>,
}
impl InputQueue {
pub(crate) fn new() -> Self {
let (mailbox_tx, _) = watch::channel(());
Self {
mailbox_tx,
mailbox_pending_mails: Mutex::new(VecDeque::new()),
idle_pending_input: Mutex::new(Vec::new()),
}
}
pub(crate) async fn subscribe_mailbox(&self) -> watch::Receiver<()> {
let mut mailbox_rx = self.mailbox_tx.subscribe();
if self.has_pending_mailbox_items().await {
mailbox_rx.mark_changed();
}
mailbox_rx
}
pub(crate) async fn enqueue_mailbox_communication(
&self,
communication: InterAgentCommunication,
) {
self.mailbox_pending_mails
.lock()
.await
.push_back(communication);
self.mailbox_tx.send_replace(());
}
pub(crate) async fn has_pending_mailbox_items(&self) -> bool {
!self.mailbox_pending_mails.lock().await.is_empty()
}
pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool {
self.mailbox_pending_mails
.lock()
.await
.iter()
.any(|mail| mail.trigger_turn)
}
pub(crate) async fn drain_mailbox_input_items(&self) -> Vec<ResponseInputItem> {
self.mailbox_pending_mails
.lock()
.await
.drain(..)
.map(|mail| mail.to_response_input_item())
.collect()
}
pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec<ResponseInputItem>) {
if items.is_empty() {
return;
}
self.idle_pending_input.lock().await.extend(items);
}
pub(crate) async fn take_queued_response_items_for_next_turn(&self) -> Vec<ResponseInputItem> {
std::mem::take(&mut *self.idle_pending_input.lock().await)
}
pub(crate) async fn has_queued_response_items_for_next_turn(&self) -> bool {
!self.idle_pending_input.lock().await.is_empty()
}
pub(crate) async fn turn_state_for_sub_id(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
sub_id: &str,
) -> Option<Arc<Mutex<TurnState>>> {
let active = active_turn.lock().await;
active.as_ref().and_then(|active_turn| {
active_turn
.tasks
.contains_key(sub_id)
.then(|| Arc::clone(&active_turn.turn_state))
})
}
/// Clear any pending waiters and input buffered for the current turn.
pub(crate) async fn clear_pending(&self, active_turn: &ActiveTurn) {
let mut turn_state = active_turn.turn_state.lock().await;
turn_state.clear_pending_waiters();
turn_state.pending_input.items.clear();
}
pub(crate) async fn defer_mailbox_delivery_to_next_turn(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
sub_id: &str,
) {
let turn_state = self.turn_state_for_sub_id(active_turn, sub_id).await;
let Some(turn_state) = turn_state else {
return;
};
let mut turn_state = turn_state.lock().await;
if !turn_state.pending_input.items.is_empty() {
return;
}
turn_state.set_mailbox_delivery_phase(MailboxDeliveryPhase::NextTurn);
}
pub(crate) async fn accept_mailbox_delivery_for_current_turn(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
sub_id: &str,
) {
let turn_state = self.turn_state_for_sub_id(active_turn, sub_id).await;
let Some(turn_state) = turn_state else {
return;
};
self.accept_mailbox_delivery_for_turn_state(turn_state.as_ref())
.await;
}
pub(super) async fn accept_mailbox_delivery_for_turn_state(
&self,
turn_state: &Mutex<TurnState>,
) {
turn_state
.lock()
.await
.accept_mailbox_delivery_for_current_turn();
}
pub(super) async fn push_pending_input_and_accept_mailbox_delivery_for_turn_state(
&self,
turn_state: &Mutex<TurnState>,
input: ResponseInputItem,
) {
let mut turn_state = turn_state.lock().await;
turn_state.pending_input.items.push(input);
turn_state.accept_mailbox_delivery_for_current_turn();
}
pub(crate) async fn extend_pending_input_for_turn_state(
&self,
turn_state: &Mutex<TurnState>,
input: Vec<ResponseInputItem>,
) {
turn_state.lock().await.pending_input.items.extend(input);
}
pub(crate) async fn take_pending_input_for_turn_state(
&self,
turn_state: &Mutex<TurnState>,
) -> Vec<ResponseInputItem> {
turn_state.lock().await.pending_input.items.split_off(0)
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub(crate) async fn inject_response_items(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
input: Vec<ResponseInputItem>,
) -> Result<(), Vec<ResponseInputItem>> {
let mut active = active_turn.lock().await;
match active.as_mut() {
Some(active_turn) => {
active_turn
.turn_state
.lock()
.await
.pending_input
.items
.extend(input);
Ok(())
}
None => Err(input),
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub(crate) async fn prepend_pending_input(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
mut input: Vec<ResponseInputItem>,
) -> Result<(), ()> {
let mut active = active_turn.lock().await;
match active.as_mut() {
Some(active_turn) => {
let mut turn_state = active_turn.turn_state.lock().await;
if !input.is_empty() {
let pending_input = &mut turn_state.pending_input;
input.append(&mut pending_input.items);
pending_input.items = input;
}
Ok(())
}
None => Err(()),
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub(crate) async fn get_pending_input(
&self,
active_turn: &Mutex<Option<ActiveTurn>>,
) -> Vec<ResponseInputItem> {
let (pending_input, accepts_mailbox_delivery) = {
let mut active = active_turn.lock().await;
match active.as_mut() {
Some(active_turn) => {
let mut turn_state = active_turn.turn_state.lock().await;
(
turn_state.pending_input.items.split_off(0),
turn_state.accepts_mailbox_delivery_for_current_turn(),
)
}
None => (Vec::new(), true),
}
};
if !accepts_mailbox_delivery {
return pending_input;
}
let mailbox_items = self.drain_mailbox_input_items().await;
if pending_input.is_empty() {
mailbox_items
} else if mailbox_items.is_empty() {
pending_input
} else {
let mut pending_input = pending_input;
pending_input.extend(mailbox_items);
pending_input
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state reads must remain atomic"
)]
pub(crate) async fn has_pending_input(&self, active_turn: &Mutex<Option<ActiveTurn>>) -> bool {
let (has_turn_pending_input, accepts_mailbox_delivery) = {
let active = active_turn.lock().await;
match active.as_ref() {
Some(active_turn) => {
let turn_state = active_turn.turn_state.lock().await;
(
!turn_state.pending_input.items.is_empty(),
turn_state.accepts_mailbox_delivery_for_current_turn(),
)
}
None => (false, true),
}
};
if has_turn_pending_input {
return true;
}
if !accepts_mailbox_delivery {
return false;
}
self.has_pending_mailbox_items().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use codex_protocol::AgentPath;
use pretty_assertions::assert_eq;
fn make_mail(
author: AgentPath,
recipient: AgentPath,
content: &str,
trigger_turn: bool,
) -> InterAgentCommunication {
InterAgentCommunication::new(
author,
recipient,
Vec::new(),
content.to_string(),
trigger_turn,
)
}
#[tokio::test]
async fn input_queue_notifies_mailbox_subscribers() {
let input_queue = InputQueue::new();
let mut mailbox_rx = input_queue.subscribe_mailbox().await;
input_queue
.enqueue_mailbox_communication(make_mail(
AgentPath::root(),
AgentPath::try_from("/root/worker").expect("agent path"),
"one",
/*trigger_turn*/ false,
))
.await;
input_queue
.enqueue_mailbox_communication(make_mail(
AgentPath::root(),
AgentPath::try_from("/root/worker").expect("agent path"),
"two",
/*trigger_turn*/ false,
))
.await;
mailbox_rx.changed().await.expect("mailbox update");
}
#[tokio::test]
async fn input_queue_drains_mailbox_in_delivery_order() {
let input_queue = InputQueue::new();
let mail_one = make_mail(
AgentPath::root(),
AgentPath::try_from("/root/worker").expect("agent path"),
"one",
/*trigger_turn*/ false,
);
let mail_two = make_mail(
AgentPath::try_from("/root/worker").expect("agent path"),
AgentPath::root(),
"two",
/*trigger_turn*/ false,
);
input_queue
.enqueue_mailbox_communication(mail_one.clone())
.await;
input_queue
.enqueue_mailbox_communication(mail_two.clone())
.await;
assert_eq!(
input_queue.drain_mailbox_input_items().await,
vec![
mail_one.to_response_input_item(),
mail_two.to_response_input_item()
]
);
assert!(!input_queue.has_pending_mailbox_items().await);
}
#[tokio::test]
async fn input_queue_tracks_pending_trigger_turn_mail() {
let input_queue = InputQueue::new();
input_queue
.enqueue_mailbox_communication(make_mail(
AgentPath::root(),
AgentPath::try_from("/root/worker").expect("agent path"),
"queued",
/*trigger_turn*/ false,
))
.await;
assert!(!input_queue.has_trigger_turn_mailbox_items().await);
input_queue
.enqueue_mailbox_communication(make_mail(
AgentPath::root(),
AgentPath::try_from("/root/worker").expect("agent path"),
"wake",
/*trigger_turn*/ true,
))
.await;
assert!(input_queue.has_trigger_turn_mailbox_items().await);
}
}
codex-rs/core/src/session/mod.rs
+22 -192
View File
@@ -10,8 +10,6 @@ use std::time::UNIX_EPOCH;
use crate::agent::AgentControl;
use crate::agent::AgentStatus;
use crate::agent::Mailbox;
use crate::agent::MailboxReceiver;
use crate::agent::agent_status_from_event;
use crate::agent::status::is_final;
use crate::attestation::AttestationProvider;
@@ -193,6 +191,7 @@ use codex_protocol::exec_output::StreamOutput;
mod config_lock;
mod handlers;
mod input_queue;
mod mcp;
mod multi_agents;
mod review;
@@ -206,6 +205,7 @@ use self::config_lock::validate_config_lock_if_configured;
#[cfg(test)]
use self::handlers::submission_dispatch_span;
use self::handlers::submission_loop;
pub(crate) use self::input_queue::TurnInputQueue;
use self::review::spawn_review_thread;
use self::session::AppServerClientMetadata;
use self::session::Session;
@@ -289,8 +289,6 @@ use crate::rollout::map_session_init_error;
use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
use crate::shell;
use crate::shell_snapshot::ShellSnapshot;
use crate::state::ActiveTurn;
use crate::state::MailboxDeliveryPhase;
use crate::state::PendingRequestPermissions;
use crate::state::SessionServices;
use crate::state::SessionState;
@@ -2947,12 +2945,14 @@ impl Session {
};
if let Some(token_info) = token_info.as_ref() {
for contributor in self.services.extensions.token_usage_contributors() {
contributor.on_token_usage(
&self.services.session_extension_data,
&self.services.thread_extension_data,
turn_context.extension_data.as_ref(),
token_info,
);
contributor
.on_token_usage(
&self.services.session_extension_data,
&self.services.thread_extension_data,
turn_context.extension_data.as_ref(),
token_info,
)
.await;
}
}
}
@@ -3020,16 +3020,6 @@ impl Session {
state.record_mcp_dependency_prompted(names);
}
pub async fn dependency_env(&self) -> HashMap<String, String> {
let state = self.state.lock().await;
state.dependency_env()
}
pub async fn set_dependency_env(&self, values: HashMap<String, String>) {
let mut state = self.state.lock().await;
state.set_dependency_env(values);
}
pub(crate) async fn set_server_reasoning_included(&self, included: bool) {
let mut state = self.state.lock().await;
state.set_server_reasoning_included(included);
@@ -3162,199 +3152,39 @@ impl Session {
.set_responsesapi_client_metadata(responsesapi_client_metadata);
}
let mut turn_state = active_turn.turn_state.lock().await;
turn_state.push_pending_input(input.into());
turn_state.accept_mailbox_delivery_for_current_turn();
self.input_queue
.push_pending_input_and_accept_mailbox_delivery_for_turn_state(
active_turn.turn_state.as_ref(),
input.into(),
)
.await;
Ok(active_turn_id.clone())
}
/// Returns the input if there was no task running to inject into.
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub async fn inject_response_items(
&self,
input: Vec<ResponseInputItem>,
) -> Result<(), Vec<ResponseInputItem>> {
let mut active = self.active_turn.lock().await;
match active.as_mut() {
Some(at) => {
let mut ts = at.turn_state.lock().await;
for item in input {
ts.push_pending_input(item);
}
Ok(())
}
None => Err(input),
}
}
pub(crate) async fn defer_mailbox_delivery_to_next_turn(&self, sub_id: &str) {
let turn_state = self.turn_state_for_sub_id(sub_id).await;
let Some(turn_state) = turn_state else {
return;
};
let mut turn_state = turn_state.lock().await;
if turn_state.has_pending_input() {
return;
}
turn_state.set_mailbox_delivery_phase(MailboxDeliveryPhase::NextTurn);
}
pub(crate) async fn accept_mailbox_delivery_for_current_turn(&self, sub_id: &str) {
let turn_state = self.turn_state_for_sub_id(sub_id).await;
let Some(turn_state) = turn_state else {
return;
};
turn_state
.lock()
self.input_queue
.inject_response_items(&self.active_turn, input)
.await
.set_mailbox_delivery_phase(MailboxDeliveryPhase::CurrentTurn);
}
pub(crate) async fn record_memory_citation_for_turn(&self, sub_id: &str) {
let turn_state = self.turn_state_for_sub_id(sub_id).await;
let turn_state = self
.input_queue
.turn_state_for_sub_id(&self.active_turn, sub_id)
.await;
let Some(turn_state) = turn_state else {
return;
};
turn_state.lock().await.has_memory_citation = true;
}
async fn turn_state_for_sub_id(
&self,
sub_id: &str,
) -> Option<Arc<tokio::sync::Mutex<crate::state::TurnState>>> {
let active = self.active_turn.lock().await;
active.as_ref().and_then(|active_turn| {
active_turn
.tasks
.contains_key(sub_id)
.then(|| Arc::clone(&active_turn.turn_state))
})
}
pub(crate) fn subscribe_mailbox_seq(&self) -> watch::Receiver<u64> {
self.mailbox.subscribe()
}
pub(crate) fn enqueue_mailbox_communication(&self, communication: InterAgentCommunication) {
self.mailbox.send(communication);
}
pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool {
self.mailbox_rx.lock().await.has_pending_trigger_turn()
}
pub(crate) async fn has_pending_mailbox_items(&self) -> bool {
self.mailbox_rx.lock().await.has_pending()
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub async fn prepend_pending_input(&self, input: Vec<ResponseInputItem>) -> Result<(), ()> {
let mut active = self.active_turn.lock().await;
match active.as_mut() {
Some(at) => {
let mut ts = at.turn_state.lock().await;
ts.prepend_pending_input(input);
Ok(())
}
None => Err(()),
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub async fn get_pending_input(&self) -> Vec<ResponseInputItem> {
let (pending_input, accepts_mailbox_delivery) = {
let mut active = self.active_turn.lock().await;
match active.as_mut() {
Some(at) => {
let mut ts = at.turn_state.lock().await;
(
ts.take_pending_input(),
ts.accepts_mailbox_delivery_for_current_turn(),
)
}
None => (Vec::new(), true),
}
};
if !accepts_mailbox_delivery {
return pending_input;
}
let mailbox_items = {
let mut mailbox_rx = self.mailbox_rx.lock().await;
mailbox_rx
.drain()
.into_iter()
.map(|mail| mail.to_response_input_item())
.collect::<Vec<_>>()
};
if pending_input.is_empty() {
mailbox_items
} else if mailbox_items.is_empty() {
pending_input
} else {
let mut pending_input = pending_input;
pending_input.extend(mailbox_items);
pending_input
}
}
/// Queue response items to be injected into the next active turn created for this session.
pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec<ResponseInputItem>) {
if items.is_empty() {
return;
}
let mut idle_pending_input = self.idle_pending_input.lock().await;
idle_pending_input.extend(items);
}
pub(crate) async fn take_queued_response_items_for_next_turn(&self) -> Vec<ResponseInputItem> {
std::mem::take(&mut *self.idle_pending_input.lock().await)
}
pub(crate) async fn has_queued_response_items_for_next_turn(&self) -> bool {
!self.idle_pending_input.lock().await.is_empty()
}
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state reads must remain atomic"
)]
pub async fn has_pending_input(&self) -> bool {
let (has_turn_pending_input, accepts_mailbox_delivery) = {
let active = self.active_turn.lock().await;
match active.as_ref() {
Some(at) => {
let ts = at.turn_state.lock().await;
(
ts.has_pending_input(),
ts.accepts_mailbox_delivery_for_current_turn(),
)
}
None => (false, true),
}
};
if has_turn_pending_input {
return true;
}
if !accepts_mailbox_delivery {
return false;
}
self.has_pending_mailbox_items().await
}
pub async fn interrupt_task(self: &Arc<Self>) {
info!("interrupt received: abort current task, if any");
let had_active_turn = self.active_turn.lock().await.is_some();
// Even without an active task, interrupt handling pauses any active goal.
self.abort_all_tasks(TurnAbortReason::Interrupted).await;
if !had_active_turn {
self.cancel_mcp_startup().await;
codex-rs/core/src/session/review.rs
+1
View File
@@ -56,6 +56,7 @@ pub(super) async fn spawn_review_thread(
.with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled)
.with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone())
.with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata)
.with_multi_agent_v2_tool_namespace(config.multi_agent_v2.tool_namespace.clone())
.with_multi_agent_v2_non_code_mode_only(config.multi_agent_v2.non_code_mode_only)
.with_goal_tools_allowed(goal_tools_supported)
.with_max_concurrent_threads_per_session(config.agent_max_threads)
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+8 -48
View File
@@ -59,7 +59,6 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ
let previous_model = "previous-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -74,11 +73,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item)];
@@ -101,7 +96,6 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif
let previous_model = "previous-rollout-model";
let mut previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -116,11 +110,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let turn_id = previous_context_item
.turn_id
@@ -925,7 +915,6 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
let previous_model = "previous-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -940,11 +929,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let previous_turn_id = previous_context_item
.turn_id
@@ -1005,7 +990,6 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
.expect("serialize seeded reference context item"),
serde_json::to_value(Some(TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -1020,11 +1004,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
}))
.expect("serialize expected reference context item")
);
@@ -1037,7 +1017,6 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu
let previous_model = "previous-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -1052,11 +1031,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let previous_turn_id = previous_context_item
.turn_id
@@ -1155,7 +1130,6 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo
let unmatched_abort_turn_id = "other-turn".to_string();
let current_context_item = TurnContextItem {
turn_id: Some(current_turn_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -1170,11 +1144,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let rollout_items = vec![
@@ -1272,7 +1242,6 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea
let previous_model = "previous-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -1287,11 +1256,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let previous_turn_id = previous_context_item
.turn_id
@@ -1428,7 +1393,6 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear
let previous_model = "previous-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -1443,11 +1407,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let previous_turn_id = previous_context_item
.turn_id
codex-rs/core/src/session/session.rs
+6 -9
View File
@@ -1,5 +1,7 @@
use super::input_queue::InputQueue;
use super::*;
use crate::goals::GoalRuntimeState;
use crate::state::ActiveTurn;
use codex_protocol::SessionId;
use codex_protocol::config_types::ServiceTier;
use codex_protocol::permissions::FileSystemPath;
@@ -27,9 +29,7 @@ pub(crate) struct Session {
pub(super) pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
pub(crate) conversation: Arc<RealtimeConversationManager>,
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
pub(super) mailbox: Mailbox,
pub(super) mailbox_rx: Mutex<MailboxReceiver>,
pub(super) idle_pending_input: Mutex<Vec<ResponseInputItem>>, // TODO (jif) merge with mailbox!
pub(crate) input_queue: InputQueue,
pub(crate) goal_runtime: GoalRuntimeState,
pub(crate) guardian_review_session: GuardianReviewSessionManager,
pub(crate) services: SessionServices,
@@ -877,7 +877,7 @@ impl Session {
config: config.as_ref(),
session_store: &session_extension_data,
thread_store: &thread_extension_data,
});
}).await;
}
let services = SessionServices {
@@ -950,7 +950,6 @@ impl Session {
let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) =
watch::channel(false);
let (mailbox, mailbox_rx) = Mailbox::new();
let sess = Arc::new(Session {
conversation_id: thread_id,
installation_id,
@@ -963,9 +962,7 @@ impl Session {
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
mailbox,
mailbox_rx: Mutex::new(mailbox_rx),
idle_pending_input: Mutex::new(Vec::new()),
input_queue: InputQueue::new(),
goal_runtime: GoalRuntimeState::new(),
guardian_review_session: GuardianReviewSessionManager::default(),
services,
@@ -1130,7 +1127,7 @@ impl Session {
};
// record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted.
sess.record_initial_history(initial_history).await;
Box::pin(sess.record_initial_history(initial_history)).await;
{
let mut state = sess.state.lock().await;
state.set_pending_session_start_source(Some(session_start_source));
codex-rs/core/src/session/tests.rs
+300 -66
View File
@@ -1874,8 +1874,9 @@ async fn record_token_usage_info_notifies_extension_contributors() {
records: Arc<std::sync::Mutex<Vec<RecordedTokenUsage>>>,
}
#[async_trait::async_trait]
impl codex_extension_api::TokenUsageContributor for TokenUsageRecorder {
fn on_token_usage(
async fn on_token_usage(
&self,
session_store: &codex_extension_api::ExtensionData,
thread_store: &codex_extension_api::ExtensionData,
@@ -2318,7 +2319,6 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() {
let previous_model = "forked-rollout-model";
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
trace_id: turn_context.trace_id.clone(),
#[allow(deprecated)]
cwd: turn_context.cwd.to_path_buf(),
current_date: turn_context.current_date.clone(),
@@ -2333,11 +2333,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() {
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
realtime_active: Some(turn_context.realtime_active),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy),
summary: codex_protocol::config_types::ReasoningSummary::Auto,
};
let turn_id = previous_context_item
.turn_id
@@ -4400,7 +4396,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
/*goal_tools_supported*/ true,
);
let (mailbox, mailbox_rx) = crate::agent::Mailbox::new();
let session = Session {
conversation_id: thread_id,
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
@@ -4413,9 +4408,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
mailbox,
mailbox_rx: Mutex::new(mailbox_rx),
idle_pending_input: Mutex::new(Vec::new()),
input_queue: super::input_queue::InputQueue::new(),
goal_runtime: crate::goals::GoalRuntimeState::new(),
guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
services,
@@ -5126,7 +5119,7 @@ async fn new_default_turn_captures_current_span_trace_id() {
&request_parent
));
let turn_context_item = async {
let turn_trace_id = async {
let expected_trace_id = Span::current()
.context()
.span()
@@ -5134,15 +5127,14 @@ async fn new_default_turn_captures_current_span_trace_id() {
.trace_id()
.to_string();
let turn_context = session.new_default_turn().await;
let turn_context_item = turn_context.to_turn_context_item();
assert_eq!(turn_context_item.trace_id, Some(expected_trace_id));
turn_context_item
assert_eq!(turn_context.trace_id, Some(expected_trace_id));
turn_context.trace_id.clone()
}
.instrument(request_span)
.await;
assert_eq!(
turn_context_item.trace_id.as_deref(),
turn_trace_id.as_deref(),
Some("00000000000000000000000000000011")
);
}
@@ -5668,8 +5660,9 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() {
expected_thread_id: ThreadId,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for ThreadStopRecorder {
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
@@ -5713,8 +5706,9 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif
expected_turn_id: String,
}
#[async_trait::async_trait]
impl codex_extension_api::ThreadLifecycleContributor<crate::config::Config> for LifecycleRecorder {
fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
@@ -5726,8 +5720,9 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif
}
}
#[async_trait::async_trait]
impl codex_extension_api::TurnLifecycleContributor for LifecycleRecorder {
fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
async fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) {
assert_eq!(
self.expected_thread_id.to_string(),
input.thread_store.level_id()
@@ -6239,7 +6234,6 @@ where
/*goal_tools_supported*/ true,
));
let (mailbox, mailbox_rx) = crate::agent::Mailbox::new();
let session = Arc::new(Session {
conversation_id: thread_id,
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
@@ -6252,9 +6246,7 @@ where
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
mailbox,
mailbox_rx: Mutex::new(mailbox_rx),
idle_pending_input: Mutex::new(Vec::new()),
input_queue: super::input_queue::InputQueue::new(),
goal_runtime: crate::goals::GoalRuntimeState::new(),
guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
services,
@@ -8085,7 +8077,7 @@ async fn steer_input_returns_active_turn_id() {
.expect("steering with matching expected turn id should succeed");
assert_eq!(turn_id, tc.sub_id);
assert!(sess.has_pending_input().await);
assert!(sess.input_queue.has_pending_input(&sess.active_turn).await);
}
#[tokio::test]
@@ -8131,7 +8123,7 @@ async fn prepend_pending_input_keeps_older_tail_ahead_of_newer_input() {
.await
.expect("inject initial pending input into active turn");
let drained = sess.get_pending_input().await;
let drained = sess.input_queue.get_pending_input(&sess.active_turn).await;
assert_eq!(drained, vec![blocked, later.clone()]);
sess.inject_response_items(vec![newer.clone()])
@@ -8140,11 +8132,15 @@ async fn prepend_pending_input_keeps_older_tail_ahead_of_newer_input() {
let mut drained_iter = drained.into_iter();
let _blocked = drained_iter.next().expect("blocked prompt should exist");
sess.prepend_pending_input(drained_iter.collect())
sess.input_queue
.prepend_pending_input(&sess.active_turn, drained_iter.collect())
.await
.expect("requeue later pending input at the front of the queue");
assert_eq!(sess.get_pending_input().await, vec![later, newer]);
assert_eq!(
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![later, newer]
);
}
#[tokio::test]
@@ -8158,7 +8154,8 @@ async fn queued_response_items_for_next_turn_move_into_next_active_turn() {
phase: None,
};
sess.queue_response_items_for_next_turn(vec![queued_item.clone()])
sess.input_queue
.queue_response_items_for_next_turn(vec![queued_item.clone()])
.await;
sess.spawn_task(
@@ -8171,7 +8168,10 @@ async fn queued_response_items_for_next_turn_move_into_next_active_turn() {
)
.await;
assert_eq!(sess.get_pending_input().await, vec![queued_item]);
assert_eq!(
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![queued_item]
);
}
#[tokio::test]
@@ -8185,13 +8185,18 @@ async fn idle_interrupt_does_not_wake_queued_next_turn_items() {
phase: None,
};
sess.queue_response_items_for_next_turn(vec![queued_item])
sess.input_queue
.queue_response_items_for_next_turn(vec![queued_item])
.await;
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
assert!(sess.active_turn.lock().await.is_none());
assert!(sess.has_queued_response_items_for_next_turn().await);
assert!(
sess.input_queue
.has_queued_response_items_for_next_turn()
.await
);
}
#[tokio::test]
@@ -8209,22 +8214,23 @@ async fn abort_empty_active_turn_preserves_pending_input() {
let active_turn = active.get_or_insert_with(ActiveTurn::default);
Arc::clone(&active_turn.turn_state)
};
turn_state
.lock()
.await
.push_pending_input(pending_item.clone());
sess.input_queue
.extend_pending_input_for_turn_state(turn_state.as_ref(), vec![pending_item.clone()])
.await;
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
assert!(sess.active_turn.lock().await.is_none());
assert_eq!(
turn_state.lock().await.take_pending_input(),
sess.input_queue
.take_pending_input_for_turn_state(turn_state.as_ref())
.await,
vec![pending_item]
);
}
#[tokio::test]
async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> {
async fn interrupt_accounts_active_goal_without_pausing() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
@@ -8254,7 +8260,7 @@ async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> {
.await?
.expect("goal should remain persisted after interrupt");
assert_eq!(
codex_protocol::protocol::ThreadGoalStatus::Paused,
codex_protocol::protocol::ThreadGoalStatus::Active,
goal.status
);
assert_eq!(70, goal.tokens_used);
@@ -8264,6 +8270,34 @@ async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test]
async fn shutdown_without_active_turn_keeps_active_goal_active() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
objective: Some("Keep improving the benchmark".to_string()),
status: None,
token_budget: None,
},
)
.await?;
assert!(sess.active_turn.lock().await.is_none());
assert!(handlers::shutdown(&sess, "shutdown".to_string()).await);
let goal = sess
.get_thread_goal()
.await?
.expect("goal should remain persisted after shutdown");
assert_eq!(
codex_protocol::protocol::ThreadGoalStatus::Active,
goal.status
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn active_goal_continuation_runs_again_after_no_tool_turn() -> anyhow::Result<()> {
let server = start_mock_server().await;
@@ -8554,7 +8588,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh
})
.await?;
let pending_input = sess.get_pending_input().await;
let pending_input = sess.input_queue.get_pending_input(&sess.active_turn).await;
let [ResponseInputItem::Message { role, content, .. }] = pending_input.as_slice() else {
panic!("expected one budget-limit steering message, got {pending_input:#?}");
};
@@ -8576,6 +8610,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted after accounting");
@@ -8599,6 +8634,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh
.await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted after follow-up accounting");
@@ -8610,6 +8646,57 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn usage_limit_runtime_stops_active_goal_and_prevents_idle_continuation() -> anyhow::Result<()>
{
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
objective: Some("Keep improving the benchmark".to_string()),
status: None,
token_budget: Some(Some(50)),
},
)
.await?;
sess.goal_runtime_apply(GoalRuntimeEvent::TurnStarted {
turn_context: tc.as_ref(),
token_usage: TokenUsage::default(),
})
.await?;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
NeverEndingTask {
kind: TaskKind::Regular,
listen_to_cancellation_token: false,
},
)
.await;
set_total_token_usage(&sess, post_goal_token_usage()).await;
sess.goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached {
turn_context: tc.as_ref(),
})
.await?;
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted after usage limiting");
assert_eq!(codex_state::ThreadGoalStatus::UsageLimited, goal.status);
assert_eq!(70, goal.tokens_used);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
sess.goal_runtime_apply(GoalRuntimeEvent::MaybeContinueIfIdle)
.await?;
assert!(sess.active_turn.lock().await.is_none());
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
@@ -8638,6 +8725,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
@@ -8646,6 +8734,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a
let previous_goal = goal.clone();
let goal_id = goal.goal_id.clone();
let updated_goal = state_db
.thread_goals()
.update_thread_goal(
sess.conversation_id,
codex_state::ThreadGoalUpdate {
@@ -8667,6 +8756,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a
assert!(sess.active_turn.lock().await.is_some());
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
@@ -8693,6 +8783,7 @@ async fn external_objective_change_steers_active_turn() -> anyhow::Result<()> {
let state_db = goal_test_state_db(sess.as_ref()).await?;
let old_goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Keep improving the benchmark",
@@ -8701,6 +8792,7 @@ async fn external_objective_change_steers_active_turn() -> anyhow::Result<()> {
)
.await?;
let new_goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Write a concise benchmark summary",
@@ -8717,7 +8809,7 @@ async fn external_objective_change_steers_active_turn() -> anyhow::Result<()> {
})
.await?;
let pending_input = sess.get_pending_input().await;
let pending_input = sess.input_queue.get_pending_input(&sess.active_turn).await;
assert!(
pending_input.iter().any(|item| {
matches!(
@@ -8758,6 +8850,7 @@ async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow:
let state_db = goal_test_state_db(sess.as_ref()).await?;
let goal = state_db
.thread_goals()
.replace_thread_goal(
sess.conversation_id,
"Keep improving the benchmark",
@@ -8791,6 +8884,7 @@ async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow:
.await?;
let goal = state_db
.thread_goals()
.get_thread_goal(sess.conversation_id)
.await?
.expect("goal should remain persisted");
@@ -8890,6 +8984,7 @@ async fn completed_goal_accounts_current_turn_tokens_before_tool_response() -> a
)
.await?;
let persisted_goal = state_db
.thread_goals()
.get_thread_goal(test.session_configured.thread_id)
.await?
.expect("goal should be persisted");
@@ -8922,19 +9017,26 @@ async fn queue_only_mailbox_mail_waits_for_next_turn_after_answer_boundary() {
)
.await;
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.enqueue_mailbox_communication(communication.clone());
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
sess.input_queue
.enqueue_mailbox_communication(communication.clone())
.await;
assert!(
!sess.has_pending_input().await,
!sess.input_queue.has_pending_input(&sess.active_turn).await,
"queue-only mailbox mail should stay buffered once the current turn emitted its answer"
);
assert_eq!(sess.get_pending_input().await, Vec::new());
assert_eq!(
sess.input_queue.get_pending_input(&sess.active_turn).await,
Vec::new()
);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
assert_eq!(
sess.get_pending_input().await,
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![communication.to_response_input_item()],
);
}
@@ -8952,23 +9054,27 @@ async fn trigger_turn_mailbox_mail_waits_for_next_turn_after_answer_boundary() {
)
.await;
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.enqueue_mailbox_communication(InterAgentCommunication::new(
AgentPath::try_from("/root/worker").expect("worker path should parse"),
AgentPath::root(),
Vec::new(),
"late trigger update".to_string(),
/*trigger_turn*/ true,
));
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
sess.input_queue
.enqueue_mailbox_communication(InterAgentCommunication::new(
AgentPath::try_from("/root/worker").expect("worker path should parse"),
AgentPath::root(),
Vec::new(),
"late trigger update".to_string(),
/*trigger_turn*/ true,
))
.await;
assert!(
!sess.has_pending_input().await,
!sess.input_queue.has_pending_input(&sess.active_turn).await,
"trigger-turn mailbox mail should not extend the current turn after its answer boundary"
);
sess.abort_all_tasks(TurnAbortReason::Replaced).await;
assert!(sess.has_trigger_turn_mailbox_items().await);
assert!(sess.input_queue.has_trigger_turn_mailbox_items().await);
}
#[tokio::test]
@@ -8991,8 +9097,12 @@ async fn steered_input_reopens_mailbox_delivery_for_current_turn() {
)
.await;
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.enqueue_mailbox_communication(communication.clone());
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
sess.input_queue
.enqueue_mailbox_communication(communication.clone())
.await;
sess.steer_input(
vec![UserInput::Text {
text: "follow up".to_string(),
@@ -9005,7 +9115,7 @@ async fn steered_input_reopens_mailbox_delivery_for_current_turn() {
.expect("steered input should be accepted");
assert_eq!(
sess.get_pending_input().await,
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![
ResponseInputItem::from(vec![UserInput::Text {
text: "follow up".to_string(),
@@ -9036,8 +9146,12 @@ async fn stale_defer_mailbox_delivery_does_not_override_steered_input() {
)
.await;
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.enqueue_mailbox_communication(communication.clone());
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
sess.input_queue
.enqueue_mailbox_communication(communication.clone())
.await;
sess.steer_input(
vec![UserInput::Text {
text: "follow up".to_string(),
@@ -9049,10 +9163,12 @@ async fn stale_defer_mailbox_delivery_does_not_override_steered_input() {
.await
.expect("steered input should be accepted");
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
assert_eq!(
sess.get_pending_input().await,
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![
ResponseInputItem::from(vec![UserInput::Text {
text: "follow up".to_string(),
@@ -9083,8 +9199,12 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() {
)
.await;
sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
sess.enqueue_mailbox_communication(communication.clone());
sess.input_queue
.defer_mailbox_delivery_to_next_turn(&sess.active_turn, &tc.sub_id)
.await;
sess.input_queue
.enqueue_mailbox_communication(communication.clone())
.await;
let item = ResponseItem::FunctionCall {
id: None,
@@ -9108,7 +9228,7 @@ async fn tool_calls_reopen_mailbox_delivery_for_current_turn() {
assert!(output.needs_follow_up);
assert!(output.tool_future.is_some());
assert_eq!(
sess.get_pending_input().await,
sess.input_queue.get_pending_input(&sess.active_turn).await,
vec![communication.to_response_input_item()],
);
}
@@ -9524,7 +9644,121 @@ async fn update_goal_tool_rejects_pausing_goal() {
};
assert_eq!(
output,
"update_goal can only mark the existing goal complete; pause, resume, and budget-limited status changes are controlled by the user or system"
"update_goal can only mark the existing goal complete or blocked; pause, resume, budget-limited, and usage-limited status changes are controlled by the user or system"
);
let goal = session
.get_thread_goal()
.await
.expect("read thread goal")
.expect("goal should still exist");
assert_eq!(goal.status, ThreadGoalStatus::Active);
}
#[tokio::test]
async fn update_goal_tool_marks_goal_blocked() {
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let create_handler = CreateGoalHandler;
let update_handler = UpdateGoalHandler;
create_handler
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
cancellation_token: CancellationToken::new(),
tracker: Arc::clone(&tracker),
call_id: "create-goal".to_string(),
tool_name: codex_tools::ToolName::plain("create_goal"),
source: ToolCallSource::Direct,
payload: ToolPayload::Function {
arguments: serde_json::json!({
"objective": "Keep the watcher alive",
"token_budget": 123,
})
.to_string(),
},
})
.await
.expect("initial create_goal should succeed");
update_handler
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
cancellation_token: CancellationToken::new(),
tracker,
call_id: "block-goal".to_string(),
tool_name: codex_tools::ToolName::plain("update_goal"),
source: ToolCallSource::Direct,
payload: ToolPayload::Function {
arguments: serde_json::json!({
"status": "blocked",
})
.to_string(),
},
})
.await
.expect("update_goal should mark the goal blocked");
let goal = session
.get_thread_goal()
.await
.expect("read thread goal")
.expect("goal should still exist");
assert_eq!(goal.status, ThreadGoalStatus::Blocked);
}
#[tokio::test]
async fn update_goal_tool_rejects_usage_limited_goal() {
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let create_handler = CreateGoalHandler;
let update_handler = UpdateGoalHandler;
create_handler
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
cancellation_token: CancellationToken::new(),
tracker: Arc::clone(&tracker),
call_id: "create-goal".to_string(),
tool_name: codex_tools::ToolName::plain("create_goal"),
source: ToolCallSource::Direct,
payload: ToolPayload::Function {
arguments: serde_json::json!({
"objective": "Keep the watcher alive",
})
.to_string(),
},
})
.await
.expect("initial create_goal should succeed");
let response = update_handler
.handle(ToolInvocation {
session: Arc::clone(&session),
turn: Arc::clone(&turn_context),
cancellation_token: CancellationToken::new(),
tracker,
call_id: "usage-limit-goal".to_string(),
tool_name: codex_tools::ToolName::plain("update_goal"),
source: ToolCallSource::Direct,
payload: ToolPayload::Function {
arguments: serde_json::json!({
"status": "usageLimited",
})
.to_string(),
},
})
.await;
let Err(FunctionCallError::RespondToModel(output)) = response else {
panic!("expected update_goal to reject usage-limiting a goal");
};
assert_eq!(
output,
"update_goal can only mark the existing goal complete or blocked; pause, resume, budget-limited, and usage-limited status changes are controlled by the user or system"
);
let goal = session
codex-rs/core/src/session/turn.rs
+72 -59
View File
@@ -9,7 +9,6 @@ use crate::build_skill_injections;
use crate::client::ModelClientSession;
use crate::client_common::Prompt;
use crate::client_common::ResponseEvent;
use crate::collect_env_var_dependencies;
use crate::collect_explicit_skill_mentions;
use crate::compact::InitialContextInjection;
use crate::compact::collect_user_messages;
@@ -20,6 +19,7 @@ use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_
use crate::connectors;
use crate::context::ContextualUserFragment;
use crate::feedback_tags;
use crate::goals::GoalRuntimeEvent;
use crate::hook_runtime::PendingInputHookDisposition;
use crate::hook_runtime::emit_hook_completed_events;
use crate::hook_runtime::inspect_pending_input;
@@ -39,7 +39,6 @@ use crate::mentions::collect_explicit_plugin_mentions;
use crate::mentions::collect_tool_mentions_from_messages;
use crate::parse_turn_item;
use crate::plugins::build_plugin_injections;
use crate::resolve_skill_dependencies_for_turn;
use crate::session::PreviousTurnSettings;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
@@ -71,6 +70,7 @@ use codex_analytics::build_track_events_context;
use codex_async_utils::OrCancelExt;
use codex_features::Feature;
use codex_git_utils::get_git_repo_root;
use codex_git_utils::get_git_repo_root_with_fs;
use codex_hooks::HookEvent;
use codex_hooks::HookEventAfterAgent;
use codex_hooks::HookPayload;
@@ -146,7 +146,7 @@ pub(crate) async fn run_turn(
prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
) -> Option<String> {
if input.is_empty() && !sess.has_pending_input().await {
if input.is_empty() && !sess.input_queue.has_pending_input(&sess.active_turn).await {
return None;
}
@@ -158,16 +158,18 @@ pub(crate) async fn run_turn(
// new user message are recorded. Estimate pending incoming items (context
// diffs/full reinjection + user input) and trigger compaction preemptively
// when they would push the thread over the compaction threshold.
let pre_sampling_compact =
match run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await {
Ok(pre_sampling_compact) => pre_sampling_compact,
Err(_) => {
error!("Failed to run pre-sampling compact");
return None;
}
};
if pre_sampling_compact.reset_client_session {
client_session.reset_websocket_session();
if let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await {
if err.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded
&& let Err(err) = sess
.goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached {
turn_context: turn_context.as_ref(),
})
.await
{
warn!("failed to usage-limit active goal after usage-limit error: {err}");
}
error!("Failed to run pre-sampling compact");
return None;
}
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
@@ -230,15 +232,6 @@ pub(crate) async fn run_turn(
&connector_slug_counts,
)
});
let config = turn_context.config.clone();
if config
.features
.enabled(Feature::SkillEnvVarDependencyPrompt)
{
let env_var_dependencies = collect_env_var_dependencies(&mentioned_skills);
resolve_skill_dependencies_for_turn(&sess, &turn_context, &env_var_dependencies).await;
}
maybe_prompt_and_install_mcp_dependencies(
sess.as_ref(),
turn_context.as_ref(),
@@ -370,8 +363,17 @@ pub(crate) async fn run_turn(
// Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains
// many turns, from the perspective of the user, it is a single turn.
#[allow(deprecated)]
let display_root = get_git_repo_root(turn_context.cwd.as_path())
.unwrap_or_else(|| turn_context.cwd.clone().into_path_buf());
let display_root = match turn_context.environments.primary() {
Some(turn_environment) => get_git_repo_root_with_fs(
turn_environment.environment.get_filesystem().as_ref(),
&turn_environment.cwd,
)
.await
.unwrap_or_else(|| turn_environment.cwd.clone())
.into_path_buf(),
None => get_git_repo_root(turn_context.cwd.as_path())
.unwrap_or_else(|| turn_context.cwd.clone().into_path_buf()),
};
let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::with_display_root(
display_root,
)));
@@ -393,7 +395,7 @@ pub(crate) async fn run_turn(
// submitted through the UI while the model was running. Though the UI
// may support this, the model might not.
let pending_input = if can_drain_pending_input {
sess.get_pending_input().await
sess.input_queue.get_pending_input(&sess.active_turn).await
} else {
Vec::new()
};
@@ -414,7 +416,10 @@ pub(crate) async fn run_turn(
} => {
let remaining_pending_input = pending_input_iter.collect::<Vec<_>>();
if !remaining_pending_input.is_empty() {
let _ = sess.prepend_pending_input(remaining_pending_input).await;
let _ = sess
.input_queue
.prepend_pending_input(&sess.active_turn, remaining_pending_input)
.await;
requeued_pending_input = true;
}
blocked_pending_input_contexts = additional_contexts;
@@ -474,7 +479,7 @@ pub(crate) async fn run_turn(
last_agent_message: sampling_request_last_agent_message,
} = sampling_request_output;
can_drain_pending_input = true;
let has_pending_input = sess.has_pending_input().await;
let has_pending_input = sess.input_queue.has_pending_input(&sess.active_turn).await;
let needs_follow_up = model_needs_follow_up || has_pending_input;
let total_usage_tokens = sess.get_total_token_usage().await;
let token_limit_reached = total_usage_tokens >= auto_compact_limit;
@@ -496,7 +501,7 @@ pub(crate) async fn run_turn(
// as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop.
if token_limit_reached && needs_follow_up {
let reset_client_session = match run_auto_compact(
match run_auto_compact(
&sess,
&turn_context,
&mut client_session,
@@ -506,12 +511,22 @@ pub(crate) async fn run_turn(
)
.await
{
Ok(reset_client_session) => reset_client_session,
Err(_) => return None,
Ok(()) => {}
Err(err) => {
if err.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded
&& let Err(err) = sess
.goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached {
turn_context: turn_context.as_ref(),
})
.await
{
warn!(
"failed to usage-limit active goal after usage-limit error: {err}"
);
}
return None;
}
};
if reset_client_session {
client_session.reset_websocket_session();
}
can_drain_pending_input = !model_needs_follow_up;
continue;
}
@@ -663,6 +678,15 @@ pub(crate) async fn run_turn(
}
Err(e) => {
info!("Turn error: {e:#}");
if e.to_codex_protocol_error() == CodexErrorInfo::UsageLimitExceeded
&& let Err(err) = sess
.goal_runtime_apply(GoalRuntimeEvent::UsageLimitReached {
turn_context: turn_context.as_ref(),
})
.await
{
warn!("failed to usage-limit active goal after usage-limit error: {err}");
}
let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None));
sess.send_event(&turn_context, event).await;
// let the user continue the conversation
@@ -722,24 +746,19 @@ async fn track_turn_resolved_config_analytics(
});
}
struct PreSamplingCompactResult {
reset_client_session: bool,
}
async fn run_pre_sampling_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
client_session: &mut ModelClientSession,
) -> CodexResult<PreSamplingCompactResult> {
) -> CodexResult<()> {
let total_usage_tokens_before_compaction = sess.get_total_token_usage().await;
let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact(
maybe_run_previous_model_inline_compact(
sess,
turn_context,
client_session,
total_usage_tokens_before_compaction,
)
.await?;
let mut reset_client_session = pre_sampling_compacted;
let total_usage_tokens = sess.get_total_token_usage().await;
let auto_compact_limit = turn_context
.model_info
@@ -747,7 +766,7 @@ async fn run_pre_sampling_compact(
.unwrap_or(i64::MAX);
// Compact if the total usage tokens are greater than the auto compact limit
if total_usage_tokens >= auto_compact_limit {
reset_client_session |= run_auto_compact(
run_auto_compact(
sess,
turn_context,
client_session,
@@ -756,27 +775,22 @@ async fn run_pre_sampling_compact(
CompactionPhase::PreTurn,
)
.await?;
pre_sampling_compacted = true;
}
Ok(PreSamplingCompactResult {
reset_client_session: pre_sampling_compacted && reset_client_session,
})
Ok(())
}
/// Runs pre-sampling compaction against the previous model when switching to a smaller
/// context-window model.
///
/// Returns `Ok(true)` when compaction ran successfully, `Ok(false)` when compaction was skipped
/// because the model/context-window preconditions were not met, and `Err(_)` only when compaction
/// was attempted and failed.
/// Returns `Err(_)` only when compaction was attempted and failed.
async fn maybe_run_previous_model_inline_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
client_session: &mut ModelClientSession,
total_usage_tokens: i64,
) -> CodexResult<bool> {
) -> CodexResult<()> {
let Some(previous_turn_settings) = sess.previous_turn_settings().await else {
return Ok(false);
return Ok(());
};
let previous_model_turn_context = Arc::new(
turn_context
@@ -785,10 +799,10 @@ async fn maybe_run_previous_model_inline_compact(
);
let Some(old_context_window) = previous_model_turn_context.model_context_window() else {
return Ok(false);
return Ok(());
};
let Some(new_context_window) = turn_context.model_context_window() else {
return Ok(false);
return Ok(());
};
let new_auto_compact_limit = turn_context
.model_info
@@ -798,7 +812,7 @@ async fn maybe_run_previous_model_inline_compact(
&& previous_model_turn_context.model_info.slug != turn_context.model_info.slug
&& old_context_window > new_context_window;
if should_run {
let _ = run_auto_compact(
run_auto_compact(
sess,
&previous_model_turn_context,
client_session,
@@ -807,9 +821,8 @@ async fn maybe_run_previous_model_inline_compact(
CompactionPhase::PreTurn,
)
.await?;
return Ok(true);
}
Ok(false)
Ok(())
}
async fn run_auto_compact(
@@ -819,7 +832,7 @@ async fn run_auto_compact(
initial_context_injection: InitialContextInjection,
reason: CompactionReason,
phase: CompactionPhase,
) -> CodexResult<bool> {
) -> CodexResult<()> {
if should_use_remote_compact_task(turn_context.provider.info()) {
if turn_context.features.enabled(Feature::RemoteCompactionV2) {
run_inline_remote_auto_compact_task_v2(
@@ -831,7 +844,7 @@ async fn run_auto_compact(
phase,
)
.await?;
return Ok(false);
return Ok(());
}
run_inline_remote_auto_compact_task(
Arc::clone(sess),
@@ -851,7 +864,7 @@ async fn run_auto_compact(
)
.await?;
}
Ok(true)
Ok(())
}
pub(super) fn collect_explicit_app_ids_from_skill_items(
@@ -2023,7 +2036,7 @@ async fn try_run_sampling_request(
}
needs_follow_up |= output_result.needs_follow_up;
// todo: remove before stabilizing multi-agent v2
if preempt_for_mailbox_mail && sess.mailbox_rx.lock().await.has_pending() {
if preempt_for_mailbox_mail && sess.input_queue.has_pending_mailbox_items().await {
break Ok(SamplingRequestResult {
needs_follow_up: true,
last_agent_message,
codex-rs/core/src/session/turn_context.rs
+3 -6
View File
@@ -220,6 +220,7 @@ impl TurnContext {
.with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled)
.with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone())
.with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata)
.with_multi_agent_v2_tool_namespace(config.multi_agent_v2.tool_namespace.clone())
.with_multi_agent_v2_non_code_mode_only(config.multi_agent_v2.non_code_mode_only)
.with_goal_tools_allowed(self.tools_config.goal_tools)
.with_max_concurrent_threads_per_session(
@@ -367,7 +368,6 @@ impl TurnContext {
pub(crate) fn to_turn_context_item(&self) -> TurnContextItem {
TurnContextItem {
turn_id: Some(self.sub_id.clone()),
trace_id: self.trace_id.clone(),
#[allow(deprecated)]
cwd: self.cwd.to_path_buf(),
current_date: self.current_date.clone(),
@@ -382,11 +382,7 @@ impl TurnContext {
collaboration_mode: Some(self.collaboration_mode.clone()),
realtime_active: Some(self.realtime_active),
effort: self.reasoning_effort,
summary: self.reasoning_summary,
user_instructions: self.user_instructions.clone(),
developer_instructions: self.developer_instructions.clone(),
final_output_json_schema: self.final_output_json_schema.clone(),
truncation_policy: Some(self.truncation_policy),
summary: ReasoningSummaryConfig::Auto,
}
}
@@ -539,6 +535,7 @@ impl Session {
.with_spawn_agent_usage_hint(per_turn_config.multi_agent_v2.usage_hint_enabled)
.with_spawn_agent_usage_hint_text(per_turn_config.multi_agent_v2.usage_hint_text.clone())
.with_hide_spawn_agent_metadata(per_turn_config.multi_agent_v2.hide_spawn_agent_metadata)
.with_multi_agent_v2_tool_namespace(per_turn_config.multi_agent_v2.tool_namespace.clone())
.with_multi_agent_v2_non_code_mode_only(per_turn_config.multi_agent_v2.non_code_mode_only)
.with_goal_tools_allowed(goal_tools_supported)
.with_max_concurrent_threads_per_session(