[codex] Inline multi-agent version resolution

This commit is contained in:
Ahmed Ibrahim
2026-05-29 00:43:23 -07:00
parent 129b77b71b
commit 581291bf68
6 changed files with 96 additions and 121 deletions

View File

@@ -60,8 +60,6 @@ pub use codex_mcp::SandboxState;
mod mcp_openai_file;
mod mcp_tool_call;
pub(crate) mod mention_syntax;
mod multi_agent_version;
mod tool_mode;
pub(crate) mod utils;
pub use mention_syntax::PLUGIN_TEXT_MENTION_SIGIL;
pub use mention_syntax::TOOL_MENTION_SIGIL;

View File

@@ -1,63 +0,0 @@
use crate::config::ManagedFeatures;
use codex_features::Feature;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::MultiAgentVersion;
pub(crate) fn resolve_multi_agent_version(
model_info: &ModelInfo,
features: &ManagedFeatures,
) -> Option<MultiAgentVersion> {
model_info
.multi_agent_version
.or_else(|| multi_agent_version_from_features(features))
}
fn multi_agent_version_from_features(features: &ManagedFeatures) -> Option<MultiAgentVersion> {
if features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use codex_features::Features;
fn features(enabled: &[Feature]) -> ManagedFeatures {
let mut features = Features::default();
for feature in enabled {
features.enable(*feature);
}
features.into()
}
fn model_info(multi_agent_version: Option<MultiAgentVersion>) -> ModelInfo {
let mut model_info = codex_models_manager::model_info::model_info_from_slug("test-model");
model_info.multi_agent_version = multi_agent_version;
model_info
}
#[test]
fn omitted_selector_follows_feature_flags() {
let multi_agent_version = resolve_multi_agent_version(
&model_info(/*multi_agent_version*/ None),
&features(&[Feature::MultiAgentV2]),
);
assert_eq!(multi_agent_version, Some(MultiAgentVersion::V2));
}
#[test]
fn explicit_selector_overrides_feature_flags() {
let multi_agent_version = resolve_multi_agent_version(
&model_info(Some(MultiAgentVersion::V1)),
&features(&[Feature::MultiAgentV2]),
);
assert_eq!(multi_agent_version, Some(MultiAgentVersion::V1));
}
}

View File

