config: enforce enterprise feature requirements

codex-rs/core/tests/common/test_codex.rs

@@ -11,12 +11,15 @@ use codex_core::ThreadManager;
 use codex_core::built_in_model_providers;
 use codex_core::config::Config;
 use codex_core::features::Feature;
+use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig;
 use codex_protocol::config_types::ServiceTier;
+use codex_protocol::openai_models::ModelsResponse;
 use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::SandboxPolicy;
 use codex_protocol::protocol::SessionConfiguredEvent;
+use codex_protocol::protocol::SessionSource;
 use codex_protocol::user_input::UserInput;
 use serde_json::Value;
 use tempfile::TempDir;
@@ -28,11 +31,13 @@ use crate::responses::output_value_to_text;
 use crate::responses::start_mock_server;
 use crate::streaming_sse::StreamingSseServer;
 use crate::wait_for_event;
+use crate::wait_for_event_match;
 use wiremock::Match;
 use wiremock::matchers::path_regex;
 
 type ConfigMutator = dyn FnOnce(&mut Config) + Send;
 type PreBuildHook = dyn FnOnce(&Path) + Send + 'static;
+const TEST_MODEL_WITH_EXPERIMENTAL_TOOLS: &str = "test-gpt-5.1-codex";
 
 /// A collection of different ways the model can output an apply_patch call
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -128,7 +133,7 @@ impl TestCodexBuilder {
         self.config_mutators.push(Box::new(move |config| {
             config.model_provider.base_url = Some(base_url_clone);
             config.experimental_realtime_ws_model = Some("realtime-test-model".to_string());
-            config.features.enable(Feature::ResponsesWebsockets);
+            let _ = config.features.enable(Feature::ResponsesWebsockets);
         }));
         self.build_with_home_and_base_url(base_url, home, None)
             .await
@@ -172,11 +177,21 @@ impl TestCodexBuilder {
         resume_from: Option<PathBuf>,
     ) -> anyhow::Result<TestCodex> {
         let auth = self.auth.clone();
-        let thread_manager = codex_core::test_support::thread_manager_with_models_provider_and_home(
-            auth.clone(),
-            config.model_provider.clone(),
-            config.codex_home.clone(),
-        );
+        let thread_manager = if let Some(model_catalog) = config.model_catalog.clone() {
+            ThreadManager::new(
+                config.codex_home.clone(),
+                codex_core::test_support::auth_manager_from_auth(auth.clone()),
+                SessionSource::Exec,
+                Some(model_catalog),
+                CollaborationModesConfig::default(),
+            )
+        } else {
+            codex_core::test_support::thread_manager_with_models_provider_and_home(
+                auth.clone(),
+                config.model_provider.clone(),
+                config.codex_home.clone(),
+            )
+        };
         let thread_manager = Arc::new(thread_manager);
 
         let new_conversation = match resume_from {
@@ -232,17 +247,50 @@ impl TestCodexBuilder {
         for mutator in mutators {
             mutator(&mut config);
         }
+        ensure_test_model_catalog(&mut config);
 
         if config.include_apply_patch_tool {
-            config.features.enable(Feature::ApplyPatchFreeform);
+            let _ = config.features.enable(Feature::ApplyPatchFreeform);
         } else {
-            config.features.disable(Feature::ApplyPatchFreeform);
+            let _ = config.features.disable(Feature::ApplyPatchFreeform);
         }
 
         Ok((config, cwd))
     }
 }
 
+fn ensure_test_model_catalog(config: &mut Config) {
+    if config.model.as_deref() != Some(TEST_MODEL_WITH_EXPERIMENTAL_TOOLS)
+        || config.model_catalog.is_some()
+    {
+        return;
+    }
+
+    let bundled_models_path =
+        codex_utils_cargo_bin::find_resource!("../models.json").expect("bundled models.json");
+    let bundled_models_contents =
+        std::fs::read_to_string(bundled_models_path).expect("read bundled models.json");
+    let bundled_models: ModelsResponse =
+        serde_json::from_str(&bundled_models_contents).expect("bundled models.json");
+    let mut model = bundled_models
+        .models
+        .iter()
+        .find(|candidate| candidate.slug == "gpt-5.1-codex")
+        .cloned()
+        .unwrap_or_else(|| panic!("missing bundled model gpt-5.1-codex"));
+    model.slug = TEST_MODEL_WITH_EXPERIMENTAL_TOOLS.to_string();
+    model.display_name = TEST_MODEL_WITH_EXPERIMENTAL_TOOLS.to_string();
+    model.experimental_supported_tools = vec![
+        "test_sync_tool".to_string(),
+        "read_file".to_string(),
+        "grep_files".to_string(),
+        "list_dir".to_string(),
+    ];
+    config.model_catalog = Some(ModelsResponse {
+        models: vec![model],
+    });
+}
+
 pub struct TestCodex {
     pub home: Arc<TempDir>,
     pub cwd: Arc<TempDir>,
@@ -334,8 +382,14 @@ impl TestCodex {
             })
             .await?;
 
-        wait_for_event(&self.codex, |event| {
-            matches!(event, EventMsg::TurnComplete(_))
+        let turn_id = wait_for_event_match(&self.codex, |event| match event {
+            EventMsg::TurnStarted(event) => Some(event.turn_id.clone()),
+            _ => None,
+        })
+        .await;
+        wait_for_event(&self.codex, |event| match event {
+            EventMsg::TurnComplete(event) => event.turn_id == turn_id,
+            _ => false,
         })
         .await;
         Ok(())

codex-rs/core/tests/common/zsh_fork.rs

@@ -24,8 +24,8 @@ impl ZshForkRuntime {
         approval_policy: AskForApproval,
         sandbox_policy: SandboxPolicy,
     ) {
-        config.features.enable(Feature::ShellTool);
-        config.features.enable(Feature::ShellZshFork);
+        let _ = config.features.enable(Feature::ShellTool);
+        let _ = config.features.enable(Feature::ShellZshFork);
         config.zsh_path = Some(self.zsh_path.clone());
         config.main_execve_wrapper_exe = Some(self.main_execve_wrapper_exe.clone());
         config.permissions.allow_login_shell = false;