Add support for choosing a model with Codex. (#9859)

## Description
<!-- Please remember to add your design buddy onto the PR for review, if
it contains any UI changes! -->

WISOTT! Codex will interrupt you with an interactive prompt if you
choose a model that is 5.4 or below so we add some configuration on the
client to set the necessary config to avoid that prompt.

## Testing
<!--
How did you test this change? What automated tests did you add? If you
didn't add any new tests, what's your justification for not adding any?
-->

Tested end to end that all the models work and don't prompt you.

## Agent Mode
- [x] Warp Agent Mode - This PR was created via Warp's AI Agent Mode

Author: abhishek-pandya

app/src/ai/agent_sdk/driver.rs

@@ -238,7 +238,7 @@ pub struct AgentDriverOptions {
     /// Selected execution harness for this run.
     pub selected_harness: Harness,
     /// Model ID for the selected harness. Only used for non-Oz harnesses.
-    pub harness_model_id: Option<String>,
+    pub third_party_harness_model_id: Option<String>,
     /// Whether to skip end-of-run snapshot upload.
     pub snapshot_disabled: Option<bool>,
     /// End-of-run snapshot upload timeout override.
@@ -308,6 +308,7 @@ pub struct AgentDriver {
     /// conversation's `parent_agent_id` field at register time so the
     /// streamer recognizes the child role in driver-hosted processes.
     parent_run_id: Option<String>,
+    third_party_harness_model_id: Option<String>,
 
     /// Async writer that records `file` declarations for paths the agent creates or edits
     /// via `RequestFileEdits`. `Some` only when `FeatureFlag::OzHandoff` is enabled, the run
@@ -504,7 +505,7 @@ impl AgentDriver {
             cloud_providers,
             environment,
             selected_harness,
-            harness_model_id,
+            third_party_harness_model_id,
             snapshot_disabled,
             snapshot_upload_timeout,
             snapshot_script_timeout,
@@ -559,7 +560,7 @@ impl AgentDriver {
         ));
         env_vars.extend(harness_model_env_vars(
             selected_harness,
-            harness_model_id.as_deref(),
+            third_party_harness_model_id.as_deref(),
         ));
 
         // Signal to third-party harnesses (e.g. Claude Code) that we're in a sandbox
@@ -632,6 +633,7 @@ impl AgentDriver {
                 .unwrap_or(snapshot::DEFAULT_DECLARATIONS_SCRIPT_TIMEOUT),
             run_conversation_id,
             parent_run_id: parent_run_id_for_self,
+            third_party_harness_model_id,
             snapshot_file_writer,
         })
     }
@@ -1614,8 +1616,13 @@ impl AgentDriver {
             }
         };
 
-        let secrets = foreground
-            .spawn(|me, _| Arc::clone(&me.secrets))
+        let (secrets, third_party_harness_model_id) = foreground
+            .spawn(|me, _| {
+                (
+                    Arc::clone(&me.secrets),
+                    me.third_party_harness_model_id.clone(),
+                )
+            })
             .await
             .map_err(|_| AgentDriverError::InvalidRuntimeState)?;
 
@@ -1654,6 +1661,7 @@ impl AgentDriver {
                 resume,
                 &resolved_env_vars,
                 &resolved_mcp_servers,
+                third_party_harness_model_id.as_deref(),
             )?
             .into();
 

app/src/ai/agent_sdk/driver/harness/claude_code.rs

