agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

app-server: Only unload threads which were unused for some time (#17398)

Browse Source 
Currently app-server may unload actively running threads once the last
connection disconnects, which is not expected.
Instead track when was the last active turn & when there were any
subscribers the last time, also add 30 minute idleness/no subscribers
timer to reduce the churn.
This commit is contained in:
Ruslan Nigmatullin
2026-04-13 12:25:26 -07:00
committed by GitHub
Unverified
parent d905376628
commit a5507b59c4
6 changed files with 495 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/codex_message_processor.rs
+325 -80
View File
@@ -329,6 +329,7 @@ use std::sync::RwLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
@@ -371,6 +372,7 @@ struct ThreadListFilters {
const LOGIN_CHATGPT_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const LOGIN_ISSUER_OVERRIDE_ENV_VAR: &str = "CODEX_APP_SERVER_LOGIN_ISSUER";
const APP_LIST_LOAD_TIMEOUT: Duration = Duration::from_secs(90);
const THREAD_UNLOADING_DELAY: Duration = Duration::from_secs(30 * 60);
enum ActiveLogin {
Browser {
@@ -460,6 +462,7 @@ struct ListenerTaskContext {
thread_manager: Arc<ThreadManager>,
thread_state_manager: ThreadStateManager,
outgoing: Arc<OutgoingMessageSender>,
pending_thread_unloads: Arc<Mutex<HashSet<ThreadId>>>,
analytics_events_client: AnalyticsEventsClient,
general_analytics_enabled: bool,
thread_watch_manager: ThreadWatchManager,
@@ -480,6 +483,110 @@ enum RefreshTokenRequestOutcome {
FailedPermanently,
}
struct UnloadingState {
delay: Duration,
has_subscribers_rx: watch::Receiver<bool>,
has_subscribers: (bool, Instant),
thread_status_rx: watch::Receiver<ThreadStatus>,
is_active: (bool, Instant),
}
impl UnloadingState {
async fn new(
listener_task_context: &ListenerTaskContext,
thread_id: ThreadId,
delay: Duration,
) -> Option<Self> {
let has_subscribers_rx = listener_task_context
.thread_state_manager
.subscribe_to_has_connections(thread_id)
.await?;
let thread_status_rx = listener_task_context
.thread_watch_manager
.subscribe(thread_id)
.await?;
let has_subscribers = (*has_subscribers_rx.borrow(), Instant::now());
let is_active = (
matches!(*thread_status_rx.borrow(), ThreadStatus::Active { .. }),
Instant::now(),
);
Some(Self {
delay,
has_subscribers_rx,
thread_status_rx,
has_subscribers,
is_active,
})
}
fn unloading_target(&self) -> Option<Instant> {
match (self.has_subscribers, self.is_active) {
((false, has_no_subscribers_since), (false, is_inactive_since)) => {
Some(std::cmp::max(has_no_subscribers_since, is_inactive_since) + self.delay)
}
_ => None,
}
}
fn sync_receiver_values(&mut self) {
let has_subscribers = *self.has_subscribers_rx.borrow();
if self.has_subscribers.0 != has_subscribers {
self.has_subscribers = (has_subscribers, Instant::now());
}
let is_active = matches!(*self.thread_status_rx.borrow(), ThreadStatus::Active { .. });
if self.is_active.0 != is_active {
self.is_active = (is_active, Instant::now());
}
}
fn should_unload_now(&mut self) -> bool {
self.sync_receiver_values();
self.unloading_target()
.is_some_and(|target| target <= Instant::now())
}
fn note_thread_activity_observed(&mut self) {
if !self.is_active.0 {
self.is_active = (false, Instant::now());
}
}
async fn wait_for_unloading_trigger(&mut self) -> bool {
loop {
self.sync_receiver_values();
let unloading_target = self.unloading_target();
if let Some(target) = unloading_target
&& target <= Instant::now()
{
return true;
}
let unloading_sleep = async {
if let Some(target) = unloading_target {
tokio::time::sleep_until(target.into()).await;
} else {
futures::future::pending::<()>().await;
}
};
tokio::select! {
_ = unloading_sleep => return true,
changed = self.has_subscribers_rx.changed() => {
if changed.is_err() {
return false;
}
self.sync_receiver_values();
},
changed = self.thread_status_rx.changed() => {
if changed.is_err() {
return false;
}
self.sync_receiver_values();
},
}
}
}
}
pub(crate) struct CodexMessageProcessorArgs {
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) thread_manager: Arc<ThreadManager>,
@@ -2149,6 +2256,7 @@ impl CodexMessageProcessor {
thread_manager: Arc::clone(&self.thread_manager),
thread_state_manager: self.thread_state_manager.clone(),
outgoing: Arc::clone(&self.outgoing),
pending_thread_unloads: Arc::clone(&self.pending_thread_unloads),
analytics_events_client: self.analytics_events_client.clone(),
general_analytics_enabled: self.config.features.enabled(Feature::GeneralAnalytics),
thread_watch_manager: self.thread_watch_manager.clone(),
@@ -3884,17 +3992,17 @@ impl CodexMessageProcessor {
self.command_exec_manager
.connection_closed(connection_id)
.await;
let thread_ids_with_no_subscribers = self
let thread_ids = self
.thread_state_manager
.remove_connection(connection_id)
.await;
for thread_id in thread_ids_with_no_subscribers {
let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
for thread_id in thread_ids {
if self.thread_manager.get_thread(thread_id).await.is_err() {
// Reconcile stale app-server bookkeeping when the thread has already been
// removed from the core manager.
self.finalize_thread_teardown(thread_id).await;
continue;
};
self.unload_thread_without_subscribers(thread_id, thread)
.await;
}
}
}
@@ -4260,13 +4368,18 @@ impl CodexMessageProcessor {
.thread_state_manager
.thread_state(existing_thread_id)
.await;
self.ensure_listener_task_running(
existing_thread_id,
existing_thread.clone(),
thread_state.clone(),
ApiVersion::V2,
)
.await;
if let Err(error) = self
.ensure_listener_task_running(
existing_thread_id,
existing_thread.clone(),
thread_state.clone(),
ApiVersion::V2,
)
.await
{
self.outgoing.send_error(request_id, error).await;
return true;
}
let config_snapshot = existing_thread.config_snapshot().await;
let mismatch_details = collect_resume_override_mismatches(params, &config_snapshot);
@@ -5653,31 +5766,23 @@ impl CodexMessageProcessor {
}
async fn unload_thread_without_subscribers(
&self,
thread_manager: Arc<ThreadManager>,
outgoing: Arc<OutgoingMessageSender>,
pending_thread_unloads: Arc<Mutex<HashSet<ThreadId>>>,
thread_state_manager: ThreadStateManager,
thread_watch_manager: ThreadWatchManager,
thread_id: ThreadId,
thread: Arc<CodexThread>,
) {
// This connection was the last subscriber. Only now do we unload the thread.
info!("thread {thread_id} has no subscribers; shutting down");
let should_start_unload_task = self.pending_thread_unloads.lock().await.insert(thread_id);
info!("thread {thread_id} has no subscribers and is idle; shutting down");
// Any pending app-server -> client requests for this thread can no longer be
// answered; cancel their callbacks before shutdown/unload.
self.outgoing
outgoing
.cancel_requests_for_thread(thread_id, /*error*/ None)
.await;
self.thread_state_manager
.remove_thread_state(thread_id)
.await;
thread_state_manager.remove_thread_state(thread_id).await;
if !should_start_unload_task {
return;
}
let outgoing = self.outgoing.clone();
let pending_thread_unloads = self.pending_thread_unloads.clone();
let thread_manager = self.thread_manager.clone();
let thread_watch_manager = self.thread_watch_manager.clone();
tokio::spawn(async move {
match Self::wait_for_thread_shutdown(&thread).await {
ThreadShutdownResult::Complete => {
@@ -5726,7 +5831,7 @@ impl CodexMessageProcessor {
}
};
let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
if self.thread_manager.get_thread(thread_id).await.is_err() {
// Reconcile stale app-server bookkeeping when the thread has already been
// removed from the core manager. This keeps loaded-status/subscription state
// consistent with the source of truth before reporting NotLoaded.
@@ -5746,30 +5851,14 @@ impl CodexMessageProcessor {
.thread_state_manager
.unsubscribe_connection_from_thread(thread_id, request_id.connection_id)
.await;
if !was_subscribed {
self.outgoing
.send_response(
request_id,
ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::NotSubscribed,
},
)
.await;
return;
}
if !self.thread_state_manager.has_subscribers(thread_id).await {
self.unload_thread_without_subscribers(thread_id, thread)
.await;
}
let status = if was_subscribed {
ThreadUnsubscribeStatus::Unsubscribed
} else {
ThreadUnsubscribeStatus::NotSubscribed
};
self.outgoing
.send_response(
request_id,
ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::Unsubscribed,
},
)
.send_response(request_id, ThreadUnsubscribeResponse { status })
.await;
}
@@ -7514,6 +7603,7 @@ impl CodexMessageProcessor {
thread_manager: Arc::clone(&self.thread_manager),
thread_state_manager: self.thread_state_manager.clone(),
outgoing: Arc::clone(&self.outgoing),
pending_thread_unloads: Arc::clone(&self.pending_thread_unloads),
analytics_events_client: self.analytics_events_client.clone(),
general_analytics_enabled: self.config.features.enabled(Feature::GeneralAnalytics),
thread_watch_manager: self.thread_watch_manager.clone(),
@@ -7549,21 +7639,45 @@ impl CodexMessageProcessor {
});
}
};
let Some(thread_state) = listener_task_context
.thread_state_manager
.try_ensure_connection_subscribed(conversation_id, connection_id, raw_events_enabled)
.await
else {
return Ok(EnsureConversationListenerResult::ConnectionClosed);
let thread_state = {
let pending_thread_unloads = listener_task_context.pending_thread_unloads.lock().await;
if pending_thread_unloads.contains(&conversation_id) {
return Err(JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!(
"thread {conversation_id} is closing; retry after the thread is closed"
),
data: None,
});
}
let Some(thread_state) = listener_task_context
.thread_state_manager
.try_ensure_connection_subscribed(
conversation_id,
connection_id,
raw_events_enabled,
)
.await
else {
return Ok(EnsureConversationListenerResult::ConnectionClosed);
};
thread_state
};
Self::ensure_listener_task_running_task(
listener_task_context,
if let Err(error) = Self::ensure_listener_task_running_task(
listener_task_context.clone(),
conversation_id,
conversation,
thread_state,
api_version,
)
.await;
.await
{
let _ = listener_task_context
.thread_state_manager
.unsubscribe_connection_from_thread(conversation_id, connection_id)
.await;
return Err(error);
}
Ok(EnsureConversationListenerResult::Attached)
}
@@ -7597,12 +7711,13 @@ impl CodexMessageProcessor {
conversation: Arc<CodexThread>,
thread_state: Arc<Mutex<ThreadState>>,
api_version: ApiVersion,
) {
) -> Result<(), JSONRPCErrorError> {
Self::ensure_listener_task_running_task(
ListenerTaskContext {
thread_manager: Arc::clone(&self.thread_manager),
thread_state_manager: self.thread_state_manager.clone(),
outgoing: Arc::clone(&self.outgoing),
pending_thread_unloads: Arc::clone(&self.pending_thread_unloads),
analytics_events_client: self.analytics_events_client.clone(),
general_analytics_enabled: self.config.features.enabled(Feature::GeneralAnalytics),
thread_watch_manager: self.thread_watch_manager.clone(),
@@ -7614,7 +7729,7 @@ impl CodexMessageProcessor {
thread_state,
api_version,
)
.await;
.await
}
async fn ensure_listener_task_running_task(
@@ -7623,12 +7738,27 @@ impl CodexMessageProcessor {
conversation: Arc<CodexThread>,
thread_state: Arc<Mutex<ThreadState>>,
api_version: ApiVersion,
) {
) -> Result<(), JSONRPCErrorError> {
let (cancel_tx, mut cancel_rx) = oneshot::channel();
let Some(mut unloading_state) = UnloadingState::new(
&listener_task_context,
conversation_id,
THREAD_UNLOADING_DELAY,
)
.await
else {
return Err(JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!(
"thread {conversation_id} is closing; retry after the thread is closed"
),
data: None,
});
};
let (mut listener_command_rx, listener_generation) = {
let mut thread_state = thread_state.lock().await;
if thread_state.listener_matches(&conversation) {
return;
return Ok(());
}
thread_state.set_listener(cancel_tx, &conversation)
};
@@ -7636,6 +7766,7 @@ impl CodexMessageProcessor {
outgoing,
thread_manager,
thread_state_manager,
pending_thread_unloads,
analytics_events_client: _,
general_analytics_enabled: _,
thread_watch_manager,
@@ -7646,10 +7777,28 @@ impl CodexMessageProcessor {
tokio::spawn(async move {
loop {
tokio::select! {
biased;
_ = &mut cancel_rx => {
// Listener was superseded or the thread is being torn down.
break;
}
listener_command = listener_command_rx.recv() => {
let Some(listener_command) = listener_command else {
break;
};
handle_thread_listener_command(
conversation_id,
&conversation,
codex_home.as_path(),
&thread_state_manager,
&thread_state,
&thread_watch_manager,
&outgoing_for_task,
&pending_thread_unloads,
listener_command,
)
.await;
}
event = conversation.next_event() => {
let event = match event {
Ok(event) => event,
@@ -7704,21 +7853,38 @@ impl CodexMessageProcessor {
)
.await;
}
listener_command = listener_command_rx.recv() => {
let Some(listener_command) = listener_command else {
unloading_watchers_open = unloading_state.wait_for_unloading_trigger() => {
if !unloading_watchers_open {
break;
};
handle_thread_listener_command(
}
if !unloading_state.should_unload_now() {
continue;
}
if matches!(conversation.agent_status().await, AgentStatus::Running) {
unloading_state.note_thread_activity_observed();
continue;
}
{
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);
}
Self::unload_thread_without_subscribers(
thread_manager.clone(),
outgoing_for_task.clone(),
pending_thread_unloads.clone(),
thread_state_manager.clone(),
thread_watch_manager.clone(),
conversation_id,
&conversation,
codex_home.as_path(),
&thread_state_manager,
&thread_state,
&thread_watch_manager,
&outgoing_for_task,
listener_command,
conversation.clone(),
)
.await;
break;
}
}
}
@@ -7728,6 +7894,7 @@ impl CodexMessageProcessor {
thread_state.clear_listener();
}
});
Ok(())
}
async fn git_diff_to_origin(&self, request_id: ConnectionRequestId, cwd: PathBuf) {
let diff = git_diff_to_remote(&cwd).await;
@@ -8218,6 +8385,7 @@ async fn handle_thread_listener_command(
thread_state: &Arc<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
outgoing: &Arc<OutgoingMessageSender>,
pending_thread_unloads: &Arc<Mutex<HashSet<ThreadId>>>,
listener_command: ThreadListenerCommand,
) {
match listener_command {
@@ -8230,6 +8398,7 @@ async fn handle_thread_listener_command(
thread_state,
thread_watch_manager,
outgoing,
pending_thread_unloads,
*resume_request,
)
.await;
@@ -8259,6 +8428,7 @@ async fn handle_pending_thread_resume_request(
thread_state: &Arc<Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
outgoing: &Arc<OutgoingMessageSender>,
pending_thread_unloads: &Arc<Mutex<HashSet<ThreadId>>>,
pending: crate::thread_state::PendingThreadResumeRequest,
) {
let active_turn = {
@@ -8312,6 +8482,37 @@ async fn handle_pending_thread_resume_request(
has_live_in_progress_turn,
);
{
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,
JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!(
"thread {conversation_id} is closing; retry thread/resume after the thread is closed"
),
data: None,
},
)
.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;
}
}
let ThreadConfigSnapshot {
model,
model_provider_id,
@@ -8340,9 +8541,6 @@ async fn handle_pending_thread_resume_request(
outgoing
.replay_requests_to_connection_for_thread(connection_id, conversation_id)
.await;
let _attached = thread_state_manager
.try_add_connection_to_thread(conversation_id, connection_id)
.await;
}
enum ThreadTurnSource<'a> {
@@ -10137,6 +10335,53 @@ mod tests {
Ok(())
}
#[tokio::test]
async fn adding_connection_to_thread_updates_has_connections_watcher() -> Result<()> {
let manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let connection_a = ConnectionId(1);
let connection_b = ConnectionId(2);
manager.connection_initialized(connection_a).await;
manager.connection_initialized(connection_b).await;
manager
.try_ensure_connection_subscribed(
thread_id,
connection_a,
/*experimental_raw_events*/ false,
)
.await
.expect("connection_a should be live");
let mut has_connections = manager
.subscribe_to_has_connections(thread_id)
.await
.expect("thread should have a has-connections watcher");
assert!(*has_connections.borrow());
assert!(
manager
.unsubscribe_connection_from_thread(thread_id, connection_a)
.await
);
tokio::time::timeout(Duration::from_secs(1), has_connections.changed())
.await
.expect("timed out waiting for no-subscriber update")
.expect("has-connections watcher should remain open");
assert!(!*has_connections.borrow());
assert!(
manager
.try_add_connection_to_thread(thread_id, connection_b)
.await
);
tokio::time::timeout(Duration::from_secs(1), has_connections.changed())
.await
.expect("timed out waiting for subscriber update")
.expect("has-connections watcher should remain open");
assert!(*has_connections.borrow());
Ok(())
}
#[tokio::test]
async fn closed_connection_cannot_be_reintroduced_by_auto_subscribe() -> Result<()> {
let manager = ThreadStateManager::new();