codex: add model multi-agent system overlay

This commit is contained in:
Ahmed Ibrahim
2026-05-29 15:38:50 -07:00
parent e7e4330e6a
commit f32ec83dc2
19 changed files with 289 additions and 73 deletions
@@ -53,6 +53,7 @@ fn preset_to_info(preset: &ModelPreset, priority: i32) -> ModelInfo {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}
}
@@ -99,6 +99,7 @@ async fn models_client_hits_models_endpoint() {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}],
};
+8 -2
View File
@@ -258,6 +258,7 @@ impl AgentControl {
&config,
session_source,
options.multi_agent_version,
/*model_info*/ None,
)
});
let notification_source = session_source.clone();
@@ -384,8 +385,12 @@ impl AgentControl {
inherited_shell_snapshot: Option<Arc<ShellSnapshot>>,
inherited_exec_policy: Option<Arc<crate::exec_policy::ExecPolicyManager>>,
) -> CodexResult<crate::thread_manager::NewThread> {
let multi_agent_version =
resolve_multi_agent_version(&config, &session_source, options.multi_agent_version);
let multi_agent_version = resolve_multi_agent_version(
&config,
&session_source,
options.multi_agent_version,
/*model_info*/ None,
);
if options.fork_parent_spawn_call_id.is_none() {
return Err(CodexErr::Fatal(
"spawn_agent fork requires a parent spawn call id".to_string(),
@@ -618,6 +623,7 @@ impl AgentControl {
&config,
&session_source,
stored_thread.multi_agent_version,
/*model_info*/ None,
);
if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) = &session_source
&& *depth >= config.agent_max_depth
+1
View File
@@ -544,6 +544,7 @@ impl Codex {
conversation_history
.get_multi_agent_version()
.or(parent_multi_agent_version),
Some(&model_info),
);
if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) = session_source
&& depth >= config.agent_max_depth
+5
View File
@@ -7237,6 +7237,7 @@ async fn build_initial_context_adds_multi_agent_v2_usage_hint_when_selector_is_v
#[tokio::test]
async fn spawned_child_multi_agent_version_follows_parent_system() {
let (_session, turn_context) = make_session_and_context().await;
let mut model_info = turn_context.model_info.clone();
let child_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: ThreadId::new(),
depth: 1,
@@ -7244,16 +7245,20 @@ async fn spawned_child_multi_agent_version_follows_parent_system() {
agent_nickname: None,
agent_role: None,
});
model_info.multi_agent_version = Some(MultiAgentVersion::V1);
let resolved_from_v2_parent = resolve_multi_agent_version(
turn_context.config.as_ref(),
&child_source,
Some(MultiAgentVersion::V2),
Some(&model_info),
);
model_info.multi_agent_version = Some(MultiAgentVersion::V2);
let resolved_from_v1_parent = resolve_multi_agent_version(
turn_context.config.as_ref(),
&child_source,
Some(MultiAgentVersion::V1),
Some(&model_info),
);
assert_eq!(
+12 -9
View File
@@ -110,6 +110,7 @@ pub(crate) fn resolve_multi_agent_version(
config: &Config,
session_source: &SessionSource,
inherited_multi_agent_version: Option<MultiAgentVersion>,
model_info: Option<&ModelInfo>,
) -> Option<MultiAgentVersion> {
if is_guardian_reviewer_source(session_source)
|| matches!(
@@ -119,15 +120,17 @@ pub(crate) fn resolve_multi_agent_version(
{
return None;
}
inherited_multi_agent_version.or_else(|| {
if config.features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if config.features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
})
inherited_multi_agent_version
.or_else(|| model_info.and_then(|model_info| model_info.multi_agent_version))
.or_else(|| {
if config.features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if config.features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
})
}
impl TurnContext {
+1
View File
@@ -922,6 +922,7 @@ impl ThreadManager {
&config,
&session_source,
history.get_multi_agent_version(),
/*model_info*/ None,
);
let interrupted_marker =
InterruptedTurnHistoryMarker::from_config(&config, multi_agent_version);
@@ -886,6 +886,7 @@ async fn resolved_multi_agent_version_controls_runtime_behavior() {
turn.config.as_ref(),
&turn.session_source,
/*inherited_multi_agent_version*/ None,
Some(&turn.model_info),
);
})
.await;
@@ -906,6 +907,7 @@ async fn resolved_multi_agent_version_controls_runtime_behavior() {
turn.config.as_ref(),
&turn.session_source,
/*inherited_multi_agent_version*/ None,
Some(&turn.model_info),
);
})
.await;
@@ -931,6 +933,7 @@ async fn resolved_multi_agent_version_controls_runtime_behavior() {
turn.config.as_ref(),
&turn.session_source,
Some(MultiAgentVersion::V1),
Some(&turn.model_info),
);
})
.await;
@@ -1,40 +1,31 @@
use anyhow::Result;
use codex_core::config::Config;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_models_manager::manager::RefreshStrategy;
use codex_models_manager::manager::SharedModelsManager;
use codex_models_manager::model_info::model_info_from_slug;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelPreset;
use codex_protocol::openai_models::ModelVisibility;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::openai_models::MultiAgentVersion;
use codex_protocol::openai_models::ToolMode;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::mount_models_once;
use core_test_support::responses::mount_sse_once;
use core_test_support::responses::sse;
use core_test_support::skip_if_no_network;
use core_test_support::submit_thread_settings;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::Value;
use tokio::time::Duration;
use tokio::time::Instant;
use tokio::time::sleep;
fn remote_model(slug: &str) -> ModelInfo {
fn catalog_model(slug: &str) -> ModelInfo {
ModelInfo {
visibility: ModelVisibility::List,
used_fallback_model_metadata: false,
supports_search_tool: false,
..model_info_from_slug(slug)
}
}
@@ -56,38 +47,31 @@ fn tool_names(body: &Value) -> Vec<String> {
.unwrap_or_default()
}
async fn wait_for_model_available(manager: &SharedModelsManager, slug: &str) -> ModelPreset {
let deadline = Instant::now() + Duration::from_secs(2);
loop {
if let Some(model) = manager
.list_models(RefreshStrategy::Online)
.await
.iter()
.find(|model| model.model == slug)
.cloned()
{
return model;
}
if Instant::now() >= deadline {
panic!("timed out waiting for the remote model {slug} to appear");
}
sleep(Duration::from_millis(25)).await;
}
fn selected_tool_names(body: &Value, selected: &[&str]) -> Vec<String> {
tool_names(body)
.into_iter()
.filter(|name| selected.contains(&name.as_str()))
.collect()
}
async fn response_body_for_remote_model(
remote_model: ModelInfo,
fn tool_description<'a>(body: &'a Value, name: &str) -> Option<&'a str> {
body.get("tools")
.and_then(Value::as_array)
.and_then(|tools| {
tools
.iter()
.find(|tool| tool.get("name").and_then(Value::as_str) == Some(name))
})
.and_then(|tool| tool.get("description"))
.and_then(Value::as_str)
}
async fn response_body_for_catalog_model(
catalog_model: ModelInfo,
configure: impl FnOnce(&mut Config) + Send + 'static,
) -> Result<Value> {
let server = responses::start_mock_server().await;
let model_slug = remote_model.slug.clone();
let models_mock = mount_models_once(
&server,
ModelsResponse {
models: vec![remote_model],
},
)
.await;
let model_slug = catalog_model.slug.clone();
let response_mock = mount_sse_once(
&server,
sse(vec![
@@ -98,23 +82,14 @@ async fn response_body_for_remote_model(
)
.await;
let mut builder = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(configure);
let mut builder = test_codex().with_config(move |config| {
config.model = Some(model_slug);
config.model_catalog = Some(ModelsResponse {
models: vec![catalog_model],
});
configure(config);
});
let test = builder.build(&server).await?;
let models_manager = test.thread_manager.get_models_manager();
let available_model = wait_for_model_available(&models_manager, &model_slug).await;
assert_eq!(available_model.model, model_slug);
assert_eq!(models_mock.requests().len(), 1);
submit_thread_settings(
&test.codex,
ThreadSettingsOverrides {
model: Some(model_slug),
..Default::default()
},
)
.await?;
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
@@ -137,12 +112,10 @@ async fn response_body_for_remote_model(
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> {
skip_if_no_network!(Ok(()));
let mut direct_model = remote_model("test-tool-mode-direct");
async fn catalog_tool_mode_selector_overrides_feature_flags() -> Result<()> {
let mut direct_model = catalog_model("test-tool-mode-direct");
direct_model.tool_mode = Some(ToolMode::Direct);
let direct_body = response_body_for_remote_model(direct_model, |config| {
let direct_body = response_body_for_catalog_model(direct_model, |config| {
config
.features
.enable(Feature::CodeModeOnly)
@@ -158,9 +131,9 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> {
"direct mode should override enabled code mode flags: {direct_tools:?}"
);
let mut code_mode_only_model = remote_model("test-tool-mode-code-mode-only");
let mut code_mode_only_model = catalog_model("test-tool-mode-code-mode-only");
code_mode_only_model.tool_mode = Some(ToolMode::CodeModeOnly);
let code_mode_only_body = response_body_for_remote_model(code_mode_only_model, |_| {}).await?;
let code_mode_only_body = response_body_for_catalog_model(code_mode_only_model, |_| {}).await?;
assert_eq!(
tool_names(&code_mode_only_body),
vec![
@@ -171,3 +144,167 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn catalog_multi_agent_version_selector_overrides_feature_flags() -> Result<()> {
let mut v1_model = catalog_model("test-multi-agent-v1");
v1_model.multi_agent_version = Some(MultiAgentVersion::V1);
let v1_body = response_body_for_catalog_model(v1_model, |config| {
config
.features
.enable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
config.multi_agent_v2.root_agent_usage_hint_text =
Some("V2 guidance must not reach v1 models.".to_string());
})
.await?;
assert_eq!(
selected_tool_names(
&v1_body,
&[
"spawn_agent",
"send_input",
"resume_agent",
"wait_agent",
"close_agent",
"send_message",
"followup_task",
"list_agents",
],
),
vec![
"spawn_agent".to_string(),
"send_input".to_string(),
"resume_agent".to_string(),
"wait_agent".to_string(),
"close_agent".to_string(),
]
);
assert!(
!v1_body
.to_string()
.contains("V2 guidance must not reach v1 models."),
"v1 models should not receive v2 usage hints: {v1_body:?}"
);
let mut v2_model = catalog_model("test-multi-agent-v2");
v2_model.multi_agent_version = Some(MultiAgentVersion::V2);
let v2_body = response_body_for_catalog_model(v2_model, |config| {
config
.features
.disable(Feature::Collab)
.expect("test config should allow feature update");
config
.features
.disable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
config.multi_agent_v2.max_concurrent_threads_per_session = 17;
config.multi_agent_v2.root_agent_usage_hint_text =
Some("V2 guidance should reach v2 models.".to_string());
})
.await?;
assert_eq!(
selected_tool_names(
&v2_body,
&[
"spawn_agent",
"send_input",
"resume_agent",
"wait_agent",
"close_agent",
"send_message",
"followup_task",
"list_agents",
],
),
vec![
"spawn_agent".to_string(),
"send_message".to_string(),
"followup_task".to_string(),
"wait_agent".to_string(),
"close_agent".to_string(),
"list_agents".to_string(),
]
);
assert!(
tool_description(&v2_body, "spawn_agent").is_some_and(
|description| description.contains("max_concurrent_threads_per_session = 17")
),
"v2 spawn_agent should advertise the configured concurrency cap: {v2_body:?}"
);
assert!(
v2_body
.to_string()
.contains("V2 guidance should reach v2 models."),
"v2 models should receive v2 usage hints: {v2_body:?}"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn catalog_omitted_and_unknown_multi_agent_versions_follow_feature_flags() -> Result<()> {
let omitted_body =
response_body_for_catalog_model(catalog_model("test-multi-agent-omitted"), |config| {
config
.features
.enable(Feature::Collab)
.expect("test config should allow feature update");
config
.features
.disable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
})
.await?;
let mut unknown_model =
serde_json::to_value(catalog_model("test-multi-agent-unknown-version"))?;
unknown_model["multi_agent_version"] = Value::String("future_multi_agent_version".to_string());
let unknown_model = serde_json::from_value::<ModelInfo>(unknown_model)?;
let unknown_body = response_body_for_catalog_model(unknown_model, |config| {
config
.features
.enable(Feature::Collab)
.expect("test config should allow feature update");
config
.features
.disable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
})
.await?;
let expected_v1_tools = vec![
"spawn_agent".to_string(),
"send_input".to_string(),
"resume_agent".to_string(),
"wait_agent".to_string(),
"close_agent".to_string(),
];
assert_eq!(
(
selected_tool_names(
&omitted_body,
&[
"spawn_agent",
"send_input",
"resume_agent",
"wait_agent",
"close_agent",
],
),
selected_tool_names(
&unknown_body,
&[
"spawn_agent",
"send_input",
"resume_agent",
"wait_agent",
"close_agent",
],
),
),
(expected_v1_tools.clone(), expected_v1_tools)
);
Ok(())
}
@@ -113,6 +113,7 @@ fn test_model_info(
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
@@ -931,6 +932,7 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result<
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
@@ -371,5 +371,6 @@ fn test_remote_model(slug: &str, priority: i32) -> ModelInfo {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}
}
+2
View File
@@ -593,6 +593,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
};
let _models_mock = mount_models_once(
@@ -704,6 +705,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() -
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
};
let _models_mock = mount_models_once(
@@ -478,6 +478,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
@@ -728,6 +729,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
@@ -1212,6 +1214,7 @@ fn test_remote_model_with_policy(
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
+1
View File
@@ -1350,6 +1350,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}],
},
)
@@ -61,6 +61,7 @@ fn test_model_info(
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers,
+1
View File
@@ -1356,6 +1356,7 @@ async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> an
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
priority: 1,
additional_speed_tiers: Vec::new(),
service_tiers: Vec::new(),
@@ -100,6 +100,7 @@ pub fn model_info_from_slug(slug: &str) -> ModelInfo {
used_fallback_model_metadata: true, // this is the fallback model metadata
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}
}
+45
View File
@@ -352,6 +352,12 @@ pub struct ModelInfo {
deserialize_with = "deserialize_optional_model_selector"
)]
pub tool_mode: Option<ToolMode>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
deserialize_with = "deserialize_optional_model_selector"
)]
pub multi_agent_version: Option<MultiAgentVersion>,
}
impl ModelInfo {
@@ -647,6 +653,7 @@ mod tests {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}
}
@@ -865,6 +872,7 @@ mod tests {
assert_eq!(model.web_search_tool_type, WebSearchToolType::Text);
assert!(!model.supports_search_tool);
assert_eq!(model.tool_mode, None);
assert_eq!(model.multi_agent_version, None);
}
#[test]
@@ -904,6 +912,43 @@ mod tests {
assert!(!object.contains_key("tool_mode"));
}
#[test]
fn model_info_deserializes_known_multi_agent_version() {
let mut value =
serde_json::to_value(test_model(/*spec*/ None)).expect("serialize test model");
let object = value
.as_object_mut()
.expect("model info should be an object");
object.insert(
"multi_agent_version".to_string(),
serde_json::Value::String("v2".to_string()),
);
let model = serde_json::from_value::<ModelInfo>(value).expect("deserialize model info");
assert_eq!(model.multi_agent_version, Some(MultiAgentVersion::V2));
}
#[test]
fn model_info_treats_unknown_multi_agent_version_as_omitted() {
let mut value =
serde_json::to_value(test_model(/*spec*/ None)).expect("serialize test model");
let object = value
.as_object_mut()
.expect("model info should be an object");
object.insert(
"multi_agent_version".to_string(),
serde_json::Value::String("future_multi_agent_version".to_string()),
);
let model = serde_json::from_value::<ModelInfo>(value).expect("deserialize model info");
assert_eq!(model.multi_agent_version, None);
let serialized = serde_json::to_value(model).expect("serialize model info");
let object = serialized
.as_object()
.expect("model info should be an object");
assert!(!object.contains_key("multi_agent_version"));
}
#[test]
fn resolved_context_window_prefers_context_window() {
let model = ModelInfo {
+1
View File
@@ -45,6 +45,7 @@ fn model_with_shell_type(shell_type: ConfigShellToolType) -> ModelInfo {
used_fallback_model_metadata: false,
supports_search_tool: false,
tool_mode: None,
multi_agent_version: None,
}
}