@@ -1,6 +1,6 @@
use super::*;
use crate::multi_agent_version::resolve_multi_agent_version;
use crate::tool_mode::resolve_tool_mode;
use codex_protocol::openai_models::MultiAgentVersion;
use codex_protocol::openai_models::ToolMode;
use std::sync::atomic::AtomicBool;
/// Spawn a review thread using the given prompt.
@@ -49,8 +49,24 @@ pub(super) async fn spawn_review_thread(
let mut per_turn_config = (*config).clone();
per_turn_config.model = Some(model.clone());
per_turn_config.features = review_features.clone();
let tool_mode = resolve_tool_mode(&model_info, &per_turn_config.features);
let multi_agent_version = resolve_multi_agent_version(&model_info, &per_turn_config.features);
let tool_mode = model_info.tool_mode.unwrap_or_else(|| {
if per_turn_config.features.enabled(Feature::CodeModeOnly) {
ToolMode::CodeModeOnly
} else if per_turn_config.features.enabled(Feature::CodeMode) {
ToolMode::CodeMode
} else {
ToolMode::Direct
}
});
let multi_agent_version = model_info.multi_agent_version.or_else(|| {
if per_turn_config.features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if per_turn_config.features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
});
if let Err(err) = per_turn_config.web_search_mode.set(review_web_search_mode) {
let fallback_value = per_turn_config.web_search_mode.value();
tracing::warn!(

View File

@@ -2,8 +2,6 @@ use super::*;
use crate::SkillLoadOutcome;
use crate::config::GhostSnapshotConfig;
use crate::environment_selection::ResolvedTurnEnvironments;
use crate::multi_agent_version::resolve_multi_agent_version;
use crate::tool_mode::resolve_tool_mode;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_protocol::SessionId;
@@ -168,14 +166,6 @@ impl TurnContext {
self.goal_tools_supported && self.features.get().enabled(Feature::Goals)
}
pub(crate) fn code_mode_enabled(&self) -> bool {
matches!(self.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly)
}
pub(crate) fn code_mode_only_enabled(&self) -> bool {
self.tool_mode == ToolMode::CodeModeOnly
}
pub(crate) async fn with_model(
&self,
model: String,
@@ -186,8 +176,24 @@ impl TurnContext {
let model_info = models_manager
.get_model_info(model.as_str(), &config.to_models_manager_config())
.await;
let tool_mode = resolve_tool_mode(&model_info, &config.features);
let multi_agent_version = resolve_multi_agent_version(&model_info, &config.features);
let tool_mode = model_info.tool_mode.unwrap_or_else(|| {
if config.features.enabled(Feature::CodeModeOnly) {
ToolMode::CodeModeOnly
} else if config.features.enabled(Feature::CodeMode) {
ToolMode::CodeMode
} else {
ToolMode::Direct
}
});
let multi_agent_version = 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
}
});
let truncation_policy = model_info.truncation_policy.into();
let supported_reasoning_levels = model_info
.supported_reasoning_levels
@@ -493,9 +499,24 @@ impl Session {
);
let mut per_turn_config = per_turn_config;
let tool_mode = resolve_tool_mode(&model_info, &per_turn_config.features);
let multi_agent_version =
resolve_multi_agent_version(&model_info, &per_turn_config.features);
let tool_mode = model_info.tool_mode.unwrap_or_else(|| {
if per_turn_config.features.enabled(Feature::CodeModeOnly) {
ToolMode::CodeModeOnly
} else if per_turn_config.features.enabled(Feature::CodeMode) {
ToolMode::CodeMode
} else {
ToolMode::Direct
}
});
let multi_agent_version = model_info.multi_agent_version.or_else(|| {
if per_turn_config.features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if per_turn_config.features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
});
per_turn_config.service_tier = get_service_tier(
per_turn_config.service_tier,
per_turn_config.features.enabled(Feature::FastMode),

View File

@@ -61,6 +61,7 @@ use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::openai_models::ConfigShellToolType;
use codex_protocol::openai_models::InputModality;
use codex_protocol::openai_models::MultiAgentVersion;
use codex_protocol::openai_models::ToolMode;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_tools::DiscoverableTool;
@@ -231,8 +232,10 @@ fn spec_for_model_request(
exposure: ToolExposure,
spec: ToolSpec,
) -> ToolSpec {
if code_mode_enabled(turn_context)
&& exposure != ToolExposure::DirectModelOnly
if matches!(
turn_context.tool_mode,
ToolMode::CodeMode | ToolMode::CodeModeOnly
) && exposure != ToolExposure::DirectModelOnly
&& codex_code_mode::is_code_mode_nested_tool(spec.name())
{
codex_tools::augment_tool_spec_for_code_mode(spec)
@@ -283,22 +286,6 @@ fn namespace_tools_enabled(turn_context: &TurnContext) -> bool {
turn_context.provider.capabilities().namespace_tools
}
fn code_mode_enabled(turn_context: &TurnContext) -> bool {
turn_context.code_mode_enabled()
}
fn code_mode_only_enabled(turn_context: &TurnContext) -> bool {
turn_context.code_mode_only_enabled()
}
fn multi_agent_v2_enabled(turn_context: &TurnContext) -> bool {
turn_context.multi_agent_version == Some(MultiAgentVersion::V2)
}
fn collab_tools_enabled(turn_context: &TurnContext) -> bool {
turn_context.multi_agent_version.is_some()
}
fn goal_tools_enabled(turn_context: &TurnContext) -> bool {
turn_context.goal_tools_enabled()
&& !matches!(
@@ -357,7 +344,7 @@ fn standalone_image_generation_available(
}
fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOptions {
if multi_agent_v2_enabled(turn_context) {
if turn_context.multi_agent_version == Some(MultiAgentVersion::V2) {
return WaitAgentTimeoutOptions {
default_timeout_ms: turn_context.config.multi_agent_v2.default_wait_timeout_ms,
min_timeout_ms: turn_context.config.multi_agent_v2.min_wait_timeout_ms,
@@ -373,7 +360,7 @@ fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOpt
}
fn max_concurrent_threads_per_session(turn_context: &TurnContext) -> Option<usize> {
multi_agent_v2_enabled(turn_context).then_some(
(turn_context.multi_agent_version == Some(MultiAgentVersion::V2)).then_some(
turn_context
.config
.multi_agent_v2
@@ -399,7 +386,7 @@ fn is_hidden_by_code_mode_only(
tool_name: &ToolName,
exposure: ToolExposure,
) -> bool {
code_mode_only_enabled(turn_context)
turn_context.tool_mode == ToolMode::CodeModeOnly
&& exposure != ToolExposure::DirectModelOnly
&& codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name(
tool_name,
@@ -411,7 +398,10 @@ fn build_code_mode_executors(
executors: &[Arc<dyn CoreToolRuntime>],
deferred_tools_available: bool,
) -> Vec<Arc<dyn CoreToolRuntime>> {
if !code_mode_enabled(turn_context) {
if !matches!(
turn_context.tool_mode,
ToolMode::CodeMode | ToolMode::CodeModeOnly
) {
return vec![];
}
@@ -445,7 +435,7 @@ fn build_code_mode_executors(
create_code_mode_tool(
&enabled_tools,
&namespace_descriptions,
code_mode_only_enabled(turn_context),
turn_context.tool_mode == ToolMode::CodeModeOnly,
deferred_tools_available,
),
code_mode_nested_tool_specs,
@@ -649,8 +639,8 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) {
let turn_context = context.turn_context;
if collab_tools_enabled(turn_context) {
if multi_agent_v2_enabled(turn_context) {
if turn_context.multi_agent_version.is_some() {
if turn_context.multi_agent_version == Some(MultiAgentVersion::V2) {
let exposure = if turn_context.config.multi_agent_v2.non_code_mode_only {
ToolExposure::DirectModelOnly
} else {
@@ -848,7 +838,10 @@ fn append_extension_tool_executors(
.iter()
.map(|executor| executor.tool_name())
.collect::<HashSet<_>>();
if code_mode_enabled(turn_context) {
if matches!(
turn_context.tool_mode,
ToolMode::CodeMode | ToolMode::CodeModeOnly
) {
reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME));
reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME));
}

View File

@@ -31,10 +31,8 @@ use codex_tools::ToolSpec;
use pretty_assertions::assert_eq;
use serde_json::json;
use crate::multi_agent_version::resolve_multi_agent_version;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tool_mode::resolve_tool_mode;
use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE;
use crate::tools::router::ToolRouter;
use crate::tools::router::ToolRouterParams;
@@ -219,7 +217,24 @@ fn set_feature(turn: &mut TurnContext, feature: Feature, enabled: bool) {
.expect("test feature should be disableable in config");
}
turn.config = Arc::new(config);
resolve_model_selectors_for_turn(turn);
turn.tool_mode = turn.model_info.tool_mode.unwrap_or_else(|| {
if turn.config.features.enabled(Feature::CodeModeOnly) {
ToolMode::CodeModeOnly
} else if turn.config.features.enabled(Feature::CodeMode) {
ToolMode::CodeMode
} else {
ToolMode::Direct
}
});
turn.multi_agent_version = turn.model_info.multi_agent_version.or_else(|| {
if turn.config.features.enabled(Feature::MultiAgentV2) {
Some(MultiAgentVersion::V2)
} else if turn.config.features.enabled(Feature::Collab) {
Some(MultiAgentVersion::V1)
} else {
None
}
});
}
fn set_features(turn: &mut TurnContext, features: &[Feature]) {
@@ -234,11 +249,6 @@ fn update_config(turn: &mut TurnContext, update: impl FnOnce(&mut crate::config:
turn.config = Arc::new(config);
}
fn resolve_model_selectors_for_turn(turn: &mut TurnContext) {
turn.tool_mode = resolve_tool_mode(&turn.model_info, &turn.config.features);
turn.multi_agent_version = resolve_multi_agent_version(&turn.model_info, &turn.config.features);
}
fn set_web_search_mode(turn: &mut TurnContext, mode: WebSearchMode) {
update_config(turn, |config| {
config
@@ -812,7 +822,7 @@ async fn tool_mode_selector_overrides_feature_flags() {
let direct = probe(|turn| {
set_features(turn, &[Feature::CodeMode, Feature::CodeModeOnly]);
turn.model_info.tool_mode = Some(ToolMode::Direct);
resolve_model_selectors_for_turn(turn);
turn.tool_mode = ToolMode::Direct;
})
.await;
direct.assert_visible_lacks(&[
@@ -826,7 +836,7 @@ async fn multi_agent_version_selector_overrides_feature_flags() {
let v1 = probe(|turn| {
set_feature(turn, Feature::MultiAgentV2, /*enabled*/ true);
turn.model_info.multi_agent_version = Some(MultiAgentVersion::V1);
resolve_model_selectors_for_turn(turn);
turn.multi_agent_version = Some(MultiAgentVersion::V1);
})
.await;
@@ -843,7 +853,7 @@ async fn multi_agent_version_selector_overrides_feature_flags() {
set_feature(turn, Feature::Collab, /*enabled*/ false);
set_feature(turn, Feature::MultiAgentV2, /*enabled*/ false);
turn.model_info.multi_agent_version = Some(MultiAgentVersion::V2);
resolve_model_selectors_for_turn(turn);
turn.multi_agent_version = Some(MultiAgentVersion::V2);
})
.await;