Allow injecting the models manager into ThreadManager (#32911)

## Why

Embedding callers need to control whether model catalogs are persisted to disk.

## What changed

- Accept a shared models manager when constructing `ThreadManager` instead of always creating a cached manager internally.
- Add provider APIs for creating uncached model managers. OpenAI-compatible providers fetch on each `OnlineIfUncached` refresh without reading or writing `models_cache.json`.
- Keep existing app server, MCP server, sample, and test callers on the standard cached manager.

## Testing

- Verify an uncached manager fetches on every refresh.
- Verify an injected uncached manager controls thread refresh behavior and does not create `models_cache.json`.

GitOrigin-RevId: 8bb646054ecec4ccb865b2bf56249384916ea9db
This commit is contained in:
richardopenai
2026-07-14 01:06:11 +00:00
committed by copyberry
parent c39520f3d1
commit 80c6cd3014
17 changed files with 224 additions and 21 deletions
+1
View File
@@ -226,6 +226,7 @@ mod tests {
ThreadManager::new(
&good_config,
auth_manager.clone(),
codex_core::build_models_manager(&good_config, auth_manager.clone()),
SessionSource::Exec,
Arc::clone(&environment_manager),
thread_extensions(
@@ -262,6 +262,7 @@ impl MessageProcessor {
ThreadManager::new(
config.as_ref(),
auth_manager.clone(),
codex_core::build_models_manager(config.as_ref(), auth_manager.clone()),
session_source,
environment_manager,
thread_extensions(
+1
View File
@@ -34,6 +34,7 @@ pub use codex_core::StartThreadOptions;
pub use codex_core::StateDbHandle;
pub use codex_core::ThreadManager;
pub use codex_core::ThreadShutdownReport;
pub use codex_core::build_models_manager;
pub use codex_core::config::Config;
pub use codex_core::config::Constrained;
pub use codex_core::config::ExtraConfig;
+3 -1
View File
@@ -2393,9 +2393,11 @@ async fn spawn_thread_subagent_uses_role_specific_nickname_candidates() {
async fn resume_thread_subagent_restores_stored_metadata() {
let (home, config) = test_config().await;
let thread_store = Arc::new(InMemoryThreadStore::default());
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy"));
let manager = ThreadManager::new(
&config,
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),
auth_manager.clone(),
crate::thread_manager::build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
+1
View File
@@ -44,6 +44,7 @@ pub async fn build_prompt_input(
let thread_manager = ThreadManager::new(
&config,
Arc::clone(&auth_manager),
crate::thread_manager::build_models_manager(&config, Arc::clone(&auth_manager)),
SessionSource::Exec,
Arc::new(
EnvironmentManager::from_codex_home(
+2 -1
View File
@@ -305,6 +305,7 @@ impl ThreadManager {
pub fn new(
config: &Config,
auth_manager: Arc<AuthManager>,
models_manager: SharedModelsManager,
session_source: SessionSource,
environment_manager: Arc<EnvironmentManager>,
extensions: Arc<ExtensionRegistry<Config>>,
@@ -337,7 +338,7 @@ impl ThreadManager {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
thread_created_tx,
models_manager: build_models_manager(config, auth_manager.clone()),
models_manager,
environment_manager,
skills_service,
plugins_manager,
+73 -4
View File
@@ -596,9 +596,12 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(recorder.clone());
extensions.mcp_server_contributor(recorder);
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let manager = ThreadManager::new(
&config,
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()),
auth_manager.clone(),
build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
Arc::new(extensions.build()),
@@ -818,6 +821,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -943,7 +947,8 @@ async fn explicit_installation_id_skips_codex_home_file() {
let thread_store = thread_store_from_config(&config, state_db.clone());
let manager = ThreadManager::new(
&config,
auth_manager,
auth_manager.clone(),
build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -985,6 +990,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1045,6 +1051,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1112,6 +1119,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1203,7 +1211,8 @@ async fn subtree_listing_uses_injected_graph_store_without_state_db() {
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let manager = ThreadManager::new(
&config,
auth_manager,
auth_manager.clone(),
build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1249,6 +1258,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1353,7 +1363,8 @@ async fn new_uses_active_provider_for_model_refresh() {
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let manager = ThreadManager::new(
&config,
auth_manager,
auth_manager.clone(),
build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1375,6 +1386,61 @@ async fn new_uses_active_provider_for_model_refresh() {
assert_eq!(models_mock.requests().len(), 1);
}
#[tokio::test]
async fn injected_models_manager_controls_refresh_policy() {
let server = MockServer::start().await;
let _ = mount_models_once(&server, ModelsResponse { models: vec![] }).await;
let _ = mount_models_once(&server, ModelsResponse { models: vec![] }).await;
let temp_dir = tempdir().expect("tempdir");
let mut config = test_config().await;
config.codex_home = temp_dir.path().join("codex-home").abs();
config.cwd = config.codex_home.abs();
std::fs::create_dir_all(&config.codex_home).expect("create codex home");
config.model_catalog = None;
config.model_provider.base_url = Some(server.uri());
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let provider = create_model_provider(
config.model_provider.clone(),
Some(Arc::clone(&auth_manager)),
);
let models_manager = provider.models_manager_without_cache(config.model_catalog.clone());
let manager = ThreadManager::new(
&config,
auth_manager,
models_manager,
SessionSource::Custom("test-embedder".to_string()),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
);
let http_client_factory = crate::test_support::default_http_client_factory();
let _ = manager
.list_models(
RefreshStrategy::OnlineIfUncached,
http_client_factory.clone(),
)
.await;
let _ = manager
.list_models(RefreshStrategy::OnlineIfUncached, http_client_factory)
.await;
assert_eq!(
server.received_requests().await.unwrap_or_default().len(),
2
);
assert!(!config.codex_home.join("models_cache.json").exists());
}
#[test]
fn interrupted_fork_snapshot_appends_interrupt_boundary() {
let committed_history =
@@ -1593,6 +1659,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1703,6 +1770,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() {
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -1804,6 +1872,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
build_models_manager(&config, auth_manager.clone()),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
@@ -4293,9 +4293,11 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr
.enable(Feature::Sqlite)
.expect("test config should allow sqlite");
let state_db = init_state_db(&config).await;
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy"));
let manager = ThreadManager::new(
&config,
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),
auth_manager.clone(),
crate::thread_manager::build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
+3 -1
View File
@@ -621,9 +621,11 @@ impl TestCodexBuilder {
config.codex_home.clone(),
))
});
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
let thread_manager = ThreadManager::new(
&config,
codex_core::test_support::auth_manager_from_auth(auth.clone()),
auth_manager.clone(),
codex_core::build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::clone(&environment_manager),
Arc::clone(&self.extensions),
+2 -1
View File
@@ -1579,7 +1579,8 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() {
.expect("resolve installation id");
let thread_manager = ThreadManager::new(
&config,
auth_manager,
auth_manager.clone(),
codex_core::build_models_manager(&config, auth_manager),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
+2 -1
View File
@@ -74,7 +74,8 @@ impl MessageProcessor {
);
let thread_manager = Arc::new(ThreadManager::new(
config.as_ref(),
auth_manager,
Arc::clone(&auth_manager),
codex_core::build_models_manager(config.as_ref(), auth_manager),
SessionSource::Mcp,
environment_manager,
Arc::new(extensions.build()),
@@ -172,6 +172,16 @@ impl ModelProvider for AmazonBedrockModelProvider {
config_model_catalog.map_or_else(static_model_catalog, with_default_only_service_tier),
))
}
fn models_manager_without_cache(
&self,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
Arc::new(StaticModelsManager::new(
/*auth_manager*/ None,
config_model_catalog.map_or_else(static_model_catalog, with_default_only_service_tier),
))
}
}
#[cfg(test)]
+45
View File
@@ -199,6 +199,20 @@ pub trait ModelProvider: fmt::Debug + Send + Sync {
codex_home: PathBuf,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager;
/// Creates a model manager with caching disabled.
///
/// Providers that fetch model catalogs should override this method. The default uses an
/// authoritative in-memory catalog so hosted callers cannot accidentally write to disk.
fn models_manager_without_cache(
&self,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
let model_catalog = config_model_catalog
.or_else(|| codex_models_manager::bundled_models_response().ok())
.unwrap_or_default();
Arc::new(StaticModelsManager::new(self.auth_manager(), model_catalog))
}
}
pub type ModelProviderFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -334,6 +348,28 @@ impl ModelProvider for ConfiguredModelProvider {
}
}
}
fn models_manager_without_cache(
&self,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
match config_model_catalog {
Some(model_catalog) => Arc::new(StaticModelsManager::new(
self.auth_manager.clone(),
model_catalog,
)),
None => {
let endpoint = Arc::new(OpenAiModelsEndpoint::new(
self.info.clone(),
self.auth_manager.clone(),
));
Arc::new(OpenAiModelsManager::new_without_cache(
endpoint,
self.auth_manager.clone(),
))
}
}
}
}
#[cfg(test)]
@@ -656,6 +692,8 @@ mod tests {
);
let manager =
provider.models_manager(test_codex_home(), /*config_model_catalog*/ None);
let uncached_manager =
provider.models_manager_without_cache(/*config_model_catalog*/ None);
let catalog = manager
.raw_model_catalog(
@@ -663,6 +701,13 @@ mod tests {
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
.await;
let uncached_catalog = uncached_manager
.raw_model_catalog(
RefreshStrategy::Online,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
.await;
assert_eq!(uncached_catalog, catalog);
let models = catalog
.models
.iter()
+2 -2
View File
@@ -30,8 +30,8 @@ impl ModelsCacheManager {
/// Attempt to load a fresh cache entry. Returns `None` if the cache doesn't exist or is stale.
pub(crate) async fn load_fresh(&self, expected_version: &str) -> Option<ModelsCache> {
info!(
cache_path = %self.cache_path.display(),
expected_version,
cache_path = %self.cache_path.display(),
expected_version,
"models cache: attempting load_fresh"
);
let cache = match self.load().await {
+33 -7
View File
@@ -215,7 +215,7 @@ pub type SharedModelsManager = Arc<dyn ModelsManager>;
pub struct OpenAiModelsManager {
remote_models: RwLock<Vec<ModelInfo>>,
etag: RwLock<Option<String>>,
cache_manager: ModelsCacheManager,
cache_manager: Option<ModelsCacheManager>,
endpoint_client: SharedModelsEndpointClient,
auth_manager: Option<Arc<AuthManager>>,
}
@@ -235,7 +235,26 @@ impl OpenAiModelsManager {
auth_manager: Option<Arc<AuthManager>>,
) -> Self {
let cache_path = codex_home.join(MODEL_CACHE_FILE);
let cache_manager = ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL);
Self::new_with_cache_manager(
Some(ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL)),
endpoint_client,
auth_manager,
)
}
/// Construct an OpenAI-compatible model manager with caching disabled.
pub fn new_without_cache(
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> Self {
Self::new_with_cache_manager(/*cache_manager*/ None, endpoint_client, auth_manager)
}
fn new_with_cache_manager(
cache_manager: Option<ModelsCacheManager>,
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> Self {
let remote_models = load_remote_models_from_file().unwrap_or_default();
Self {
remote_models: RwLock::new(remote_models),
@@ -319,7 +338,9 @@ impl OpenAiModelsManager {
async fn refresh_if_new_etag(&self, etag: String, http_client_factory: HttpClientFactory) {
let current_etag = self.get_etag().await;
if current_etag.clone().is_some() && current_etag.as_deref() == Some(etag.as_str()) {
if let Err(err) = self.cache_manager.renew_cache_ttl().await {
if let Some(cache_manager) = self.cache_manager.as_ref()
&& let Err(err) = cache_manager.renew_cache_ttl().await
{
error!("failed to renew cache TTL: {err}");
}
return;
@@ -381,9 +402,11 @@ impl OpenAiModelsManager {
.await?;
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.await;
if let Some(cache_manager) = self.cache_manager.as_ref() {
cache_manager
.persist_cache(&models, etag, client_version)
.await;
}
Ok(())
}
@@ -429,13 +452,16 @@ impl OpenAiModelsManager {
/// Attempt to satisfy the refresh from the cache when it matches the provider and TTL.
async fn try_load_cache(&self) -> bool {
let Some(cache_manager) = self.cache_manager.as_ref() else {
return false;
};
let _timer =
codex_otel::start_global_timer("codex.remote_models.load_cache.duration_ms", &[]);
let client_version = crate::client_version_to_whole();
info!(client_version, "models cache: evaluating cache eligibility");
// TODO(celia-oai): Include provider identity in cache eligibility so switching
// providers does not reuse a fresh models_cache.json entry from another provider.
let cache = match self.cache_manager.load_fresh(&client_version).await {
let cache = match cache_manager.load_fresh(&client_version).await {
Some(cache) => cache,
None => {
info!("models cache: no usable cache entry");
+39 -1
View File
@@ -209,6 +209,36 @@ fn static_manager_for_tests(model_catalog: ModelsResponse) -> StaticModelsManage
StaticModelsManager::new(/*auth_manager*/ None, model_catalog)
}
#[tokio::test]
async fn manager_without_cache_fetches_on_every_refresh() {
let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)];
let endpoint = TestModelsEndpoint::new(vec![remote_models.clone(), remote_models.clone()]);
let manager = OpenAiModelsManager::new_without_cache(
endpoint.clone(),
Some(AuthManager::from_auth_for_testing(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
)),
);
let catalog = manager
.raw_model_catalog(
RefreshStrategy::OnlineIfUncached,
DEFAULT_HTTP_CLIENT_FACTORY,
)
.await;
let second_catalog = manager
.raw_model_catalog(
RefreshStrategy::OnlineIfUncached,
DEFAULT_HTTP_CLIENT_FACTORY,
)
.await;
assert_eq!(catalog.models, remote_models);
assert_eq!(second_catalog, catalog);
assert_eq!(manager.get_remote_models().await, remote_models);
assert_eq!(endpoint.fetch_count(), 2);
}
async fn chatgpt_auth_tokens_for_tests(codex_home: &Path) -> CodexAuth {
let auth_dot_json = codex_login::AuthDotJson {
auth_mode: Some(AuthMode::ChatgptAuthTokens),
@@ -719,6 +749,8 @@ async fn refresh_available_models_refetches_when_cache_stale() {
// Rewrite cache with an old timestamp so it is treated as stale.
manager
.cache_manager
.as_ref()
.expect("cached model manager")
.manipulate_cache_for_test(|fetched_at| {
*fetched_at = Utc::now() - chrono::Duration::hours(1);
})
@@ -758,6 +790,8 @@ async fn refresh_available_models_refetches_when_version_mismatch() {
manager
.cache_manager
.as_ref()
.expect("cached model manager")
.mutate_cache_for_test(|cache| {
let client_version = crate::client_version_to_whole();
cache.client_version = Some(format!("{client_version}-mismatch"));
@@ -795,7 +829,11 @@ async fn refresh_available_models_drops_removed_remote_models() {
)];
let endpoint = TestModelsEndpoint::new(vec![initial_models, refreshed_models]);
let mut manager = openai_manager_for_tests(codex_home.path().to_path_buf(), endpoint.clone());
manager.cache_manager.set_ttl(Duration::ZERO);
manager
.cache_manager
.as_mut()
.expect("cached model manager")
.set_ttl(Duration::ZERO);
manager
.refresh_available_models(
+3 -1
View File
@@ -55,6 +55,7 @@ use codex_core_api::UriBasedFileOpener;
use codex_core_api::UserInput;
use codex_core_api::WebSearchMode;
use codex_core_api::arg0_dispatch_or_else;
use codex_core_api::build_models_manager;
use codex_core_api::built_in_model_providers;
use codex_core_api::find_codex_home;
use codex_core_api::init_state_db;
@@ -132,7 +133,8 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
});
let thread_manager = ThreadManager::new(
&config,
auth_manager,
Arc::clone(&auth_manager),
build_models_manager(&config, auth_manager),
SessionSource::Exec,
environment_manager,
Arc::new(extensions.build()),