Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions codex-rs/core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ wildmatch = { workspace = true }

[features]
deterministic_process_ids = []
test-support = []


[target.'cfg(target_os = "linux")'.dependencies]
Expand Down
86 changes: 52 additions & 34 deletions codex-rs/core/src/codex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use crate::compact_remote::run_inline_remote_auto_compact_task;
use crate::exec_policy::load_exec_policy_for_features;
use crate::features::Feature;
use crate::features::Features;
use crate::openai_models::model_family::ModelFamily;
use crate::openai_models::models_manager::ModelsManager;
use crate::parse_command::parse_command;
use crate::parse_turn_item;
Expand Down Expand Up @@ -398,35 +399,39 @@ pub(crate) struct SessionSettingsUpdate {
}

impl Session {
fn build_per_turn_config(session_configuration: &SessionConfiguration) -> Config {
let config = session_configuration.original_config_do_not_use.clone();
let mut per_turn_config = (*config).clone();
per_turn_config.model = session_configuration.model.clone();
per_turn_config.model_reasoning_effort = session_configuration.model_reasoning_effort;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not super stoked we are keeping these. Is it possible to apply overrides to model but don't expose it on config an let everyone read random values?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non blocking, just ranting.

per_turn_config.model_reasoning_summary = session_configuration.model_reasoning_summary;
per_turn_config.features = config.features.clone();
per_turn_config
}

#[allow(clippy::too_many_arguments)]
fn make_turn_context(
auth_manager: Option<Arc<AuthManager>>,
models_manager: Arc<ModelsManager>,
otel_event_manager: &OtelEventManager,
provider: ModelProviderInfo,
session_configuration: &SessionConfiguration,
mut per_turn_config: Config,
model_family: ModelFamily,
conversation_id: ConversationId,
sub_id: String,
) -> TurnContext {
let config = session_configuration.original_config_do_not_use.clone();
let features = &config.features;
let mut per_turn_config = (*config).clone();
per_turn_config.model = session_configuration.model.clone();
per_turn_config.model_reasoning_effort = session_configuration.model_reasoning_effort;
per_turn_config.model_reasoning_summary = session_configuration.model_reasoning_summary;
per_turn_config.features = features.clone();
let model_family =
models_manager.construct_model_family(&per_turn_config.model, &per_turn_config);
if let Some(model_info) = get_model_info(&model_family) {
per_turn_config.model_context_window = Some(model_info.context_window);
}

let otel_event_manager = otel_event_manager.clone().with_model(
session_configuration.model.as_str(),
session_configuration.model.as_str(),
model_family.slug.as_str(),
);

let per_turn_config = Arc::new(per_turn_config);
let client = ModelClient::new(
Arc::new(per_turn_config.clone()),
per_turn_config.clone(),
auth_manager,
model_family.clone(),
otel_event_manager,
Expand All @@ -439,7 +444,7 @@ impl Session {

let tools_config = ToolsConfig::new(&ToolsConfigParams {
model_family: &model_family,
features,
features: &per_turn_config.features,
});

TurnContext {
Expand All @@ -452,14 +457,14 @@ impl Session {
user_instructions: session_configuration.user_instructions.clone(),
approval_policy: session_configuration.approval_policy,
sandbox_policy: session_configuration.sandbox_policy.clone(),
shell_environment_policy: config.shell_environment_policy.clone(),
shell_environment_policy: per_turn_config.shell_environment_policy.clone(),
tools_config,
final_output_json_schema: None,
codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(),
codex_linux_sandbox_exe: per_turn_config.codex_linux_sandbox_exe.clone(),
tool_call_gate: Arc::new(ReadinessFlag::new()),
exec_policy: session_configuration.exec_policy.clone(),
truncation_policy: TruncationPolicy::new(
&per_turn_config,
per_turn_config.as_ref(),
model_family.truncation_policy,
),
}
Expand Down Expand Up @@ -545,7 +550,9 @@ impl Session {
});
}

let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = models_manager
.construct_model_family(&config.model, &config)
.await;
// todo(aibrahim): why are we passing model here while it can change?
let otel_event_manager = OtelEventManager::new(
conversation_id,
Expand Down Expand Up @@ -768,12 +775,19 @@ impl Session {
session_configuration
};

let per_turn_config = Self::build_per_turn_config(&session_configuration);
let model_family = self
.services
.models_manager
.construct_model_family(&per_turn_config.model, &per_turn_config)
.await;
let mut turn_context: TurnContext = Self::make_turn_context(
Some(Arc::clone(&self.services.auth_manager)),
Arc::clone(&self.services.models_manager),
&self.services.otel_event_manager,
session_configuration.provider.clone(),
&session_configuration,
per_turn_config,
model_family,
self.conversation_id,
sub_id,
);
Expand Down Expand Up @@ -1907,7 +1921,8 @@ async fn spawn_review_thread(
let review_model_family = sess
.services
.models_manager
.construct_model_family(&model, &config);
.construct_model_family(&model, &config)
.await;
// For reviews, disable web_search and view_image regardless of global settings.
let mut review_features = sess.features.clone();
review_features
Expand Down Expand Up @@ -2738,15 +2753,12 @@ mod tests {
fn otel_event_manager(
conversation_id: ConversationId,
config: &Config,
models_manager: &ModelsManager,
model_family: &ModelFamily,
) -> OtelEventManager {
OtelEventManager::new(
conversation_id,
config.model.as_str(),
models_manager
.construct_model_family(&config.model, config)
.slug
.as_str(),
model_family.slug.as_str(),
None,
Some("test@test.com".to_string()),
Some(AuthMode::ChatGPT),
Expand All @@ -2769,9 +2781,6 @@ mod tests {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager.clone()));
let otel_event_manager =
otel_event_manager(conversation_id, config.as_ref(), &models_manager);

let session_configuration = SessionConfiguration {
provider: config.model_provider.clone(),
model: config.model.clone(),
Expand All @@ -2788,6 +2797,11 @@ mod tests {
exec_policy: Arc::new(RwLock::new(ExecPolicy::empty())),
session_source: SessionSource::Exec,
};
let per_turn_config = Session::build_per_turn_config(&session_configuration);
let model_family =
ModelsManager::construct_model_family_offline(&per_turn_config.model, &per_turn_config);
let otel_event_manager =
otel_event_manager(conversation_id, config.as_ref(), &model_family);

let state = SessionState::new(session_configuration.clone());

Expand All @@ -2801,16 +2815,17 @@ mod tests {
show_raw_agent_reasoning: config.show_raw_agent_reasoning,
auth_manager: auth_manager.clone(),
otel_event_manager: otel_event_manager.clone(),
models_manager: models_manager.clone(),
models_manager,
tool_approvals: Mutex::new(ApprovalStore::default()),
};

let turn_context = Session::make_turn_context(
Some(Arc::clone(&auth_manager)),
models_manager,
&otel_event_manager,
session_configuration.provider.clone(),
&session_configuration,
per_turn_config,
model_family,
conversation_id,
"turn_id".to_string(),
);
Expand Down Expand Up @@ -2848,9 +2863,6 @@ mod tests {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager.clone()));
let otel_event_manager =
otel_event_manager(conversation_id, config.as_ref(), &models_manager);

let session_configuration = SessionConfiguration {
provider: config.model_provider.clone(),
model: config.model.clone(),
Expand All @@ -2867,6 +2879,11 @@ mod tests {
exec_policy: Arc::new(RwLock::new(ExecPolicy::empty())),
session_source: SessionSource::Exec,
};
let per_turn_config = Session::build_per_turn_config(&session_configuration);
let model_family =
ModelsManager::construct_model_family_offline(&per_turn_config.model, &per_turn_config);
let otel_event_manager =
otel_event_manager(conversation_id, config.as_ref(), &model_family);

let state = SessionState::new(session_configuration.clone());

Expand All @@ -2880,16 +2897,17 @@ mod tests {
show_raw_agent_reasoning: config.show_raw_agent_reasoning,
auth_manager: Arc::clone(&auth_manager),
otel_event_manager: otel_event_manager.clone(),
models_manager: models_manager.clone(),
models_manager,
tool_approvals: Mutex::new(ApprovalStore::default()),
};

let turn_context = Arc::new(Session::make_turn_context(
Some(Arc::clone(&auth_manager)),
models_manager,
&otel_event_manager,
session_configuration.provider.clone(),
&session_configuration,
per_turn_config,
model_family,
conversation_id,
"turn_id".to_string(),
));
Expand Down
9 changes: 8 additions & 1 deletion codex-rs/core/src/openai_models/models_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,14 @@ impl ModelsManager {
Ok(models)
}

pub fn construct_model_family(&self, model: &str, config: &Config) -> ModelFamily {
pub async fn construct_model_family(&self, model: &str, config: &Config) -> ModelFamily {
find_family_for_model(model)
.with_config_overrides(config)
.with_remote_overrides(self.remote_models.read().await.clone())
}

#[cfg(any(test, feature = "test-support"))]
pub fn construct_model_family_offline(model: &str, config: &Config) -> ModelFamily {
find_family_for_model(model).with_config_overrides(config)
}

Expand Down
4 changes: 3 additions & 1 deletion codex-rs/core/src/sandboxing/assessment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,9 @@ pub(crate) async fn assess_command(
output_schema: Some(sandbox_assessment_schema()),
};

let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = models_manager
.construct_model_family(&config.model, &config)
.await;

let child_otel = parent_otel.with_model(config.model.as_str(), model_family.slug.as_str());

Expand Down
6 changes: 1 addition & 5 deletions codex-rs/core/tests/chat_completions_payload.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
use std::sync::Arc;

use codex_app_server_protocol::AuthMode;
use codex_core::AuthManager;
use codex_core::CodexAuth;
use codex_core::ContentItem;
use codex_core::LocalShellAction;
use codex_core::LocalShellExecAction;
Expand Down Expand Up @@ -73,9 +71,7 @@ async fn run_request(input: Vec<ResponseItem>) -> Value {
let config = Arc::new(config);

let conversation_id = ConversationId::new();
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager));
let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);
let otel_event_manager = OtelEventManager::new(
conversation_id,
config.model.as_str(),
Expand Down
5 changes: 2 additions & 3 deletions codex-rs/core/tests/chat_completions_sse.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use assert_matches::assert_matches;
use codex_core::AuthManager;
use codex_core::openai_models::models_manager::ModelsManager;
use std::sync::Arc;
use tracing_test::traced_test;

Expand All @@ -12,6 +11,7 @@ use codex_core::Prompt;
use codex_core::ResponseEvent;
use codex_core::ResponseItem;
use codex_core::WireApi;
use codex_core::openai_models::models_manager::ModelsManager;
use codex_otel::otel_event_manager::OtelEventManager;
use codex_protocol::ConversationId;
use codex_protocol::models::ReasoningItemContent;
Expand Down Expand Up @@ -74,8 +74,7 @@ async fn run_stream_with_bytes(sse_body: &[u8]) -> Vec<ResponseEvent> {
let conversation_id = ConversationId::new();
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let auth_mode = auth_manager.get_auth_mode();
let models_manager = Arc::new(ModelsManager::new(auth_manager));
let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);
let otel_event_manager = OtelEventManager::new(
conversation_id,
config.model.as_str(),
Expand Down
2 changes: 1 addition & 1 deletion codex-rs/core/tests/common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ path = "lib.rs"
anyhow = { workspace = true }
assert_cmd = { workspace = true }
base64 = { workspace = true }
codex-core = { workspace = true }
codex-core = { workspace = true, features = ["test-support"] }
codex-protocol = { workspace = true }
notify = { workspace = true }
regex-lite = { workspace = true }
Expand Down
16 changes: 6 additions & 10 deletions codex-rs/core/tests/responses_headers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,7 @@ async fn responses_stream_includes_subagent_header_on_review() {

let conversation_id = ConversationId::new();
let auth_mode = AuthMode::ChatGPT;
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager));
let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);
let otel_event_manager = OtelEventManager::new(
conversation_id,
config.model.as_str(),
Expand Down Expand Up @@ -157,9 +155,7 @@ async fn responses_stream_includes_subagent_header_on_other() {

let conversation_id = ConversationId::new();
let auth_mode = AuthMode::ChatGPT;
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager));
let model_family = models_manager.construct_model_family(&config.model, &config);
let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);

let otel_event_manager = OtelEventManager::new(
conversation_id,
Expand Down Expand Up @@ -250,16 +246,16 @@ async fn responses_respects_model_family_overrides_from_config() {
let config = Arc::new(config);

let conversation_id = ConversationId::new();
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager.clone()));
let model_family = models_manager.construct_model_family(&config.model, &config);
let auth_mode =
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key")).get_auth_mode();
let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);
let otel_event_manager = OtelEventManager::new(
conversation_id,
config.model.as_str(),
model_family.slug.as_str(),
None,
Some("test@test.com".to_string()),
auth_manager.get_auth_mode(),
auth_mode,
false,
"test".to_string(),
);
Expand Down
4 changes: 1 addition & 3 deletions codex-rs/core/tests/suite/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1015,11 +1015,9 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
let effort = config.model_reasoning_effort;
let summary = config.model_reasoning_summary;
let config = Arc::new(config);

let model_family = ModelsManager::construct_model_family_offline(&config.model, &config);
let conversation_id = ConversationId::new();
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key"));
let models_manager = Arc::new(ModelsManager::new(auth_manager.clone()));
let model_family = models_manager.construct_model_family(&config.model, &config);
let otel_event_manager = OtelEventManager::new(
conversation_id,
config.model.as_str(),
Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/tests/suite/prompt_caching.rs
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> {
let base_instructions = conversation_manager
.get_models_manager()
.construct_model_family(&config.model, &config)
.await
.base_instructions
.clone();

Expand Down
1 change: 1 addition & 0 deletions codex-rs/tui/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ arboard = { workspace = true }


[dev-dependencies]
codex-core = { workspace = true, features = ["test-support"] }
assert_matches = { workspace = true }
chrono = { workspace = true, features = ["serde"] }
insta = { workspace = true }
Expand Down
Loading
Loading