Codex Handbook
core/tests/suite/model_runtime_selectors.rs 307 lines
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::InputModality;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::ToolMode;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::MultiAgentVersion;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;const CHILD_MODEL: &str = "test-multi-agent-child";const ROOT_MODEL: &str = "test-multi-agent-root";const ROOT_PROMPT: &str = "spawn a child";fn remote_model(slug: &str) -> ModelInfo {    ModelInfo {        visibility: ModelVisibility::List,        used_fallback_model_metadata: false,        ..model_info_from_slug(slug)    }}fn tool_names(body: &Value) -> Vec<String> {    body.get("tools")        .and_then(Value::as_array)        .map(|tools| {            tools                .iter()                .filter_map(|tool| {                    tool.get("name")                        .or_else(|| tool.get("type"))                        .and_then(Value::as_str)                        .map(str::to_string)                })                .collect()        })        .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;    }}async fn response_body_for_remote_model(    remote_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 response_mock = mount_sse_once(        &server,        sse(vec![            ev_response_created("resp-1"),            ev_assistant_message("msg-1", "done"),            ev_completed("resp-1"),        ]),    )    .await;    let mut builder = test_codex()        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())        .with_config(configure);    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 {                text: "list tools".into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: Default::default(),        })        .await?;    wait_for_event(&test.codex, |event| {        matches!(event, EventMsg::TurnComplete(_))    })    .await;    Ok(response_mock.single_request().body_json())}#[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");    direct_model.tool_mode = Some(ToolMode::Direct);    let direct_body = response_body_for_remote_model(direct_model, |config| {        config            .features            .enable(Feature::CodeModeOnly)            .expect("test config should allow feature update");    })    .await?;    let direct_tools = tool_names(&direct_body);    assert!(        direct_tools            .iter()            .all(|name| name != codex_code_mode::PUBLIC_TOOL_NAME                && name != codex_code_mode::WAIT_TOOL_NAME),        "direct mode should override enabled code mode flags: {direct_tools:?}"    );    let mut code_mode_only_model = remote_model("test-tool-mode-code-mode-only");    code_mode_only_model.tool_mode = Some(ToolMode::CodeModeOnly);    code_mode_only_model.input_modalities = vec![InputModality::Text, InputModality::Image];    let code_mode_only_body = response_body_for_remote_model(code_mode_only_model, |_| {}).await?;    assert_eq!(        tool_names(&code_mode_only_body),        vec![            // Code-mode entrypoints.            codex_code_mode::PUBLIC_TOOL_NAME.to_string(),            codex_code_mode::WAIT_TOOL_NAME.to_string(),            "request_user_input".to_string(),            // Hosted Responses tools.            "web_search".to_string(),            "image_generation".to_string(),        ]    );    Ok(())}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn remote_multi_agent_selector_overrides_feature_flags() -> Result<()> {    skip_if_no_network!(Ok(()));    let mut v2_model = remote_model("test-multi-agent-v2");    v2_model.multi_agent_version = Some(MultiAgentVersion::V2);    let v2_body = response_body_for_remote_model(v2_model, |config| {        config.agent_max_threads = Some(3);        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?;    assert!(tool_names(&v2_body).contains(&"send_message".to_string()));    let mut disabled_model = remote_model("test-multi-agent-disabled");    disabled_model.multi_agent_version = Some(MultiAgentVersion::Disabled);    let disabled_body = response_body_for_remote_model(disabled_model, |config| {        config            .features            .enable(Feature::MultiAgentV2)            .expect("test config should allow feature update");    })    .await?;    let disabled_tools = tool_names(&disabled_body);    assert!(disabled_tools.iter().all(|name| !matches!(        name.as_str(),        "multi_agent_v1" | "spawn_agent" | "send_message" | "wait_agent" | "list_agents"    )));    Ok(())}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn remote_multi_agent_selector_uses_model_selected_before_first_turn() -> Result<()> {    skip_if_no_network!(Ok(()));    let server = wiremock::MockServer::start().await;    let mut initial_model = remote_model(ROOT_MODEL);    initial_model.multi_agent_version = Some(MultiAgentVersion::V1);    let mut selected_model = remote_model(CHILD_MODEL);    selected_model.multi_agent_version = Some(MultiAgentVersion::V2);    let models_mock = mount_models_once(        &server,        ModelsResponse {            models: vec![initial_model, selected_model],        },    )    .await;    let response_mock = mount_sse_once(        &server,        sse(vec![            ev_response_created("resp-1"),            ev_assistant_message("msg-1", "done"),            ev_completed("resp-1"),        ]),    )    .await;    let mut builder = test_codex()        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())        .with_config(|config| {            config.model = Some(ROOT_MODEL.to_string());        });    let test = builder.build(&server).await?;    assert_eq!(        (            models_mock.requests().len(),            test.codex.multi_agent_version(),        ),        (1, None)    );    submit_thread_settings(        &test.codex,        ThreadSettingsOverrides {            model: Some(CHILD_MODEL.to_string()),            ..Default::default()        },    )    .await?;    assert_eq!(test.codex.multi_agent_version(), None);    test.codex        .submit(Op::UserInput {            items: vec![UserInput::Text {                text: ROOT_PROMPT.into(),                text_elements: Vec::new(),            }],            final_output_json_schema: None,            responsesapi_client_metadata: None,            additional_context: Default::default(),            thread_settings: Default::default(),        })        .await?;    wait_for_event(&test.codex, |event| {        matches!(event, EventMsg::TurnComplete(_))    })    .await;    assert_eq!(        (            models_mock.requests().len(),            test.codex.multi_agent_version(),            tool_names(                &response_mock                    .last_request()                    .expect("expected response request")                    .body_json(),            )            .contains(&"send_message".to_string()),        ),        (1, Some(MultiAgentVersion::V2), true)    );    Ok(())}