@@ -103,6 +103,7 @@ impl ThirdPartyHarness for ClaudeHarness {
         resume: Option<ResumePayload>,
         resolved_env_vars: &HashMap<OsString, OsString>,
         resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
+        _third_party_harness_model_id: Option<&str>,
     ) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
         // Prepare the environment config files.
         prepare_claude_environment_config(working_dir, resolved_env_vars).map_err(|error| {

app/src/ai/agent_sdk/driver/harness/codex.rs

@@ -87,13 +87,15 @@ impl ThirdPartyHarness for CodexHarness {
         resume: Option<ResumePayload>,
         resolved_env_vars: &HashMap<OsString, OsString>,
         resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
+        third_party_harness_model_id: Option<&str>,
     ) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
         // Prepare the environment config files.
         prepare_codex_environment_config(
             working_dir,
             system_prompt,
             resolved_env_vars,
             resolved_mcp_servers,
+            third_party_harness_model_id,
         )
         .map_err(|error| AgentDriverError::HarnessConfigSetupFailed {
             harness: self.cli_agent().command_prefix().to_owned(),
@@ -434,6 +436,15 @@ const CODEX_TRUST_LEVEL_TRUSTED: &str = "trusted";
 /// Top-level config key codex reads to override the built-in `openai` provider's base URL
 /// (codex `core/src/config/mod.rs`).
 const CODEX_OPENAI_BASE_URL_KEY: &str = "openai_base_url";
+const CODEX_MODEL_KEY: &str = "model";
+/// Target model for the `[notice.model_migrations]` table that suppresses Codex's
+/// "choose a newer model" upgrade prompt at session launch. We stamp this for any
+/// pinned model id (even when it already matches the target) so the unattended
+/// cloud run never blocks on the prompt.
+///
+/// TODO: Ideally, we would make this server-driven so we don't depend on a client
+/// release to change this.
+const CODEX_MODEL_MIGRATIONS_TARGET: &str = "gpt-5.4";
 /// US data-residency endpoint. Our OpenAI keys are issued under a US-residency project,
 /// which rejects requests to the global host with `401 incorrect_hostname`.
 /// TODO(REMOTE-1509): plumb a region-tagged auth secret instead of hardcoding the URL.
@@ -444,6 +455,7 @@ fn prepare_codex_environment_config(
     system_prompt: Option<&str>,
     resolved_env_vars: &HashMap<OsString, OsString>,
     resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
+    third_party_harness_model_id: Option<&str>,
 ) -> Result<()> {
     let home_dir =
         dirs::home_dir().ok_or_else(|| anyhow::anyhow!("could not determine home directory"))?;
@@ -462,6 +474,7 @@ fn prepare_codex_environment_config(
         &codex_dir.join(CODEX_CONFIG_TOML_FILE_NAME),
         working_dir,
         resolved_mcp_servers,
+        third_party_harness_model_id,
     )?;
     Ok(())
 }
@@ -570,10 +583,14 @@ fn resolve_openai_api_key(resolved_env_vars: &HashMap<OsString, OsString>) -> Op
 ///   set the projects to `trusted`.
 /// - base URL: set `openai_base_url = "<US data-residency endpoint>"` so we
 ///   hit the regional host our API keys require.
+/// - model override: when a non-default `third_party_harness_model_id` is
+///   supplied, write the top-level `model` key so Codex pins the chosen model
+///   for new sessions.
 fn prepare_codex_config_toml(
     config_toml_path: &Path,
     working_dir: &Path,
     resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
+    third_party_harness_model_id: Option<&str>,
 ) -> Result<()> {
     let existing = match fs::read_to_string(config_toml_path) {
         Ok(content) => content,
@@ -593,6 +610,7 @@ fn prepare_codex_config_toml(
     })?;
 
     set_codex_openai_base_url(&mut doc, CODEX_OPENAI_BASE_URL);
+    set_codex_model(&mut doc, third_party_harness_model_id);
 
     let canonical = working_dir.canonicalize().with_context(|| {
         format!(
@@ -631,6 +649,45 @@ fn set_codex_openai_base_url(doc: &mut toml_edit::DocumentMut, base_url: &str) {
     doc[CODEX_OPENAI_BASE_URL_KEY] = toml_edit::value(base_url);
 }
 
+fn set_codex_model(doc: &mut toml_edit::DocumentMut, third_party_harness_model_id: Option<&str>) {
+    let Some(model_id) =
+        third_party_harness_model_id.filter(|id| !id.is_empty() && *id != "default")
+    else {
+        return;
+    };
+    doc[CODEX_MODEL_KEY] = toml_edit::value(model_id);
+
+    // Codex's TUI prompts the user to upgrade older models on session launch even when
+    // a `model` key has been pinned. Stamping a migration entry keyed on the chosen
+    // model id suppresses that prompt for the unattended cloud run. We do this
+    // unconditionally rather than enumerating a list of "old" models on the client:
+    // mapping the migration target to itself (e.g. `gpt-5.4 = "gpt-5.4"`) is a no-op
+    // for Codex, and keeping the client free of model-version knowledge means we
+    // don't have to ship a client update every time Anthropic/OpenAI ages out a model.
+    set_codex_model_migration(doc, model_id, CODEX_MODEL_MIGRATIONS_TARGET);
+}
+
+fn set_codex_model_migration(
+    doc: &mut toml_edit::DocumentMut,
+    from_model_id: &str,
+    to_model_id: &str,
+) {
+    if !doc.contains_table("notice") {
+        let mut notice_tbl = toml_edit::Table::new();
+        notice_tbl.set_implicit(true);
+        doc.insert("notice", toml_edit::Item::Table(notice_tbl));
+    }
+    let migrations_tbl = doc["notice"]
+        .as_table_mut()
+        .expect("notice table inserted above")
+        .entry("model_migrations")
+        .or_insert_with(toml_edit::table)
+        .as_table_mut()
+        .expect("model_migrations entry is a table");
+    migrations_tbl.set_implicit(false);
+    migrations_tbl[from_model_id] = toml_edit::value(to_model_id);
+}
+
 /// Return immediate subdirectories of `dir` that contain a `.git`.
 fn find_child_git_repos(dir: &Path) -> Vec<std::path::PathBuf> {
     let Ok(entries) = fs::read_dir(dir) else {

app/src/ai/agent_sdk/driver/harness/codex_tests.rs

@@ -178,7 +178,7 @@ fn prepare_codex_config_toml_writes_fresh_config() {
     let working_dir = tmp.path().join("workspace/proj");
     fs::create_dir_all(&working_dir).unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
 
     let canonical = working_dir.canonicalize().unwrap();
     let key = canonical.to_string_lossy().into_owned();
@@ -202,7 +202,8 @@ fn prepare_codex_config_toml_preserves_unrelated_keys() {
     )
     .unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    // Pass `None` for the model id so the helper preserves the user's existing `model = ...`.
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
 
     let canonical = working_dir.canonicalize().unwrap();
     let key = canonical.to_string_lossy().into_owned();
@@ -225,9 +226,9 @@ fn prepare_codex_config_toml_is_idempotent() {
     let working_dir = tmp.path().join("workspace");
     fs::create_dir_all(&working_dir).unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
     let after_first = fs::read_to_string(&config_path).unwrap();
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
     let after_second = fs::read_to_string(&config_path).unwrap();
 
     assert_eq!(after_first, after_second);
@@ -256,7 +257,7 @@ fn prepare_codex_config_toml_upgrades_untrusted_entry() {
     )
     .unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     assert_eq!(
@@ -275,7 +276,7 @@ fn prepare_codex_config_toml_trusts_multiple_child_repos() {
     fs::create_dir_all(repo_a.join(".git")).unwrap();
     fs::create_dir_all(repo_b.join(".git")).unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     let projects = cfg["projects"].as_table().unwrap();
@@ -303,7 +304,7 @@ fn prepare_codex_config_toml_overwrites_stale_openai_base_url() {
     )
     .unwrap();
 
-    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new()).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     assert_eq!(cfg["openai_base_url"].as_str(), Some(CODEX_OPENAI_BASE_URL));
@@ -327,7 +328,7 @@ fn write_codex_mcp_servers_cli_server() {
             },
         },
     )]);
-    prepare_codex_config_toml(&config_path, &working_dir, &servers).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &servers, None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     let mcp = &cfg["mcp_servers"]["my-mcp"];
@@ -358,7 +359,7 @@ fn write_codex_mcp_servers_sse_server() {
             },
         },
     )]);
-    prepare_codex_config_toml(&config_path, &working_dir, &servers).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &servers, None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     let mcp = &cfg["mcp_servers"]["remote-mcp"];
@@ -384,7 +385,7 @@ fn write_codex_mcp_servers_cli_server_with_cwd() {
             },
         },
     )]);
-    prepare_codex_config_toml(&config_path, &working_dir, &servers).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &servers, None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     let mcp = &cfg["mcp_servers"]["my-mcp"];
@@ -410,13 +411,100 @@ fn write_codex_mcp_servers_cli_server_without_cwd_omits_key() {
             },
         },
     )]);
