app-server: keep the model cache warm (#28699)

## Why

The app server is long-lived, but its shared model cache otherwise
refreshes only when a caller needs it. Once the five-minute cache
expires, starting a thread or calling `model/list` can wait for
`/models` on the request path.

Refresh the cache in the background before it expires so foreground
callers normally use fresh local state.

## What changed

- Start an app-server worker that refreshes models immediately and then
every three minutes using the existing models-manager API.
- Hold only a weak reference to the models manager between refreshes, so
the worker does not extend its lifetime.
- Stop scheduling refreshes when the app-server lifecycle handle is shut
down or dropped. A refresh already in progress is allowed to finish.
- Adjust affected app-server test fixtures to distinguish the background
`/models` probe from the connection they are testing.

The existing models-manager cache, refresh strategies, auth handling,
ETag behavior, and concurrency semantics are unchanged.

## Testing

-
`models_refresh_worker::tests::refreshes_immediately_periodically_and_stops_when_dropped`
-
`suite::v2::remote_control::listen_off_honors_persisted_remote_control_enable`
-
`suite::v2::attestation::attestation_generate_round_trip_adds_header_to_responses_websocket_handshake`
This commit is contained in:
jif
2026-06-17 16:18:39 +02:00
committed by GitHub
parent 45f603302c
commit 5935a90619
7 changed files with 238 additions and 58 deletions
+1
View File
@@ -101,6 +101,7 @@ pub mod in_process;
mod mcp_refresh;
mod message_processor;
mod models;
mod models_refresh_worker;
mod outgoing_message;
mod request_processors;
mod request_serialization;
@@ -92,6 +92,8 @@ use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use crate::models_refresh_worker::ModelsRefreshWorker;
const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);
const CONNECTION_RPC_DRAIN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 30);
@@ -182,6 +184,7 @@ impl ExternalAuth for ExternalAuthRefreshBridge {
pub(crate) struct MessageProcessor {
outgoing: Arc<OutgoingMessageSender>,
models_refresh_worker: ModelsRefreshWorker,
skills_watcher: Arc<SkillsWatcher>,
account_processor: AccountRequestProcessor,
apps_processor: AppsRequestProcessor,
@@ -371,6 +374,8 @@ impl MessageProcessor {
)),
)
});
let models_manager = thread_manager.get_models_manager();
let models_refresh_worker = crate::models_refresh_worker::spawn(&models_manager);
thread_manager
.plugins_manager()
.set_analytics_events_client(analytics_events_client.clone());
@@ -537,6 +542,7 @@ impl MessageProcessor {
Self {
outgoing,
models_refresh_worker,
skills_watcher,
account_processor,
apps_processor,
@@ -566,6 +572,7 @@ impl MessageProcessor {
pub(crate) fn clear_runtime_references(&self) {
self.account_processor.clear_external_auth();
self.apps_processor.shutdown();
self.models_refresh_worker.shutdown();
self.skills_watcher.shutdown();
}
@@ -742,6 +749,7 @@ impl MessageProcessor {
}
pub(crate) async fn drain_background_tasks(&self) {
self.models_refresh_worker.shutdown();
self.thread_processor.drain_background_tasks().await;
}
@@ -0,0 +1,65 @@
use std::sync::Arc;
use std::time::Duration;
use codex_models_manager::manager::RefreshStrategy;
use codex_models_manager::manager::SharedModelsManager;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
const MODELS_REFRESH_INTERVAL: Duration = Duration::from_secs(3 * 60);
#[derive(Debug)]
pub(crate) struct ModelsRefreshWorker {
shutdown: CancellationToken,
_task: JoinHandle<()>,
}
impl ModelsRefreshWorker {
pub(crate) fn shutdown(&self) {
self.shutdown.cancel();
}
}
impl Drop for ModelsRefreshWorker {
fn drop(&mut self) {
self.shutdown();
}
}
pub(crate) fn spawn(models_manager: &SharedModelsManager) -> ModelsRefreshWorker {
spawn_with_interval(models_manager, MODELS_REFRESH_INTERVAL)
}
fn spawn_with_interval(
models_manager: &SharedModelsManager,
refresh_interval: Duration,
) -> ModelsRefreshWorker {
let models_manager = Arc::downgrade(models_manager);
let shutdown = CancellationToken::new();
let worker_shutdown = shutdown.clone();
let task = tokio::spawn(async move {
loop {
if worker_shutdown.is_cancelled() {
break;
}
let Some(models_manager) = models_manager.upgrade() else {
break;
};
models_manager.list_models(RefreshStrategy::Online).await;
drop(models_manager);
tokio::select! {
_ = worker_shutdown.cancelled() => break,
_ = tokio::time::sleep(refresh_interval) => {}
}
}
});
ModelsRefreshWorker {
shutdown,
_task: task,
}
}
#[cfg(test)]
#[path = "models_refresh_worker_tests.rs"]
mod tests;
@@ -0,0 +1,90 @@
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_models_manager::manager::ModelsEndpointClient;
use codex_models_manager::manager::ModelsEndpointFuture;
use codex_models_manager::manager::OpenAiModelsManager;
use codex_models_manager::manager::SharedModelsManager;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CoreResult;
use codex_protocol::openai_models::ModelInfo;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
use tokio::sync::Notify;
use super::*;
#[derive(Debug)]
struct TestModelsEndpoint {
fetch_count: AtomicUsize,
fetched: Notify,
release_second_fetch: Notify,
}
impl TestModelsEndpoint {
fn new() -> Arc<Self> {
Arc::new(Self {
fetch_count: AtomicUsize::new(0),
fetched: Notify::new(),
release_second_fetch: Notify::new(),
})
}
async fn wait_for_fetch_count(&self, expected: usize) {
tokio::time::timeout(Duration::from_secs(1), async {
while self.fetch_count.load(Ordering::SeqCst) < expected {
self.fetched.notified().await;
}
})
.await
.unwrap_or_else(|_| panic!("expected {expected} model fetches"));
}
}
impl ModelsEndpointClient for TestModelsEndpoint {
fn has_command_auth(&self) -> bool {
true
}
fn uses_codex_backend(&self) -> ModelsEndpointFuture<'_, bool> {
Box::pin(async { false })
}
fn list_models<'a>(
&'a self,
_client_version: &'a str,
) -> ModelsEndpointFuture<'a, CoreResult<(Vec<ModelInfo>, Option<String>)>> {
Box::pin(async move {
let fetch_index = self.fetch_count.fetch_add(1, Ordering::SeqCst);
self.fetched.notify_one();
if fetch_index == 0 {
return Err(CodexErr::Io(std::io::Error::other("test failure")));
}
if fetch_index == 1 {
self.release_second_fetch.notified().await;
}
Ok((Vec::new(), None))
})
}
}
#[tokio::test]
async fn refreshes_immediately_periodically_and_stops_when_dropped() {
let codex_home = tempdir().expect("temp dir");
let endpoint = TestModelsEndpoint::new();
let models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new(
codex_home.path().to_path_buf(),
endpoint.clone(),
/*auth_manager*/ None,
));
let worker = spawn_with_interval(&models_manager, Duration::from_millis(10));
endpoint.wait_for_fetch_count(/*expected*/ 2).await;
drop(worker);
endpoint.release_second_fetch.notify_one();
tokio::time::sleep(Duration::from_millis(30)).await;
assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 2);
}