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

Preserve model catalogs when starting child agents (#60049)

Browse Source 
## Why

Child startup can lose remote model capabilities such as `tool_search` when the selected catalog is uninitialized or disk data is older than the parent's in-memory snapshot.

## What changed

- Initialize the final provider, auth, and originator catalog before resolving a child's startup model, then read it from memory.
- Preserve loaded catalogs, including empty responses and updates that failed to persist, without overwriting concurrent refreshes.
- Reuse catalogs across retry and WebSocket transport overrides while retaining routing and header distinctions.

## Testing

Add regression coverage for child tool search with warm, cold, and disabled caches; failed cache writes; concurrent refreshes; empty catalogs; and Guardian reviewer catalog reuse.

GitOrigin-RevId: 01d995a65b623cc7d6f26e2dc9b6e8cf380de782
This commit is contained in:
evanz-oai
2026-11-07 03:16:52 +00:00
committed by copyberry
parent 358e464528
commit 44ccf39457
15 changed files with 682 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/mod.rs
+4 -5
View File
@@ -605,15 +605,14 @@ impl Session {
session_source: &SessionSource,
allow_provider_model_fallback: bool,
) -> CodexResult<(String, CatalogModel)> {
// Child catalogs are initialized by their selector. Reading only memory here
// preserves the loaded snapshot shared with the parent.
let refresh_strategy = if session_source.is_non_root_agent() {
RefreshStrategy::Offline
RefreshStrategy::InMemory
} else {
RefreshStrategy::OnlineIfUncached
};
if config.features.enabled(Feature::RuntimeAgentInstances)
|| config.model.is_none()
|| refresh_strategy != RefreshStrategy::Offline
{
if !session_source.is_non_root_agent() {
let _ = models_manager
.list_models(refresh_strategy, config.http_client_factory())
.await;
codex-rs/core/src/session/step_activation_tests.rs
+4
View File
@@ -360,6 +360,10 @@ impl GatedModelsManager {
}
impl ModelsManager for GatedModelsManager {
fn ensure_initialized(&self, factory: HttpClientFactory) -> ModelsManagerFuture<'_, ()> {
self.inner.ensure_initialized(factory)
}
fn model_catalog_with_origins(
&self,
strategy: RefreshStrategy,
codex-rs/core/src/thread_manager.rs
+38 -20
View File
@@ -1116,12 +1116,16 @@ impl ThreadManager {
/*parent_originator*/ None,
)
.await;
let manager = self.state.models_for_thread(
&options.config,
&self.state.auth_manager,
&self.state.runtime_auth,
&originator,
);
let manager = self
.state
.models_for_thread(
&options.config,
&self.state.auth_manager,
&self.state.runtime_auth,
&originator,
&source,
)
.await;
let (_, info) = Session::resolve_startup_model(
manager.as_ref(),
&options.config,
@@ -1919,12 +1923,13 @@ impl ThreadManager {
}
impl ThreadManagerState {
fn models_for_thread(
async fn models_for_thread(
&self,
config: &Config,
auth_manager: &Arc<AuthManager>,
runtime_auth: &crate::runtime_auth::RuntimeAuthSlot,
originator: &str,
session_source: &SessionSource,
) -> SharedModelsManager {
let runtime = crate::runtime_auth::runtime_for_config(
runtime_auth,
@@ -1932,13 +1937,16 @@ impl ThreadManagerState {
Arc::clone(auth_manager),
&self.installation_id,
);
let manager = self.model_catalog.for_thread(
config,
Arc::clone(auth_manager),
runtime,
Arc::ptr_eq(auth_manager, &self.auth_manager),
);
manager.for_originator(originator).unwrap_or(manager)
self.model_catalog
.for_thread(
config,
Arc::clone(auth_manager),
runtime,
Arc::ptr_eq(auth_manager, &self.auth_manager),
originator,
session_source,
)
.await
}
pub(crate) fn shared_thread_instructions_provider(
@@ -2775,12 +2783,22 @@ impl ThreadManagerState {
} else {
Arc::new(mcp_manager.for_plugins_manager(Arc::clone(&plugins_manager)))
};
let models_manager = self.models_for_thread(
&config,
&auth_manager,
&runtime_auth,
&startup_context.originator,
);
let models_manager = tokio::select! {
biased;
_ = agent_control.runtime().shutdown.cancelled() => {
if let Some(startup_guard) = startup_guard {
startup_guard.cleanup().await;
}
return Err(CodexErr::TurnAborted);
}
manager = self.models_for_thread(
&config,
&auth_manager,
&runtime_auth,
&startup_context.originator,
&session_source,
) => manager,
};
let spawn_result = Session::spawn(SessionSpawnArgs {
attached_thread,
startup,
codex-rs/core/src/thread_manager/model_catalog.rs
+29 -4
View File
@@ -12,6 +12,7 @@ use codex_model_provider_info::ModelProviderInfo;
use codex_models_manager::manager::ModelCatalogSettings;
use codex_models_manager::manager::SharedModelsManager;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::protocol::SessionSource;
use codex_runtime_agent_instance::RuntimeAuthManager;
use std::sync::Arc;
@@ -115,23 +116,35 @@ impl ThreadModelCatalog {
}
}
pub(super) fn for_thread(
#[allow(clippy::too_many_arguments)]
pub(super) async fn for_thread(
&self,
config: &Config,
credentials: Arc<AuthManager>,
runtime: Option<Arc<RuntimeAuthManager>>,
same_owner: bool,
originator: &str,
session_source: &SessionSource,
) -> SharedModelsManager {
let models = config.model_catalog.as_ref().or_else(|| {
self.catalog
.as_ref()
.filter(|_| self.provider == config.model_provider)
.filter(|_| self.matches_catalog_provider(&config.model_provider))
});
let auth = ProviderAuthContext {
credentials,
runtime_auth: runtime.map(|runtime| (runtime, config.runtime_auth_mode)),
};
self.for_config(config, models, auth, same_owner)
let manager = self.for_config(config, models, auth, same_owner);
let manager = manager.for_originator(originator).unwrap_or(manager);
if session_source.is_non_root_agent() {
// The final provider/auth/originator catalog may be new even when the base
// manager was reused. Preserve loaded snapshots, including unsaved updates.
manager
.ensure_initialized(config.http_client_factory())
.await;
}
manager
}
pub(super) fn for_config(
@@ -145,7 +158,7 @@ impl ThreadModelCatalog {
config.features.enabled(Feature::ApiKeyModelDiscovery);
if same_owner
&& self.runtime_mode == auth.runtime_auth.as_ref().map(|(_, mode)| *mode)
&& self.provider == config.model_provider
&& self.matches_catalog_provider(&config.model_provider)
&& self.catalog.as_ref() == models
&& self.api_key_model_discovery_default == api_key_model_discovery_default
{
@@ -158,6 +171,18 @@ impl ThreadModelCatalog {
)
}
fn matches_catalog_provider(&self, provider: &ModelProviderInfo) -> bool {
// Guardian overrides inference transport and retries without changing the catalog.
// Keep all other provider fields, including routing and headers, in the comparison.
self.provider
== ModelProviderInfo {
request_max_retries: self.provider.request_max_retries,
stream_max_retries: self.provider.stream_max_retries,
supports_websockets: self.provider.supports_websockets,
..provider.clone()
}
}
pub(super) fn set_api_key_model_discovery_enabled(&self, enabled: bool) {
self.factory.set_api_key_model_discovery_enabled(enabled);
}
codex-rs/core/src/thread_manager_tests.rs
+73 -6
View File
@@ -2248,6 +2248,8 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new()));
let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new()));
let mcp_loaded = Arc::new(std::sync::atomic::AtomicUsize::new(0));
// This fixture counts explicit MCP projections, excluding background websocket prewarming.
config.model_provider.supports_websockets = false;
let recorder = Arc::new(InitialDataRecorder {
lifecycle_observed: Arc::clone(&lifecycle_observed),
mcp_observed: Arc::clone(&mcp_observed),
@@ -3432,8 +3434,10 @@ async fn metadata_update_without_result_reads_only_when_the_caller_needs_the_thr
);
}
#[test_case::test_case(true; "disk_cache")]
#[test_case::test_case(false; "without_cache")]
#[tokio::test]
async fn new_uses_active_provider_for_model_refresh() {
async fn new_uses_active_provider_for_model_refresh(with_cache: bool) {
let server = MockServer::start().await;
let models_mock = mount_models_once(&server, ModelsResponse { models: vec![] }).await;
@@ -3453,7 +3457,11 @@ async fn new_uses_active_provider_for_model_refresh() {
&config,
auth_manager.clone(),
TEST_INSTALLATION_ID.to_string(),
Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())),
Arc::new(if with_cache {
ModelCatalogSettings::disk(config.codex_home.to_path_buf())
} else {
ModelCatalogSettings::default()
}),
),
crate::CodexAppsToolsCache::default(),
SessionSource::Exec,
@@ -3476,6 +3484,16 @@ async fn new_uses_active_provider_for_model_refresh() {
.await;
assert_eq!(models_mock.requests().len(), 1);
let mut thread_config = config.clone();
thread_config.model = Some("remote-search-model".to_string());
thread_config.model_provider.request_max_retries = Some(1);
thread_config.model_provider.stream_max_retries = Some(1);
thread_config
.model_provider
.http_headers
.get_or_insert_default()
.insert("X-OpenAI-Product-Sku".to_string(), "test-product".into());
// Thread startup must select the originator catalog before resolving model metadata.
for originator in [
"codex_work_cca",
@@ -3484,24 +3502,32 @@ async fn new_uses_active_provider_for_model_refresh() {
"codex_work_mobile",
"chatgpt_cca",
] {
let mut model = codex_models_manager::model_info::model_info_from_slug("gpt-5.5");
let mut model =
codex_models_manager::model_info::model_info_from_slug("remote-search-model");
model.description = Some(originator.to_string());
model.supports_search_tool = true;
model.tool_mode = Some(codex_protocol::openai_models::ToolMode::CodeModeOnly);
model.use_responses_lite = true;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/models"))
.and(wiremock::matchers::header("originator", originator))
.and(wiremock::matchers::header(
"X-OpenAI-Product-Sku",
"test-product",
))
.respond_with(
wiremock::ResponseTemplate::new(200).set_body_json(ModelsResponse {
models: vec![model],
}),
)
.expect(1)
.expect(if with_cache { 1..=1 } else { 1..=2 })
.mount(&server)
.await;
let started = manager
.start_thread(StartThreadOptions {
metrics_service_name: Some(originator.to_string()),
environments: Some(Vec::new()),
..StartThreadOptions::new(config.clone())
..StartThreadOptions::new(thread_config.clone())
})
.await
.expect("start scoped thread");
@@ -3510,9 +3536,50 @@ async fn new_uses_active_provider_for_model_refresh() {
.session
.services
.models_manager
.get_model_info("gpt-5.5", &config.to_models_manager_config())
.get_model_info(
"remote-search-model",
&thread_config.to_models_manager_config(),
)
.await;
assert_eq!(model.description.as_deref(), Some(originator));
assert!(model.supports_search_tool);
let child = manager
.state
.spawn_child_thread(
StartThreadOptions {
session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: started.thread_id,
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: None,
})),
environments: Some(Vec::new()),
..StartThreadOptions::new(thread_config.clone())
},
started
.thread
.session
.services
.local_agent_runtime
.control(started.thread.session.session_id()),
Some(started.thread_id),
/*inherited_exec_policy*/ None,
)
.await
.expect("spawn child with the parent's provider headers");
let (turn, _) = child
.thread
.session
.new_turn_with_sub_id(
"child-turn".to_string(),
SessionSettingsUpdate::default(),
Default::default(),
)
.await
.expect("prepare child's first turn");
assert_eq!(turn.model_info().as_ref(), &model);
}
let legacy = manager
.start_thread(StartThreadOptions {
codex-rs/core/tests/suite/injected_models_cache.rs
+61 -7
View File
@@ -129,14 +129,14 @@ impl ModelsCache for TestModelsCache {
#[derive(Debug)]
struct TestModelsEndpoint {
models: Vec<ModelInfo>,
models: Mutex<Vec<ModelInfo>>,
fetch_count: AtomicUsize,
}
impl TestModelsEndpoint {
fn new(models: Vec<ModelInfo>) -> Arc<Self> {
Arc::new(Self {
models,
models: Mutex::new(models),
fetch_count: AtomicUsize::new(0),
})
}
@@ -163,7 +163,7 @@ impl ModelsEndpointClient for TestModelsEndpoint {
Box::pin(async move {
self.fetch_count.fetch_add(1, Ordering::SeqCst);
Ok(ModelsEndpointResponse {
models: self.models.clone(),
models: self.models.lock().expect("model response lock").clone(),
etag: None,
identity: self.identity().expect("test endpoint identity"),
})
@@ -424,6 +424,57 @@ async fn threads_follow_provider_changes_and_preserve_host_cache_policy(
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn initialization_preserves_a_concurrent_refresh() -> Result<()> {
let (store_entered, entered) = oneshot::channel();
let (release_store, release) = oneshot::channel();
let cache = Arc::new(TestModelsCache {
entry: None,
load_error: false,
stored_entries: Mutex::new(Vec::new()),
store_gate: Mutex::new(Some(StoreGate {
entered: store_entered,
release,
})),
});
let mut model = remote_model("gpt-5.4");
model.default_reasoning_summary = ReasoningSummary::Concise;
let endpoint = TestModelsEndpoint::new(vec![model.clone()]);
let manager = models_manager(cache, endpoint.clone());
let initialization = tokio::spawn({
let manager = manager.clone();
async move {
manager
.ensure_initialized(codex_core::test_support::default_http_client_factory())
.await;
}
});
timeout(Duration::from_secs(/*secs*/ 10), entered).await??;
// A normal refresh must proceed while the initial load is waiting on storage.
model.default_reasoning_summary = ReasoningSummary::Detailed;
*endpoint.models.lock().expect("model response lock") = vec![model.clone()];
timeout(
Duration::from_secs(/*secs*/ 10),
manager.list_models(
RefreshStrategy::Online,
codex_core::test_support::default_http_client_factory(),
),
)
.await
.expect("refresh must not wait for an initialization cache write");
assert_eq!(manager.get_remote_models().await, vec![model.clone()]);
release_store.send(()).expect("release initial cache write");
timeout(Duration::from_secs(/*secs*/ 10), initialization).await??;
assert_eq!(manager.get_remote_models().await, vec![model.clone()]);
assert_eq!(
manager.refreshed_model_catalog().as_deref(),
Some(std::slice::from_ref(&model))
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn() -> Result<()> {
let server = wiremock::MockServer::start().await;
@@ -517,12 +568,15 @@ async fn account_switch_during_cache_store_preserves_new_catalog_for_next_turn()
},
)
.await;
manager
.list_models(
timeout(
Duration::from_secs(/*secs*/ 10),
manager.list_models(
RefreshStrategy::Online,
codex_core::test_support::default_http_client_factory(),
)
.await;
),
)
.await
.expect("new account catalog must load while the old account's cache write is paused");
assert_eq!(manager.get_remote_models().await, vec![new_model.clone()]);
release_store.send(()).unwrap();
timeout(Duration::from_secs(10), old_refresh).await??;
codex-rs/core/tests/suite/items.rs
+2 -1
View File
@@ -309,10 +309,11 @@ async fn web_search_item_is_emitted() -> anyhow::Result<()> {
let server = start_mock_server().await;
let TestCodex { codex, .. } = test_codex()
let test = test_codex()
.with_history_mode(ThreadHistoryMode::Paginated)
.build(&server)
.await?;
let codex = test.codex.clone();
let web_search_added = ev_web_search_call_added_partial("web-search-1", "in_progress");
let web_search_done = ev_web_search_call_done("web-search-1", "completed", "weather seattle");
codex-rs/core/tests/suite/models_cache_auth.rs
+11 -2
View File
@@ -25,6 +25,8 @@ use codex_protocol::openai_models::ToolMode;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
@@ -385,8 +387,10 @@ async fn auth_and_provider_switches_do_not_reuse_chatgpt_catalog() -> Result<()>
Ok(())
}
#[test_case(SessionSource::VSCode; "root")]
#[test_case(SessionSource::SubAgent(SubAgentSource::Review); "child")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn host_model_info_uses_the_thread_catalog() -> Result<()> {
async fn host_model_info_uses_the_thread_catalog(source: SessionSource) -> Result<()> {
let slug = "provider-default";
let server = MockServer::start().await;
let test = test_codex()
@@ -403,7 +407,11 @@ async fn host_model_info_uses_the_thread_catalog() -> Result<()> {
.model_catalog
.take()
.expect("runtime catalog");
let mut model = catalog.models.into_iter().find(|m| m.slug == slug).unwrap();
let mut model = catalog
.models
.into_iter()
.find(|m| m.slug == slug)
.expect("runtime catalog contains the selected model");
model.comp_hash = Some("thread-revision".into());
let mut next_model = model.clone();
next_model.slug = "after-update".into();
@@ -422,6 +430,7 @@ async fn host_model_info_uses_the_thread_catalog() -> Result<()> {
options.config.model = None;
options.config.model_context_window = Some(32_000);
options.metrics_service_name = Some("codex_work_cca".into());
options.session_source = Some(source);
let selected = test.thread_manager.resolve_startup_model(&options).await?;
assert_eq!(selected.slug, slug);
assert_eq!(selected.comp_hash.as_deref(), Some("thread-revision"));
codex-rs/core/tests/suite/models_cache_ttl.rs
+71 -13
View File
@@ -1,3 +1,4 @@
use codex_core::StartThreadOptions;
use codex_core::TurnInputRequest;
use core_test_support::test_codex::local_requests;
use std::path::Path;
@@ -10,6 +11,7 @@ use chrono::Utc;
use codex_core::config::Constrained;
use codex_login::CodexAuth;
use codex_models_manager::client_version_to_whole;
use codex_models_manager::manager::ModelCatalogSettings;
use codex_models_manager::manager::RefreshStrategy;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::config_types::CollaborationMode;
@@ -28,6 +30,8 @@ use codex_protocol::openai_models::TruncationPolicyConfig;
use codex_protocol::openai_models::default_input_modalities;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
@@ -46,7 +50,11 @@ use pretty_assertions::assert_eq;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path_regex;
const ETAG: &str = "\"models-etag-ttl\"";
const CACHE_FILE: &str = "models_cache.json";
@@ -55,11 +63,19 @@ const VERSIONED_MODEL: &str = "codex-test-versioned";
const MISSING_VERSION_MODEL: &str = "codex-test-missing-version";
const DIFFERENT_VERSION_MODEL: &str = "codex-test-different-version";
#[test_case::test_case(true; "expired_cache")]
#[test_case::test_case(false; "disabled_cache")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> {
async fn guardian_reused_reviewer_avoids_stale_catalog_lookup(with_cache: bool) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = MockServer::start().await;
let websocket_fallback = Mock::given(method("GET"))
.and(path_regex(".*/responses$"))
.respond_with(ResponseTemplate::new(426))
.expect(1)
.mount_as_scoped(&server)
.await;
let bundled = codex_models_manager::bundled_models_response()?;
let catalog = ModelsResponse {
models: bundled
@@ -77,23 +93,30 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> {
.with_model("gpt-5.5");
builder = builder.with_config(|config| {
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config.approvals_reviewer = ApprovalsReviewer::AutoReview;
config.approvals_reviewer = ApprovalsReviewer::User;
config.model_provider.request_max_retries = Some(0);
config.model_provider.supports_websockets = true;
});
if !with_cache {
builder = builder.with_model_catalog_settings(ModelCatalogSettings::default());
}
let test = builder.build_with_auto_env(&server).await?;
test.thread_manager
.get_models_manager()
.list_models(
RefreshStrategy::OnlineIfUncached,
codex_core::test_support::default_http_client_factory(),
)
.await;
assert_eq!(initial_models.requests().len(), 1);
rewrite_cache_timestamp(
&test.config.codex_home.join(CACHE_FILE),
Utc::now() - chrono::Duration::hours(1),
// Complete the parent's fallback before Guardian can prewarm its own session.
tokio::time::timeout(
std::time::Duration::from_secs(5),
websocket_fallback.wait_until_satisfied(),
)
.await?;
if with_cache {
rewrite_cache_timestamp(
&test.config.codex_home.join(CACHE_FILE),
Utc::now() - chrono::Duration::hours(1),
)
.await?;
} else {
assert!(!test.config.codex_home.join(CACHE_FILE).exists());
}
let repeat_models = responses::mount_models_once_with_etag(&server, catalog, ETAG).await;
let deny = r#"{"risk_level":"high","user_authorization":"low","outcome":"deny"}"#;
let mut sequence = Vec::new();
@@ -121,7 +144,22 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> {
ev_completed("parent-done"),
]));
let response_mock = mount_sse_sequence(&server, sequence).await;
test.submit_text_turn("Review both shell requests.").await?;
test.codex
.start_or_steer_turn(
TurnInputRequest::user_input(vec![UserInput::Text {
text: "Review both shell requests.".into(),
text_elements: Vec::new(),
}])
.with_thread_settings(ThreadSettingsOverrides {
approvals_reviewer: Some(ApprovalsReviewer::AutoReview),
..Default::default()
}),
)
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let requests = response_mock.requests();
assert_eq!(
@@ -151,6 +189,26 @@ async fn guardian_reused_reviewer_avoids_stale_catalog_lookup() -> Result<()> {
Some(guardian_thread.as_str()),
"the second approval must reuse the Guardian reviewer"
);
// Parallel reviewers inherit HTTP fallback as well as Guardian's retry overrides.
let fork_reviewer = || async {
let mut config = test.config.clone();
config.model_provider.request_max_retries = Some(1);
config.model_provider.stream_max_retries = Some(1);
config.model_provider.supports_websockets = false;
test.thread_manager
.fork_internal_session(
test.session_configured.session_id.into(),
StartThreadOptions {
session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)),
..StartThreadOptions::new(config)
},
Vec::new(),
)
.await
};
let (first, second) = tokio::try_join!(fork_reviewer(), fork_reviewer())?;
assert_ne!(first.thread_id, second.thread_id);
assert_eq!(repeat_models.requests().len(), 0);
Ok(())
}
codex-rs/core/tests/suite/remote_models.rs
+246
View File
@@ -1,11 +1,17 @@
#![cfg(not(target_os = "windows"))]
use anyhow::Result;
use codex_core::StartThreadOptions;
use codex_core::TurnInputRequest;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::built_in_model_providers;
use codex_models_manager::bundled_models_response;
use codex_models_manager::cache::ModelsCache;
use codex_models_manager::cache::ModelsCacheEntry;
use codex_models_manager::cache::ModelsCacheError;
use codex_models_manager::cache::ModelsCacheFuture;
use codex_models_manager::manager::ModelCatalogSettings;
use codex_models_manager::manager::RefreshStrategy;
use codex_models_manager::manager::SharedModelsManager;
use codex_models_manager::model_info::BASE_INSTRUCTIONS;
@@ -32,6 +38,7 @@ use core_test_support::load_default_config_for_test;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_function_call_with_namespace;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::mount_models_once;
use core_test_support::responses::mount_models_once_with_delay;
@@ -49,6 +56,10 @@ use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tempfile::TempDir;
use test_case::test_case;
use tokio::time::Duration;
@@ -64,6 +75,241 @@ use wiremock::matchers::path;
const REMOTE_MODEL_SLUG: &str = "codex-test";
enum ChildCatalogCache {
Warm,
Cold,
Disabled,
}
#[test_case(ChildCatalogCache::Warm; "warm_cache")]
#[test_case(ChildCatalogCache::Cold; "cold_cache")]
#[test_case(ChildCatalogCache::Disabled; "without_cache")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn child_loads_remote_tool_search_with_thread_provider_headers(
cache: ChildCatalogCache,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = MockServer::start().await;
let cache_home = TempDir::new()?;
let catalog_settings = match cache {
ChildCatalogCache::Warm | ChildCatalogCache::Cold => {
ModelCatalogSettings::disk(cache_home.path().to_path_buf())
}
ChildCatalogCache::Disabled => ModelCatalogSettings::default(),
};
let mut model = test_remote_model(
REMOTE_MODEL_SLUG,
ModelVisibility::List,
/*priority*/ 0,
);
model.supports_search_tool = true;
let _models = Mock::given(method("GET"))
.and(path("/v1/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(ModelsResponse {
models: vec![model],
}))
.expect(match cache {
ChildCatalogCache::Warm => 2..=2,
ChildCatalogCache::Cold | ChildCatalogCache::Disabled => 2..=3,
})
.mount_as_scoped(&server)
.await;
let mut test = test_codex()
.with_model(REMOTE_MODEL_SLUG)
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_model_catalog_settings(catalog_settings)
.with_config(|config| {
config
.features
.enable(Feature::Collab)
.expect("enable collaboration");
config
.features
.disable(Feature::MultiAgentV2)
.expect("disable multi-agent v2");
config.agent_max_depth = 2;
})
.build_with_auto_env(&server)
.await?;
test.codex.shutdown_and_wait().await?;
test.config
.model_provider
.http_headers
.get_or_insert_default()
.insert("X-OpenAI-Product-Sku".to_string(), "test-product".into());
let root = test
.thread_manager
.start_thread(StartThreadOptions {
metrics_service_name: Some("codex_work_cca".to_string()),
..test.start_thread_options().await
})
.await?;
test.codex = root.thread;
test.session_configured = root.session_configured;
if matches!(cache, ChildCatalogCache::Cold) {
tokio::fs::remove_file(cache_home.path().join("models_cache.codex_work_cca.json")).await?;
}
assert_child_retains_tool_search(&mut test, &server).await
}
#[derive(Debug, Default)]
struct FailingWritesCache {
entry: Mutex<Option<ModelsCacheEntry>>,
reject_writes: AtomicBool,
}
impl ModelsCache for FailingWritesCache {
fn load<'a>(
&'a self,
_client_version: &'a str,
) -> ModelsCacheFuture<'a, Result<Option<ModelsCacheEntry>, ModelsCacheError>> {
Box::pin(async { Ok(self.entry.lock().expect("cache lock").clone()) })
}
fn store<'a>(
&'a self,
entry: &'a ModelsCacheEntry,
) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> {
Box::pin(async move {
if self.reject_writes.load(Ordering::SeqCst) {
return Err(ModelsCacheError::new("injected cache write failure"));
}
*self.entry.lock().expect("cache lock") = Some(entry.clone());
Ok(())
})
}
fn refresh_ttl<'a>(
&'a self,
_client_version: &'a str,
_identity: &'a str,
_etag: &'a str,
) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> {
Box::pin(async { Err(ModelsCacheError::new("unexpected TTL refresh")) })
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn child_preserves_parent_catalog_after_cache_write_failure() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = MockServer::start().await;
let cache = Arc::new(FailingWritesCache::default());
let mut model = test_remote_model(
REMOTE_MODEL_SLUG,
ModelVisibility::List,
/*priority*/ 0,
);
model.supports_search_tool = false;
let initial = mount_models_once(
&server,
ModelsResponse {
models: vec![model.clone()],
},
)
.await;
let mut test = test_codex()
.with_model(REMOTE_MODEL_SLUG)
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_model_catalog_settings(ModelCatalogSettings::with_cache(cache.clone()))
.with_config(|config| {
config
.features
.enable(Feature::Collab)
.expect("enable collaboration");
config
.features
.disable(Feature::MultiAgentV2)
.expect("disable multi-agent v2");
config.agent_max_depth = 2;
})
.build_with_auto_env(&server)
.await?;
assert_eq!(initial.requests().len(), 1);
model.supports_search_tool = true;
let updated = mount_models_once(
&server,
ModelsResponse {
models: vec![model],
},
)
.await;
cache.reject_writes.store(/*val*/ true, Ordering::SeqCst);
let manager = test.thread_manager.get_models_manager();
manager
.list_models(RefreshStrategy::Online, test.config.http_client_factory())
.await;
let model_config = test.config.to_models_manager_config();
let before = manager
.get_model_info(REMOTE_MODEL_SLUG, &model_config)
.await;
assert!(before.supports_search_tool);
assert!(
!cache
.entry
.lock()
.expect("cache lock")
.as_ref()
.expect("older cached catalog")
.models[0]
.supports_search_tool
);
assert_child_retains_tool_search(&mut test, &server).await?;
assert_eq!(
manager
.get_model_info(REMOTE_MODEL_SLUG, &model_config)
.await,
before
);
assert_eq!(updated.requests().len(), 1);
Ok(())
}
async fn assert_child_retains_tool_search(test: &mut TestCodex, server: &MockServer) -> Result<()> {
let responses = mount_sse_sequence(
server,
vec![
sse(vec![
ev_response_created("parent-spawn"),
ev_function_call_with_namespace(
"spawn-child",
"multi_agent_v1",
"spawn_agent",
&json!({"message": "Say done."}).to_string(),
),
ev_completed("parent-spawn"),
]),
sse(vec![ev_response_created("done-1"), ev_completed("done-1")]),
sse(vec![ev_response_created("done-2"), ev_completed("done-2")]),
],
)
.await;
let mut created = test.thread_manager.subscribe_thread_created();
test.submit_turn("Spawn a child.").await?;
let child_id = timeout(Duration::from_secs(10), created.recv()).await??;
let child = test.thread_manager.get_thread(child_id).await?;
wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await;
let requests = responses.requests();
let child_request = requests
.iter()
.find(|request| request.header("thread-id") == Some(child_id.to_string()))
.expect("child's first Responses request");
for (thread, request) in [("root", &requests[0]), ("child", child_request)] {
assert_eq!(request.body_json()["model"], REMOTE_MODEL_SLUG);
assert!(
request.body_json()["tools"]
.as_array()
.expect("request tools array")
.iter()
.any(|tool| { tool["type"] == "tool_search" }),
"{thread} must retain deferred tool search"
);
}
child.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn disabled_update_plan_preserves_custom_catalog_instructions() -> Result<()> {
skip_if_no_network!(Ok(()));
codex-rs/core/tests/suite/wake_reservation.rs
+4
View File
@@ -49,6 +49,10 @@ struct GatedModelsManager {
}
impl ModelsManager for GatedModelsManager {
fn ensure_initialized(&self, factory: HttpClientFactory) -> ModelsManagerFuture<'_, ()> {
self.inner.ensure_initialized(factory)
}
fn model_catalog_with_origins(
&self,
strategy: RefreshStrategy,
codex-rs/models-manager/src/cache_identity_tests.rs
+4 -4
View File
@@ -23,12 +23,12 @@ async fn mismatched_and_legacy_cache_entries_fetch_the_current_catalog() {
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
)),
);
manager
.ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY)
.await;
assert_eq!(
manager
.raw_model_catalog(
RefreshStrategy::OnlineIfUncached,
DEFAULT_HTTP_CLIENT_FACTORY
)
.raw_model_catalog(RefreshStrategy::InMemory, DEFAULT_HTTP_CLIENT_FACTORY)
.await
.models,
vec![current.clone()]
codex-rs/models-manager/src/catalog.rs
+4
View File
@@ -29,6 +29,8 @@ impl<'a> SelectedCatalog<'a> {
pub(crate) struct CatalogState {
pub entry: ModelsCacheEntry,
pub provider_model_ids: HashSet<String>,
// Distinguishes a successfully loaded empty catalog from the initial bundled seed.
pub initialized: bool,
}
impl CatalogState {
@@ -36,6 +38,7 @@ impl CatalogState {
Self {
entry,
provider_model_ids: HashSet::new(),
initialized: false,
}
}
@@ -49,6 +52,7 @@ impl CatalogState {
Self {
entry,
provider_model_ids,
initialized: true,
}
}
codex-rs/models-manager/src/manager.rs
+107 -18
View File
@@ -129,6 +129,13 @@ impl fmt::Display for RefreshStrategy {
}
}
/// Initialization preserves a loaded snapshot; explicit refreshes follow their requested policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CatalogLoadMode {
Initialize,
Refresh(RefreshStrategy),
}
type SharedModelsEndpointClient = Arc<dyn ModelsEndpointClient>;
/// A successful in-process provider response, before bundled fallback is merged in.
@@ -140,6 +147,14 @@ pub struct RefreshedModelCatalog {
/// Coordinates model discovery plus cached metadata on disk.
pub trait ModelsManager: fmt::Debug + Send + Sync {
/// Initialize a newly selected catalog without replacing an already-loaded snapshot.
/// Stateful managers must check readiness for their current provider and credentials;
/// a successful empty catalog is initialized too. Static catalogs need no work.
fn ensure_initialized(
&self,
http_client_factory: HttpClientFactory,
) -> ModelsManagerFuture<'_, ()>;
/// Observe successful provider refreshes and invalidations, never cache-only reads.
fn model_catalog_updates(&self) -> Option<watch::Receiver<Option<Arc<RefreshedModelCatalog>>>> {
None
@@ -490,6 +505,25 @@ impl StaticModelsManager {
}
impl ModelsManager for OpenAiModelsManager {
fn ensure_initialized(
&self,
http_client_factory: HttpClientFactory,
) -> ModelsManagerFuture<'_, ()> {
Box::pin(async move {
// Resolve lazy credentials before checking the selected catalog's identity.
self.should_refresh_models().await;
if self.catalog_is_initialized(&*self.remote_models.read().await) {
return;
}
if let Err(err) = self
.load_catalog(CatalogLoadMode::Initialize, &http_client_factory)
.await
{
error!("failed to initialize model catalog: {err}");
}
})
}
fn model_catalog_updates(&self) -> Option<watch::Receiver<Option<Arc<RefreshedModelCatalog>>>> {
Some(self.catalog_updates.subscribe())
}
@@ -672,6 +706,12 @@ impl ModelsManager for OpenAiModelsManager {
}
impl OpenAiModelsManager {
fn catalog_is_initialized(&self, catalog: &CatalogState) -> bool {
catalog.initialized
&& catalog.entry.identity.is_some()
&& catalog.entry.identity == self.endpoint_client.identity()
}
/// Selects the catalog available under the current provider and auth identity.
fn models_for_entry<'a>(&self, entry: &'a ModelsCacheEntry) -> SelectedCatalog<'a> {
if !self.api_key_discovery_disabled()
@@ -719,6 +759,18 @@ impl OpenAiModelsManager {
&self,
refresh_strategy: RefreshStrategy,
http_client_factory: &HttpClientFactory,
) -> CoreResult<()> {
self.load_catalog(
CatalogLoadMode::Refresh(refresh_strategy),
http_client_factory,
)
.await
}
async fn load_catalog(
&self,
mode: CatalogLoadMode,
http_client_factory: &HttpClientFactory,
) -> CoreResult<()> {
let _refresh_guard = self.catalog_source.refresh_lock().await;
// API-key discovery must be enabled and supported before reusing a remote catalog.
@@ -727,34 +779,49 @@ impl OpenAiModelsManager {
if self.api_key_discovery_disabled() {
return Ok(());
}
if !self.should_refresh_models().await {
let should_refresh = self.should_refresh_models().await;
if mode == CatalogLoadMode::Initialize
&& self.catalog_is_initialized(&*self.remote_models.read().await)
{
return Ok(());
}
if !should_refresh {
if matches!(
refresh_strategy,
RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached
mode,
CatalogLoadMode::Initialize
| CatalogLoadMode::Refresh(
RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached
)
) {
self.try_load_cache().await;
self.try_load_cache(mode).await;
}
return Ok(());
}
match refresh_strategy {
RefreshStrategy::InMemory => Ok(()),
RefreshStrategy::Offline => {
self.try_load_cache().await;
match mode {
CatalogLoadMode::Refresh(RefreshStrategy::InMemory) => Ok(()),
CatalogLoadMode::Refresh(RefreshStrategy::Offline) => {
self.try_load_cache(mode).await;
Ok(())
}
RefreshStrategy::OnlineIfUncached => {
if self.try_load_cache().await {
CatalogLoadMode::Initialize
| CatalogLoadMode::Refresh(RefreshStrategy::OnlineIfUncached) => {
if self.try_load_cache(mode).await {
return Ok(());
}
self.fetch_and_update_models(http_client_factory).await
self.fetch_and_update_models(http_client_factory, mode)
.await
}
CatalogLoadMode::Refresh(RefreshStrategy::Online) => {
self.fetch_and_update_models(http_client_factory, mode)
.await
}
RefreshStrategy::Online => self.fetch_and_update_models(http_client_factory).await,
}
}
async fn fetch_and_update_models(
&self,
http_client_factory: &HttpClientFactory,
mode: CatalogLoadMode,
) -> CoreResult<()> {
let client_version = crate::client_version_to_whole();
let request_identity = self.endpoint_client.identity();
@@ -787,6 +854,7 @@ impl OpenAiModelsManager {
let CatalogState {
entry: current,
provider_model_ids,
initialized,
} = &mut *current;
if request_identity != self.endpoint_client.identity() {
return Err(error);
@@ -797,6 +865,7 @@ impl OpenAiModelsManager {
provider_model_ids.clear();
current.etag = None;
current.identity.clone_from(&request_identity);
*initialized = true;
self.catalog_updates.send_replace(None);
}
if let Some(identity) = request_identity
@@ -823,7 +892,10 @@ impl OpenAiModelsManager {
return Err(error);
}
};
if Some(&identity) != self.endpoint_client.identity().as_ref() {
if Some(&identity) != self.endpoint_client.identity().as_ref()
|| (mode == CatalogLoadMode::Initialize
&& self.catalog_is_initialized(&*self.remote_models.read().await))
{
return Ok(());
}
let refreshed = Arc::new(RefreshedModelCatalog {
@@ -843,7 +915,7 @@ impl OpenAiModelsManager {
error!("failed to write models cache: {err}");
}
let mut current = self.remote_models.write().await;
if self.apply_remote_models(&mut current, entry) {
if self.apply_remote_models(&mut current, entry, mode) {
self.catalog_updates.send_replace(Some(refreshed));
}
Ok(())
@@ -878,8 +950,17 @@ impl OpenAiModelsManager {
}
/// Publish only while the request identity still matches, including after async storage.
fn apply_remote_models(&self, current: &mut CatalogState, entry: ModelsCacheEntry) -> bool {
if entry.identity != self.endpoint_client.identity() {
fn apply_remote_models(
&self,
current: &mut CatalogState,
entry: ModelsCacheEntry,
mode: CatalogLoadMode,
) -> bool {
// A refresh may finish while initialization awaits cache or network I/O.
// Check under the write lock so initialization cannot replace its loaded snapshot.
if entry.identity != self.endpoint_client.identity()
|| (mode == CatalogLoadMode::Initialize && self.catalog_is_initialized(current))
{
return false;
}
let mut catalog = CatalogState::provider(entry);
@@ -912,7 +993,7 @@ impl OpenAiModelsManager {
}
/// Attempt to satisfy the refresh from the cache when it matches the provider and TTL.
async fn try_load_cache(&self) -> bool {
async fn try_load_cache(&self, mode: CatalogLoadMode) -> bool {
let Some(cache) = self.settings.cache.as_ref() else {
return false;
};
@@ -961,11 +1042,19 @@ impl OpenAiModelsManager {
return false;
}
let mut current = self.remote_models.write().await;
self.apply_remote_models(&mut current, cache_entry)
(mode == CatalogLoadMode::Initialize && self.catalog_is_initialized(&current))
|| self.apply_remote_models(&mut current, cache_entry, mode)
}
}
impl ModelsManager for StaticModelsManager {
fn ensure_initialized(
&self,
_http_client_factory: HttpClientFactory,
) -> ModelsManagerFuture<'_, ()> {
Box::pin(std::future::ready(()))
}
fn model_catalog_with_origins(
&self,
_refresh_strategy: RefreshStrategy,
codex-rs/models-manager/src/manager_tests.rs
+24
View File
@@ -566,6 +566,9 @@ async fn injected_cache_write_error_does_not_fail_remote_refresh() {
.await;
assert_eq!(catalog.models, remote_models);
manager
.ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY)
.await;
assert_eq!(
manager
.model_catalog_with_origins(RefreshStrategy::InMemory, DEFAULT_HTTP_CLIENT_FACTORY)
@@ -581,6 +584,27 @@ async fn injected_cache_write_error_does_not_fail_remote_refresh() {
assert_eq!(endpoint.fetch_count(), 1);
}
#[tokio::test]
async fn initialization_retains_loaded_catalogs_including_empty_responses() {
for models in [
Vec::new(),
vec![remote_model("remote", "Remote", /*priority*/ 0)],
] {
let endpoint = TestModelsEndpoint::new(vec![models.clone()]);
let manager =
OpenAiModelsManager::new_without_cache(endpoint.clone(), /*auth_manager*/ None)
.with_provider_catalog();
manager
.ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY)
.await;
manager
.ensure_initialized(DEFAULT_HTTP_CLIENT_FACTORY)
.await;
assert_eq!(manager.get_remote_models().await, models);
assert_eq!(endpoint.fetch_count(), 1);
}
}
#[tokio::test]
async fn injected_cache_ttl_refresh_preserves_cached_payload() {
let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)];