-    prepare_codex_config_toml(&config_path, &working_dir, &servers).unwrap();
+    prepare_codex_config_toml(&config_path, &working_dir, &servers, None).unwrap();
 
     let cfg = read_codex_config(&config_path);
     let mcp = &cfg["mcp_servers"]["my-mcp"];
     assert!(mcp.get("cwd").is_none());
 }
 
+#[test]
+fn prepare_codex_config_toml_writes_model_when_specified() {
+    // A non-default model id is written to the top-level `model` key so Codex pins it
+    // for new sessions launched from this `~/.codex/config.toml`. Even for the
+    // current target model, we stamp a self-referential migration entry so the
+    // upgrade prompt is suppressed regardless of what the user selected.
+    let tmp = TempDir::new().unwrap();
+    let config_path = tmp.path().join("config.toml");
+    let working_dir = tmp.path().join("workspace");
+    fs::create_dir_all(&working_dir).unwrap();
+
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), Some("gpt-5.5"))
+        .unwrap();
+
+    let cfg = read_codex_config(&config_path);
+    assert_eq!(cfg["model"].as_str(), Some("gpt-5.5"));
+    assert_eq!(
+        cfg["notice"]["model_migrations"]["gpt-5.5"].as_str(),
+        Some(CODEX_MODEL_MIGRATIONS_TARGET),
+    );
+}
+
+#[test]
+fn prepare_codex_config_toml_writes_model_migration_for_older_model() {
+    // For an older model id, the migration entry maps it to the current target
+    // so Codex's "choose a newer model" prompt is suppressed at session launch.
+    let tmp = TempDir::new().unwrap();
+    let config_path = tmp.path().join("config.toml");
+    let working_dir = tmp.path().join("workspace");
+    fs::create_dir_all(&working_dir).unwrap();
+
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), Some("gpt-5.2"))
+        .unwrap();
+
+    let cfg = read_codex_config(&config_path);
+    assert_eq!(cfg["model"].as_str(), Some("gpt-5.2"));
+    assert_eq!(
+        cfg["notice"]["model_migrations"]["gpt-5.2"].as_str(),
+        Some(CODEX_MODEL_MIGRATIONS_TARGET),
+    );
+}
+
+#[test]
+fn prepare_codex_config_toml_skips_model_for_default_sentinel() {
+    // The literal "default" sentinel means "let Codex pick its own default model";
+    // we should NOT write a `model` key (or a migration entry) in that case.
+    let tmp = TempDir::new().unwrap();
+    let config_path = tmp.path().join("config.toml");
+    let working_dir = tmp.path().join("workspace");
+    fs::create_dir_all(&working_dir).unwrap();
+
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), Some("default"))
+        .unwrap();
+
+    let cfg = read_codex_config(&config_path);
+    assert!(
+        cfg.get("model").is_none(),
+        "`model` should not be written for the default sentinel"
+    );
+    assert!(
+        cfg.get("notice").is_none(),
+        "`[notice]` table should not be written without a pinned model id"
+    );
+}
+
+#[test]
+fn prepare_codex_config_toml_skips_model_when_none() {
+    // No model id supplied means the user didn't pick one; we should not write a
+    // `model` key or any `[notice.model_migrations]` entries.
+    let tmp = TempDir::new().unwrap();
+    let config_path = tmp.path().join("config.toml");
+    let working_dir = tmp.path().join("workspace");
+    fs::create_dir_all(&working_dir).unwrap();
+
+    prepare_codex_config_toml(&config_path, &working_dir, &HashMap::new(), None).unwrap();
+
+    let cfg = read_codex_config(&config_path);
+    assert!(
+        cfg.get("model").is_none(),
+        "`model` should not be written when no override is supplied"
+    );
+    assert!(
+        cfg.get("notice").is_none(),
+        "`[notice]` table should not be written without a pinned model id"
+    );
+}
+
 #[test]
 fn find_child_git_repos_returns_only_repo_children() {
     let tmp = TempDir::new().unwrap();

app/src/ai/agent_sdk/driver/harness/gemini.rs

@@ -61,6 +61,7 @@ impl ThirdPartyHarness for GeminiHarness {
         _resume: Option<ResumePayload>,
         _resolved_env_vars: &HashMap<OsString, OsString>,
         _resolved_mcp_servers: &HashMap<String, JSONMCPServer>,
+        _third_party_harness_model_id: Option<&str>,
     ) -> Result<Box<dyn HarnessRunner>, AgentDriverError> {
         // Prepare the environment config files.
         prepare_gemini_environment_config(working_dir, system_prompt).map_err(|error| {