openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-25 18:28:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

app-server: reserve running resumes before idle unload

Browse Source
This commit is contained in:
Winston Howes
2026-07-07 03:22:42 -07:00
parent 69309d111e
commit 0ce1cd8c76
5 changed files with 415 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors.rs
+2
View File
@@ -545,6 +545,8 @@ use crate::error_code::invalid_request;
use crate::filters::compute_source_filters;
use crate::filters::source_kind_matches;
use crate::thread_state::ConnectionCapabilities;
use crate::thread_state::ConnectionReservationAction;
use crate::thread_state::ConnectionReservationOutcome;
use crate::thread_state::ThreadListenerCommand;
use crate::thread_state::ThreadListenerRegistration;
use crate::thread_state::ThreadState;
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
+74 -38
View File
@@ -120,6 +120,51 @@ impl UnloadingState {
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "idle unload admission holds both lifecycle guards through its final rechecks"
)]
async fn claim_idle_unload_if_permitted(
unloading_state: &mut UnloadingState,
thread_list_state_permit: Arc<Semaphore>,
pending_thread_unloads: &Arc<Mutex<HashSet<ThreadId>>>,
thread_state_manager: &ThreadStateManager,
conversation_id: ThreadId,
is_running: impl std::future::Future<Output = bool>,
) -> bool {
let Ok(permit) = thread_list_state_permit.acquire_owned().await else {
return false;
};
let mut pending = pending_thread_unloads.lock().await;
if pending.contains(&conversation_id) {
return false;
}
if !unloading_state.should_unload_now() {
return false;
}
if is_running.await {
unloading_state.note_thread_activity_observed();
return false;
}
if thread_state_manager
.has_connections_or_reservations(conversation_id)
.await
{
return false;
}
if !unloading_state.should_unload_now() {
return false;
}
pending.insert(conversation_id);
drop(pending);
drop(permit);
true
}
#[cfg(test)]
#[path = "thread_lifecycle_tests.rs"]
mod tests;
pub(super) enum ThreadShutdownResult {
Complete,
SubmitFailed,
@@ -384,15 +429,19 @@ pub(super) async fn ensure_listener_task_running(
unloading_state.note_thread_activity_observed();
continue;
}
if !claim_idle_unload_if_permitted(
&mut unloading_state,
thread_list_state_permit.clone(),
&pending_thread_unloads,
&thread_state_manager,
conversation_id,
async {
matches!(conversation.agent_status().await, AgentStatus::Running)
},
)
.await
{
let mut pending_thread_unloads = pending_thread_unloads.lock().await;
if pending_thread_unloads.contains(&conversation_id) {
continue;
}
if !unloading_state.should_unload_now() {
continue;
}
pending_thread_unloads.insert(conversation_id);
continue;
}
unload_thread_without_subscribers(
thread_manager.clone(),
@@ -519,7 +568,6 @@ pub(super) async fn handle_thread_listener_command(
conversation_id,
conversation,
codex_home,
thread_state_manager,
thread_state,
thread_watch_manager,
outgoing,
@@ -568,15 +616,10 @@ pub(super) async fn handle_thread_listener_command(
}
#[allow(clippy::too_many_arguments)]
#[expect(
clippy::await_holding_invalid_type,
reason = "running-thread resume subscription must be serialized against pending unloads"
)]
pub(super) async fn handle_pending_thread_resume_request(
conversation_id: ThreadId,
conversation: &Arc<CodexThread>,
_codex_home: &Path,
thread_state_manager: &ThreadStateManager,
thread_state: &Arc<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
outgoing: &Arc<OutgoingMessageSender>,
@@ -602,6 +645,7 @@ pub(super) async fn handle_pending_thread_resume_request(
.is_some_and(|turn| matches!(turn.status, TurnStatus::InProgress));
let request_id = pending.request_id;
let resume_handled_tx = pending.resume_handled_tx;
let connection_id = request_id.connection_id;
let mut thread = pending.thread_summary;
if pending.include_turns {
@@ -633,6 +677,7 @@ pub(super) async fn handle_pending_thread_resume_request(
Ok(page) => Some(page),
Err(error) => {
outgoing.send_error(request_id, error).await;
let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel);
return;
}
}
@@ -646,31 +691,21 @@ pub(super) async fn handle_pending_thread_resume_request(
}
}
if pending_thread_unloads
.lock()
.await
.contains(&conversation_id)
{
let pending_thread_unloads = pending_thread_unloads.lock().await;
if pending_thread_unloads.contains(&conversation_id) {
drop(pending_thread_unloads);
outgoing
.send_error(
request_id,
invalid_request(format!(
"thread {conversation_id} is closing; retry thread/resume after the thread is closed"
)),
)
.await;
return;
}
if !thread_state_manager
.try_add_connection_to_thread(conversation_id, connection_id)
.await
{
tracing::debug!(
thread_id = %conversation_id,
connection_id = ?connection_id,
"skipping running thread resume for closed connection"
);
return;
}
outgoing
.send_error(
request_id,
invalid_request(format!(
"thread {conversation_id} is closing; retry thread/resume after the thread is closed"
)),
)
.await;
let _ = resume_handled_tx.send(ConnectionReservationAction::Cancel);
return;
}
let config_snapshot = pending.config_snapshot;
@@ -714,6 +749,7 @@ pub(super) async fn handle_pending_thread_resume_request(
outgoing
.send_response_with_thread_originator(request_id, response, originator)
.await;
let _ = resume_handled_tx.send(ConnectionReservationAction::Commit);
// Match cold resume: metadata-only resume should attach the listener without
// paying the cost of turn reconstruction for historical usage replay.
if let Some(token_usage_thread) = token_usage_thread {
codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs
+135
View File
@@ -0,0 +1,135 @@
use super::*;
use pretty_assertions::assert_eq;
#[tokio::test]
async fn resume_reservation_prevents_idle_unload_before_listener_commit() {
let now = Instant::now();
let (_has_subscribers_tx, has_subscribers_rx) = watch::channel(false);
let (_thread_status_tx, thread_status_rx) = watch::channel(ThreadStatus::Idle);
let mut unloading_state = UnloadingState {
delay: Duration::ZERO,
has_subscribers_rx,
has_subscribers: (false, now),
thread_status_rx,
is_active: (false, now),
};
let pending = Arc::new(Mutex::new(HashSet::new()));
let thread_state_manager = ThreadStateManager::new();
let thread_id = ThreadId::new();
let connection_id = ConnectionId(1);
thread_state_manager
.connection_initialized(connection_id, ConnectionCapabilities::default())
.await;
let permit = Arc::new(Semaphore::new(1));
let resume_permit = permit
.clone()
.acquire_owned()
.await
.expect("permit should be open");
let (resume_handled_tx, resume_outcome_rx) = thread_state_manager
.start_connection_reservation(thread_id, connection_id)
.await
.expect("live connection should reserve the thread");
drop(resume_permit);
assert!(
!claim_idle_unload_if_permitted(
&mut unloading_state,
permit,
&pending,
&thread_state_manager,
thread_id,
std::future::ready(false),
)
.await
);
assert!(pending.lock().await.is_empty());
assert!(
thread_state_manager
.unsubscribe_connection_from_thread(thread_id, connection_id)
.await
);
let (replacement_tx, replacement_outcome_rx) = thread_state_manager
.start_connection_reservation(thread_id, connection_id)
.await
.expect("replacement reservation should succeed");
resume_handled_tx
.send(ConnectionReservationAction::Commit)
.expect("reservation owner should still be running");
assert_eq!(
resume_outcome_rx
.await
.expect("reservation owner should report an outcome"),
ConnectionReservationOutcome::Handled,
);
assert!(
thread_state_manager
.subscribed_connection_ids(thread_id)
.await
.is_empty()
);
assert!(
thread_state_manager
.has_connections_or_reservations(thread_id)
.await
);
drop(replacement_tx);
assert_eq!(
replacement_outcome_rx
.await
.expect("reservation owner should report abandonment"),
ConnectionReservationOutcome::Abandoned,
);
assert!(
!thread_state_manager
.has_connections_or_reservations(thread_id)
.await
);
assert!(
thread_state_manager
.subscribed_connection_ids(thread_id)
.await
.is_empty()
);
thread_state_manager.clear_all_listeners().await;
}
#[tokio::test]
async fn pending_resume_prevents_disconnect_from_reporting_thread_stale() {
let thread_state_manager = ThreadStateManager::new();
let thread_id = ThreadId::new();
let committed_connection_id = ConnectionId(1);
let pending_connection_id = ConnectionId(2);
for connection_id in [committed_connection_id, pending_connection_id] {
thread_state_manager
.connection_initialized(connection_id, ConnectionCapabilities::default())
.await;
}
assert!(
thread_state_manager
.try_add_connection_to_thread(thread_id, committed_connection_id)
.await
);
let (reservation_tx, reservation_outcome_rx) = thread_state_manager
.start_connection_reservation(thread_id, pending_connection_id)
.await
.expect("live connection should reserve the thread");
assert!(
thread_state_manager
.remove_connection(committed_connection_id)
.await
.is_empty(),
"a pending reservation must keep the thread from being reported stale"
);
drop(reservation_tx);
assert_eq!(
reservation_outcome_rx
.await
.expect("reservation owner should report abandonment"),
ConnectionReservationOutcome::Abandoned,
);
thread_state_manager.clear_all_listeners().await;
}
codex-rs/app-server/src/request_processors/thread_processor.rs
+36 -2
View File
@@ -371,6 +371,12 @@ pub(crate) struct ThreadRequestProcessor {
enum RunningThreadResumeResult {
/// The request was delegated to the loaded thread.
Handled,
/// The listener owns the response after the caller releases the global
/// thread-list permit.
ListenerPending {
thread_id: ThreadId,
outcome_rx: oneshot::Receiver<ConnectionReservationOutcome>,
},
/// No loaded thread handled the request.
///
/// The optional stored thread contains the history-bearing probe that cold
@@ -2672,7 +2678,7 @@ impl ThreadRequestProcessor {
let redact_resume_payloads =
should_redact_thread_resume_payloads(app_server_client_name.as_deref());
let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await {
let thread_list_state_permit = match self.acquire_thread_list_state_permit().await {
Ok(permit) => permit,
Err(error) => {
self.outgoing.send_error(request_id, error).await;
@@ -2689,6 +2695,23 @@ impl ThreadRequestProcessor {
.await
{
Ok(RunningThreadResumeResult::Handled) => return Ok(()),
Ok(RunningThreadResumeResult::ListenerPending {
thread_id,
outcome_rx,
}) => {
drop(thread_list_state_permit);
if !matches!(outcome_rx.await, Ok(ConnectionReservationOutcome::Handled)) {
self.outgoing
.send_error(
request_id,
internal_error(format!(
"failed to subscribe running thread resume for thread {thread_id}: thread listener closed before handling the request"
)),
)
.await;
}
return Ok(());
}
Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread,
Err(error) => {
self.outgoing.send_error(request_id, error).await;
@@ -3131,10 +3154,18 @@ impl ThreadRequestProcessor {
.thread_goal_processor
.pending_resume_goal_state(existing_thread.as_ref())
.await;
let Some((resume_handled_tx, outcome_rx)) = self
.thread_state_manager
.start_connection_reservation(existing_thread_id, request_id.connection_id)
.await
else {
return Ok(RunningThreadResumeResult::Handled);
};
let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse(
Box::new(crate::thread_state::PendingThreadResumeRequest {
request_id: request_id.clone(),
resume_handled_tx,
history_items,
config_snapshot,
instruction_sources,
@@ -3151,7 +3182,10 @@ impl ThreadRequestProcessor {
"failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed"
)));
}
return Ok(RunningThreadResumeResult::Handled);
return Ok(RunningThreadResumeResult::ListenerPending {
thread_id: existing_thread_id,
outcome_rx,
});
}
Ok(RunningThreadResumeResult::NotRunning(None))
}
codex-rs/app-server/src/thread_state.rs
+168 -8
View File
@@ -32,6 +32,7 @@ type PendingInterruptQueue = Vec<ConnectionRequestId>;
pub(crate) struct PendingThreadResumeRequest {
pub(crate) request_id: ConnectionRequestId,
pub(crate) resume_handled_tx: oneshot::Sender<ConnectionReservationAction>,
pub(crate) history_items: Vec<RolloutItem>,
pub(crate) config_snapshot: ThreadConfigSnapshot,
pub(crate) instruction_sources: Vec<LegacyAppPathString>,
@@ -389,6 +390,7 @@ mod tests {
struct ThreadEntry {
state: Arc<Mutex<ThreadState>>,
connection_ids: HashSet<ConnectionId>,
pending_connection_reservations: HashMap<ConnectionId, HashSet<u64>>,
has_connections_watcher: watch::Sender<bool>,
}
@@ -397,6 +399,7 @@ impl Default for ThreadEntry {
Self {
state: Arc::new(Mutex::new(ThreadState::default())),
connection_ids: HashSet::new(),
pending_connection_reservations: HashMap::new(),
has_connections_watcher: watch::channel(false).0,
}
}
@@ -417,6 +420,7 @@ struct ThreadStateManagerInner {
live_connections: HashMap<ConnectionId, ConnectionCapabilities>,
threads: HashMap<ThreadId, ThreadEntry>,
thread_ids_by_connection: HashMap<ConnectionId, HashSet<ThreadId>>,
next_connection_reservation_id: u64,
}
#[derive(Default)]
@@ -430,6 +434,25 @@ pub(crate) struct ConnectionCapabilities {
pub(crate) request_attestation: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ConnectionReservationOutcome {
Handled,
Abandoned,
}
#[derive(Clone, Copy, Debug)]
pub(crate) enum ConnectionReservationAction {
Commit,
Cancel,
}
#[derive(Clone, Copy)]
struct ConnectionReservation {
thread_id: ThreadId,
connection_id: ConnectionId,
reservation_id: u64,
}
#[derive(Clone, Default)]
pub(crate) struct ThreadStateManager {
state: Arc<Mutex<ThreadStateManagerInner>>,
@@ -748,16 +771,26 @@ impl ThreadStateManager {
return false;
}
let removed = if let Some(thread_entry) = state.threads.get_mut(&thread_id) {
let removed_subscription = thread_entry.connection_ids.remove(&connection_id);
let canceled_reservation = thread_entry
.pending_connection_reservations
.remove(&connection_id)
.is_some();
thread_entry.update_has_connections();
removed_subscription || canceled_reservation
} else {
false
};
if !removed {
return false;
}
if let Some(thread_ids) = state.thread_ids_by_connection.get_mut(&connection_id) {
thread_ids.remove(&thread_id);
if thread_ids.is_empty() {
state.thread_ids_by_connection.remove(&connection_id);
}
}
if let Some(thread_entry) = state.threads.get_mut(&thread_id) {
thread_entry.connection_ids.remove(&connection_id);
thread_entry.update_has_connections();
}
};
true
@@ -773,6 +806,18 @@ impl ThreadStateManager {
.is_some_and(|thread_entry| !thread_entry.connection_ids.is_empty())
}
pub(crate) async fn has_connections_or_reservations(&self, thread_id: ThreadId) -> bool {
self.state
.lock()
.await
.threads
.get(&thread_id)
.is_some_and(|thread_entry| {
!thread_entry.connection_ids.is_empty()
|| !thread_entry.pending_connection_reservations.is_empty()
})
}
pub(crate) async fn try_ensure_connection_subscribed(
&self,
thread_id: ThreadId,
@@ -809,6 +854,7 @@ impl ThreadStateManager {
Some(thread_state)
}
#[cfg(test)]
pub(crate) async fn try_add_connection_to_thread(
&self,
thread_id: ThreadId,
@@ -835,6 +881,117 @@ impl ThreadStateManager {
true
}
pub(crate) async fn start_connection_reservation(
&self,
thread_id: ThreadId,
connection_id: ConnectionId,
) -> Option<(
oneshot::Sender<ConnectionReservationAction>,
oneshot::Receiver<ConnectionReservationOutcome>,
)> {
let mut state = self.state.lock().await;
let listener_registry = self
.listener_commands
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if listener_registry.closing || !state.live_connections.contains_key(&connection_id) {
return None;
}
let reservation_id = state.next_connection_reservation_id;
let Some(next_reservation_id) = reservation_id.checked_add(1) else {
error!("connection reservation id space exhausted");
return None;
};
state.next_connection_reservation_id = next_reservation_id;
state
.thread_ids_by_connection
.entry(connection_id)
.or_default()
.insert(thread_id);
let thread_entry = state.threads.entry(thread_id).or_default();
thread_entry
.pending_connection_reservations
.entry(connection_id)
.or_default()
.insert(reservation_id);
let reservation = ConnectionReservation {
thread_id,
connection_id,
reservation_id,
};
let (handled_tx, handled_rx) = oneshot::channel();
let (outcome_tx, outcome_rx) = oneshot::channel();
let manager = self.clone();
let task = self.listener_tasks.spawn(async move {
let (action, outcome) = match handled_rx.await {
Ok(action) => (action, ConnectionReservationOutcome::Handled),
Err(_) => (
ConnectionReservationAction::Cancel,
ConnectionReservationOutcome::Abandoned,
),
};
manager
.finish_connection_reservation(reservation, action)
.await;
let _ = outcome_tx.send(outcome);
});
drop(task);
drop(listener_registry);
drop(state);
Some((handled_tx, outcome_rx))
}
async fn finish_connection_reservation(
&self,
reservation: ConnectionReservation,
action: ConnectionReservationAction,
) {
let ConnectionReservation {
thread_id,
connection_id,
reservation_id,
} = reservation;
let mut state = self.state.lock().await;
let connection_is_live = state.live_connections.contains_key(&connection_id);
let Some(thread_entry) = state.threads.get_mut(&thread_id) else {
return;
};
let remove_reservation_set = {
let Some(reservations) = thread_entry
.pending_connection_reservations
.get_mut(&connection_id)
else {
return;
};
if !reservations.remove(&reservation_id) {
return;
}
reservations.is_empty()
};
if remove_reservation_set {
thread_entry
.pending_connection_reservations
.remove(&connection_id);
}
let committed = matches!(action, ConnectionReservationAction::Commit) && connection_is_live;
if committed {
thread_entry.connection_ids.insert(connection_id);
}
let keep_connection_mapping = thread_entry.connection_ids.contains(&connection_id)
|| thread_entry
.pending_connection_reservations
.contains_key(&connection_id);
thread_entry.update_has_connections();
if !keep_connection_mapping
&& let Some(thread_ids) = state.thread_ids_by_connection.get_mut(&connection_id)
{
thread_ids.remove(&thread_id);
if thread_ids.is_empty() {
state.thread_ids_by_connection.remove(&connection_id);
}
}
}
pub(crate) async fn remove_connection(&self, connection_id: ConnectionId) -> Vec<ThreadId> {
{
let mut state = self.state.lock().await;
@@ -846,16 +1003,19 @@ impl ThreadStateManager {
for thread_id in &thread_ids {
if let Some(thread_entry) = state.threads.get_mut(thread_id) {
thread_entry.connection_ids.remove(&connection_id);
thread_entry
.pending_connection_reservations
.remove(&connection_id);
thread_entry.update_has_connections();
}
}
thread_ids
.into_iter()
.filter(|thread_id| {
state
.threads
.get(thread_id)
.is_some_and(|thread_entry| thread_entry.connection_ids.is_empty())
state.threads.get(thread_id).is_some_and(|thread_entry| {
thread_entry.connection_ids.is_empty()
&& thread_entry.pending_connection_reservations.is_empty()
})
})
.collect::<Vec<_>>()
}