From b7983c2a07182c3c74fbe67283b5893b762dd508 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 16 Jul 2026 18:05:47 +0000 Subject: [PATCH] Restore agent roles when reloading v2 sub-agents (#33657) ## Why Durable v2 sub-agents can be lazily reloaded after their root session resumes. The reload path restored the agent identity but did not reapply the agent's selected role configuration. ## What changed - Recover the agent role from the resumed session source and apply it before reloading the sub-agent. - Preserve the runtime approval policy, approval reviewer, working directory, and permission profile while applying the role. ## Testing Extend the cold-resume integration test to verify that the reloaded worker retains its role-defined instructions, model, provider, reasoning effort, and permission profile. GitOrigin-RevId: a26dbbd2599838ca317326a1223a6a0a9cb2a2ed --- codex-rs/core/src/agent/control/spawn.rs | 46 +++++++++- .../core/tests/suite/multi_agent_resume.rs | 88 +++++++++++++++---- 2 files changed, 113 insertions(+), 21 deletions(-) diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index bfbd506520..391058b34b 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,5 +1,7 @@ use super::residency::is_v2_resident_session_source; use super::*; +use crate::agent::role::apply_role_to_config; +use crate::config::PermissionProfileSnapshot; use codex_extension_api::ExtensionDataInit; const AGENT_NAMES: &str = include_str!("../agent_names.txt"); @@ -247,7 +249,7 @@ impl AgentControl { pub(crate) async fn ensure_v2_agent_loaded( &self, - config: Config, + mut config: Config, thread_id: ThreadId, ) -> CodexResult<()> { let state = self.upgrade()?; @@ -279,13 +281,49 @@ impl AgentControl { if initial_history.get_multi_agent_version() != Some(MultiAgentVersion::V2) { return Err(CodexErr::ThreadNotFound(thread_id)); } + let (session_source, _) = initial_history + .get_resumed_session_sources() + .unwrap_or((stored_source, None)); + if let Some(role_name) = session_source.get_agent_role() { + let runtime_approval_policy = config.permissions.approval_policy.value(); + let runtime_approvals_reviewer = config.approvals_reviewer; + let runtime_cwd = config.cwd.clone(); + let runtime_permission_profile = match config.permissions.active_permission_profile() { + Some(active_permission_profile) => { + PermissionProfileSnapshot::active_with_profile_workspace_roots( + config.permissions.permission_profile().clone(), + active_permission_profile, + config.permissions.profile_workspace_roots().to_vec(), + ) + } + None => PermissionProfileSnapshot::legacy( + config.permissions.permission_profile().clone(), + ), + }; + + apply_role_to_config(&mut config, Some(&role_name)) + .await + .map_err(CodexErr::InvalidRequest)?; + config + .permissions + .approval_policy + .set(runtime_approval_policy) + .map_err(|err| { + CodexErr::InvalidRequest(format!("approval_policy is invalid: {err}")) + })?; + config.approvals_reviewer = runtime_approvals_reviewer; + config.cwd = runtime_cwd; + config + .permissions + .set_permission_profile_from_session_snapshot(runtime_permission_profile) + .map_err(|err| { + CodexErr::InvalidRequest(format!("permission_profile is invalid: {err}")) + })?; + } let residency_slot = self .reserve_v2_residency_slot(&state, &config, Some(thread_id)) .await?; - let (session_source, _) = initial_history - .get_resumed_session_sources() - .unwrap_or((stored_source, None)); let parent_thread_id = initial_history .get_resumed_parent_thread_id() .or(stored_parent_thread_id); diff --git a/codex-rs/core/tests/suite/multi_agent_resume.rs b/codex-rs/core/tests/suite/multi_agent_resume.rs index cbb7dd42b1..a42e14dadd 100644 --- a/codex-rs/core/tests/suite/multi_agent_resume.rs +++ b/codex-rs/core/tests/suite/multi_agent_resume.rs @@ -1,5 +1,8 @@ use anyhow::Result; +use codex_core::config::AgentRoleConfig; use codex_features::Feature; +use codex_protocol::models::PermissionProfile; +use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AgentStatus; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -23,6 +26,10 @@ const INITIAL_PROMPT: &str = "spawn a durable worker"; const INITIAL_TASK: &str = "inspect the repository"; const FOLLOWUP_PROMPT: &str = "continue the durable worker"; const FOLLOWUP_TASK: &str = "inspect the tests too"; +const ROLE_NAME: &str = "durable_worker"; +const ROLE_MODEL: &str = "gpt-5.4"; +const ROLE_MODEL_PROVIDER_ID: &str = "mock"; +const ROLE_DEVELOPER_INSTRUCTIONS: &str = "Keep the durable worker role configuration."; fn decoded_body(request: &wiremock::Request) -> Option> { let is_zstd = request @@ -58,7 +65,10 @@ fn request_has_input_type(request: &wiremock::Request, input_type: &str) -> bool }) } -fn configure_multi_agent_v2(config: &mut codex_core::config::Config) { +fn configure_multi_agent_v2_with_role( + config: &mut codex_core::config::Config, + model_provider_base_url: &str, +) { config .features .enable(Feature::Collab) @@ -67,14 +77,32 @@ fn configure_multi_agent_v2(config: &mut codex_core::config::Config) { .features .enable(Feature::MultiAgentV2) .expect("test config should allow feature update"); + let role_path = config.codex_home.join("durable-worker-role.toml"); + std::fs::write( + &role_path, + format!( + "model = \"{ROLE_MODEL}\"\nmodel_reasoning_effort = \"high\"\ndeveloper_instructions = \"{ROLE_DEVELOPER_INSTRUCTIONS}\"\nsandbox_mode = \"read-only\"\nmodel_provider = \"mock\"\n\n[model_providers.mock]\nname = \"mock\"\nbase_url = \"{model_provider_base_url}\"\nenv_key = \"PATH\"\nwire_api = \"responses\"\n" + ), + ) + .expect("write durable worker role config"); + config.agent_roles.insert( + ROLE_NAME.to_string(), + AgentRoleConfig { + description: Some("Durable worker role".to_string()), + config_file: Some(role_path.to_path_buf()), + nickname_candidates: None, + }, + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup() -> Result<()> { +async fn cold_root_resume_restores_agent_identity_and_role_on_followup() -> Result<()> { let server = start_mock_server().await; let spawn_args = serde_json::to_string(&json!({ "message": INITIAL_TASK, "task_name": "worker", + "agent_type": ROLE_NAME, + "fork_turns": "none", }))?; mount_sse_once_match( &server, @@ -117,8 +145,11 @@ async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup ) .await; - let mut initial_builder = test_codex().with_config(configure_multi_agent_v2); - let initial = initial_builder.build(&server).await?; + let initial_model_provider_base_url = format!("{}/v1", server.uri()); + let mut initial_builder = test_codex().with_config(move |config| { + configure_multi_agent_v2_with_role(config, &initial_model_provider_base_url); + }); + let initial = initial_builder.build_with_auto_env(&server).await?; let root_thread_id = initial.session_configured.thread_id; let home = initial.home.clone(); let rollout_path = initial @@ -158,11 +189,25 @@ async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup } sleep(Duration::from_millis(10)).await; } - assert!( - initial_child_request - .requests() - .iter() - .any(|request| request.body_contains_text(INITIAL_TASK)) + assert!(initial_child_request.requests().iter().any(|request| { + request.body_contains_text(INITIAL_TASK) + && request.body_contains_text(ROLE_DEVELOPER_INSTRUCTIONS) + })); + let initial_worker_config = worker_thread.config_snapshot().await; + let initial_worker_role_config = ( + initial_worker_config.model, + initial_worker_config.model_provider_id, + initial_worker_config.reasoning_effort, + initial_worker_config.permission_profile, + ); + assert_eq!( + initial_worker_role_config, + ( + ROLE_MODEL.to_string(), + ROLE_MODEL_PROVIDER_ID.to_string(), + Some(ReasoningEffort::High), + PermissionProfile::Disabled, + ) ); worker_thread.flush_rollout().await?; initial.codex.flush_rollout().await?; @@ -215,7 +260,10 @@ async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup ) .await; - let mut resume_builder = test_codex().with_config(configure_multi_agent_v2); + let resumed_model_provider_base_url = format!("{}/v1", server.uri()); + let mut resume_builder = test_codex().with_config(move |config| { + configure_multi_agent_v2_with_role(config, &resumed_model_provider_base_url); + }); let resumed = resume_builder.resume(&server, home, rollout_path).await?; assert_eq!( resumed.thread_manager.list_thread_ids().await, @@ -231,17 +279,23 @@ async fn cold_root_resume_restores_agent_identity_and_reloads_target_on_followup resumed.submit_turn(FOLLOWUP_PROMPT).await?; - assert!( - followup_child_request - .requests() - .iter() - .any(|request| request.body_contains_text(FOLLOWUP_TASK)) - ); - resumed + assert!(followup_child_request.requests().iter().any(|request| { + request.body_contains_text(FOLLOWUP_TASK) + && request.body_contains_text(ROLE_DEVELOPER_INSTRUCTIONS) + })); + let reloaded_worker = resumed .thread_manager .get_thread(worker_thread_id) .await .expect("follow-up should lazily reload the original worker"); + let reloaded_worker_config = reloaded_worker.config_snapshot().await; + let reloaded_worker_role_config = ( + reloaded_worker_config.model, + reloaded_worker_config.model_provider_id, + reloaded_worker_config.reasoning_effort, + reloaded_worker_config.permission_profile, + ); + assert_eq!(reloaded_worker_role_config, initial_worker_role_config); Ok(()) }