Refresh loaded model provider sessions

This commit is contained in:
celia-oai
2026-07-09 17:20:21 -07:00
parent 056a453511
commit d83d952311
33 changed files with 1348 additions and 253 deletions
@@ -184,12 +184,15 @@ impl AccountRequestProcessor {
.set_auth_mode(self.auth_manager.get_api_auth_mode());
}
fn current_account_updated_notification(&self) -> AccountUpdatedNotification {
let auth = self.auth_manager.auth_cached();
async fn current_account_updated_notification(&self) -> AccountUpdatedNotification {
let config = self.load_latest_config().await;
let provider =
create_model_provider(config.model_provider, Some(self.auth_manager.clone()));
let auth = provider.auth().await;
AccountUpdatedNotification {
auth_mode: auth
.as_ref()
.map(CodexAuth::api_auth_mode)
.map(CodexAuth::auth_mode)
.map(auth_mode_to_api),
plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type),
}
@@ -424,7 +427,11 @@ impl AccountRequestProcessor {
self.config.auth_keyring_backend_kind(),
)
.map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?;
self.auth_manager.reload().await;
self.auth_manager
.reload_bedrock_api_key_auth()
.map_err(|err| {
internal_error(format!("failed to reload Amazon Bedrock auth: {err}"))
})?;
refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?;
Ok(LoginAccountResponse::AmazonBedrock {})
}
@@ -782,7 +789,7 @@ impl AccountRequestProcessor {
self.outgoing
.send_server_notification(ServerNotification::AccountUpdated(
self.current_account_updated_notification(),
self.current_account_updated_notification().await,
))
.await;
}
@@ -835,11 +842,9 @@ impl AccountRequestProcessor {
}
async fn logout_common(&self) -> std::result::Result<Option<AuthMode>, JSONRPCErrorError> {
let managed_bedrock_auth = matches!(
self.auth_manager.auth_cached(),
Some(CodexAuth::BedrockApiKey(_))
);
let config = self.load_latest_config().await;
let managed_bedrock_auth = config.model_provider.is_amazon_bedrock()
&& self.auth_manager.bedrock_api_key_auth_cached().is_some();
if config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth {
return Err(invalid_request(
"cannot log out while Amazon Bedrock is using AWS-managed credentials; manage those credentials through AWS or switch model providers before logging out Codex authentication",
@@ -854,16 +859,17 @@ impl AccountRequestProcessor {
}
}
match self.auth_manager.logout_with_revoke().await {
Ok(_) => {}
Err(err) => {
return Err(internal_error(format!("logout failed: {err}")));
}
}
if managed_bedrock_auth {
self.auth_manager
.logout_bedrock_api_key_auth()
.map_err(|err| internal_error(format!("logout failed: {err}")))?;
clear_user_model_provider_if_bedrock(&self.config_manager).await?;
refresh_default_model_provider(&self.config_manager, &self.thread_manager).await?;
} else {
self.auth_manager
.logout_with_revoke()
.await
.map_err(|err| internal_error(format!("logout failed: {err}")))?;
}
Self::maybe_refresh_plugin_caches_for_current_config(
@@ -1,6 +1,7 @@
use super::*;
use crate::error_code::method_not_found;
use codex_app_server_protocol::SelectedCapabilityRoot;
use codex_core::ModelProviderSelection;
use codex_extension_api::ExtensionDataInit;
use codex_protocol::config_types::MultiAgentMode;
use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS;
@@ -23,6 +24,19 @@ struct ThreadListFilters {
relation_filter: Option<StoreThreadRelationFilter>,
}
fn thread_start_model_provider_selection(
model_provider: Option<&str>,
request_config: Option<&HashMap<String, serde_json::Value>>,
) -> ModelProviderSelection {
if model_provider.is_some()
|| request_config.is_some_and(|config| config.contains_key("model_provider"))
{
ModelProviderSelection::Explicit
} else {
ModelProviderSelection::RuntimeDefault
}
}
fn collect_resume_override_mismatches(
request: &ThreadResumeParams,
config_snapshot: &ThreadConfigSnapshot,
@@ -968,6 +982,8 @@ impl ThreadRequestProcessor {
"`permissions` cannot be combined with `sandbox`",
));
}
let model_provider_selection =
thread_start_model_provider_selection(model_provider.as_deref(), config.as_ref());
let environment_selections =
resolve_turn_environment_selections(self.thread_manager.as_ref(), environments)?;
let runtime_workspace_roots = runtime_workspace_roots.map(resolve_runtime_workspace_roots);
@@ -1012,6 +1028,7 @@ impl ThreadRequestProcessor {
supports_openai_form_elicitation,
config,
typesafe_overrides,
model_provider_selection,
dynamic_tools,
selected_capability_roots.unwrap_or_default(),
history_mode.map(Into::into),
@@ -1089,6 +1106,7 @@ impl ThreadRequestProcessor {
supports_openai_form_elicitation: bool,
config_overrides: Option<HashMap<String, serde_json::Value>>,
typesafe_overrides: ConfigOverrides,
model_provider_selection: ModelProviderSelection,
dynamic_tools: Option<Vec<DynamicToolSpec>>,
selected_capability_roots: Vec<SelectedCapabilityRoot>,
history_mode: Option<ThreadHistoryMode>,
@@ -1217,25 +1235,32 @@ impl ThreadRequestProcessor {
..
} = listener_task_context
.thread_manager
.start_thread_with_options(StartThreadOptions {
config,
allow_provider_model_fallback,
initial_history: match session_start_source
.unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup)
{
codex_app_server_protocol::ThreadStartSource::Startup => InitialHistory::New,
codex_app_server_protocol::ThreadStartSource::Clear => InitialHistory::Cleared,
.start_thread_with_options_and_provider_selection(
StartThreadOptions {
config,
allow_provider_model_fallback,
initial_history: match session_start_source
.unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup)
{
codex_app_server_protocol::ThreadStartSource::Startup => {
InitialHistory::New
}
codex_app_server_protocol::ThreadStartSource::Clear => {
InitialHistory::Cleared
}
},
history_mode,
session_source: None,
thread_source,
dynamic_tools,
metrics_service_name: service_name,
parent_trace: request_trace,
environments,
thread_extension_init,
supports_openai_form_elicitation,
},
history_mode,
session_source: None,
thread_source,
dynamic_tools,
metrics_service_name: service_name,
parent_trace: request_trace,
environments,
thread_extension_init,
supports_openai_form_elicitation,
})
model_provider_selection,
)
.instrument(tracing::info_span!(
"app_server.thread_start.create_thread",
otel.name = "app_server.thread_start.create_thread",
@@ -1,3 +1,47 @@
mod model_provider_selection_tests {
use super::super::thread_start_model_provider_selection;
use codex_core::ModelProviderSelection;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::HashMap;
#[test]
fn omitted_provider_follows_the_runtime_default() {
assert_eq!(
thread_start_model_provider_selection(
/*model_provider*/ None, /*request_config*/ None
),
ModelProviderSelection::RuntimeDefault
);
let request_config = HashMap::from([("model".to_string(), json!("mock-model"))]);
assert_eq!(
thread_start_model_provider_selection(
/*model_provider*/ None,
Some(&request_config),
),
ModelProviderSelection::RuntimeDefault
);
}
#[test]
fn typed_or_config_provider_override_is_explicit() {
assert_eq!(
thread_start_model_provider_selection(Some("provider-a"), /*request_config*/ None,),
ModelProviderSelection::Explicit
);
let request_config = HashMap::from([("model_provider".to_string(), json!("provider-a"))]);
assert_eq!(
thread_start_model_provider_selection(
/*model_provider*/ None,
Some(&request_config),
),
ModelProviderSelection::Explicit
);
}
}
mod thread_list_cwd_filter_tests {
use super::super::normalize_thread_list_cwd_filters;
use codex_app_server_protocol::ThreadListCwdFilter;
@@ -1386,6 +1386,65 @@ async fn logout_managed_bedrock_restores_default_account() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn managed_bedrock_login_and_logout_preserve_cached_openai_auth() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
login_with_api_key(
codex_home.path(),
"sk-test-key",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
let mut mcp = TestAppServer::builder()
.with_codex_home(codex_home.path())
.without_auto_env()
.with_env_overrides(&[("OPENAI_API_KEY", None)])
.build()
.await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let request_id = mcp
.send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2")
.await?;
let response = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
assert_eq!(
to_response::<LoginAccountResponse>(response)?,
LoginAccountResponse::AmazonBedrock {}
);
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("account/login/completed"),
)
.await??;
assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?;
let request_id = mcp.send_logout_account_request().await?;
let response = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
assert_eq!(
to_response::<LogoutAccountResponse>(response)?,
LogoutAccountResponse {}
);
assert_account_updated(&mut mcp, Some(AuthMode::ApiKey)).await?;
assert_eq!(
read_account(&mut mcp).await?,
GetAccountResponse {
account: Some(Account::ApiKey {}),
requires_openai_auth: true,
}
);
Ok(())
}
#[tokio::test]
async fn logout_aws_managed_bedrock_errors_without_changing_auth_or_config() -> Result<()> {
let codex_home = TempDir::new()?;
@@ -17,8 +17,13 @@ use codex_app_server_protocol::ModelProviderCapabilitiesReadResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadListParams;
use codex_app_server_protocol::ThreadListResponse;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::UserInput;
use codex_app_server_protocol::WriteStatus;
use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID;
use core_test_support::responses;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
@@ -107,6 +112,53 @@ fn bedrock_capabilities() -> ModelProviderCapabilitiesReadResponse {
}
}
async fn start_provider_runtime_thread(
mcp: &mut TestAppServer,
model_provider: Option<&str>,
) -> Result<ThreadStartResponse> {
let request_id = mcp
.send_thread_start_request_with_auto_env(ThreadStartParams {
model: Some("mock-model".to_string()),
model_provider: model_provider.map(str::to_string),
..Default::default()
})
.await?;
let response: JSONRPCResponse = timeout(
DEFAULT_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
to_response(response)
}
async fn run_provider_runtime_turn(
mcp: &mut TestAppServer,
thread_id: &str,
text: &str,
) -> Result<()> {
let request_id = mcp
.send_turn_start_request(TurnStartParams {
thread_id: thread_id.to_string(),
input: vec![UserInput::Text {
text: text.to_string(),
text_elements: Vec::new(),
}],
..Default::default()
})
.await?;
timeout(
DEFAULT_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
timeout(
DEFAULT_TIMEOUT,
mcp.read_stream_until_notification_message("turn/completed"),
)
.await??;
Ok(())
}
#[tokio::test]
async fn config_write_publishes_provider_runtime_before_responding() -> Result<()> {
let codex_home = TempDir::new()?;
@@ -227,3 +279,102 @@ async fn managed_bedrock_login_and_logout_publish_provider_runtime() -> Result<(
);
Ok(())
}
#[tokio::test]
async fn loaded_threads_adopt_provider_switch_only_when_following_runtime_default() -> Result<()> {
let provider_a_server = responses::start_mock_server().await;
let provider_b_server = responses::start_mock_server().await;
let provider_a_mock = responses::mount_sse_once(
&provider_a_server,
responses::sse(vec![
responses::ev_response_created("provider-a-response"),
responses::ev_assistant_message("provider-a-message", "provider A"),
responses::ev_completed("provider-a-response"),
]),
)
.await;
let provider_b_mock = responses::mount_sse_once(
&provider_b_server,
responses::sse(vec![
responses::ev_response_created("provider-b-response"),
responses::ev_assistant_message("provider-b-message", "provider B"),
responses::ev_completed("provider-b-response"),
]),
)
.await;
let codex_home = TempDir::new()?;
std::fs::write(
codex_home.path().join("config.toml"),
format!(
r#"
model = "mock-model"
model_provider = "provider_a"
[model_providers.provider_a]
name = "Provider A"
base_url = "{}/v1"
wire_api = "responses"
request_max_retries = 0
stream_max_retries = 0
[model_providers.provider_b]
name = "Provider B"
base_url = "{}/v1"
wire_api = "responses"
request_max_retries = 0
stream_max_retries = 0
"#,
provider_a_server.uri(),
provider_b_server.uri(),
),
)?;
let mut mcp = TestAppServer::builder()
.with_codex_home(codex_home.path())
.build()
.await?;
timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??;
let runtime_default_thread = start_provider_runtime_thread(&mut mcp, None).await?;
let explicit_thread = start_provider_runtime_thread(&mut mcp, Some("provider_a")).await?;
assert_eq!(runtime_default_thread.model_provider, "provider_a");
assert_eq!(explicit_thread.model_provider, "provider_a");
let request_id = mcp
.send_config_value_write_request(ConfigValueWriteParams {
key_path: "model_provider".to_string(),
value: json!("provider_b"),
merge_strategy: MergeStrategy::Replace,
file_path: None,
expected_version: None,
})
.await?;
let response: JSONRPCResponse = timeout(
DEFAULT_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
assert_eq!(
to_response::<ConfigWriteResponse>(response)?.status,
WriteStatus::Ok
);
// Both threads were idle when the runtime changed. Selection is applied at the next turn
// boundary, so no in-flight request is expected to migrate between providers.
run_provider_runtime_turn(
&mut mcp,
&runtime_default_thread.thread.id,
"use the refreshed runtime provider",
)
.await?;
run_provider_runtime_turn(
&mut mcp,
&explicit_thread.thread.id,
"stay on the explicitly selected provider",
)
.await?;
assert_eq!(provider_b_mock.requests().len(), 1);
assert_eq!(provider_a_mock.requests().len(), 1);
Ok(())
}
+52
View File
@@ -426,6 +426,39 @@ impl ModelClient {
http_client_factory: HttpClientFactory,
) -> Self {
let model_provider = create_model_provider(provider_info, auth_manager);
Self::new_with_provider(
model_provider,
agent_identity_policy,
thread_id,
session_source,
originator,
model_verbosity,
enable_request_compression,
include_timing_metrics,
beta_features_header,
item_ids_enabled,
concurrent_reasoning_summaries_enabled,
attestation_provider,
http_client_factory,
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn new_with_provider(
model_provider: SharedModelProvider,
agent_identity_policy: AgentIdentityAuthPolicy,
thread_id: ThreadId,
session_source: SessionSource,
originator: String,
model_verbosity: Option<VerbosityConfig>,
enable_request_compression: bool,
include_timing_metrics: bool,
beta_features_header: Option<String>,
item_ids_enabled: bool,
concurrent_reasoning_summaries_enabled: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
http_client_factory: HttpClientFactory,
) -> Self {
let codex_api_key_env_enabled = model_provider
.auth_manager()
.as_ref()
@@ -458,6 +491,25 @@ impl ModelClient {
}
}
pub(crate) fn rebind_provider(&self, model_provider: SharedModelProvider) -> Self {
Self::new_with_provider(
model_provider,
self.agent_identity_policy,
self.state.thread_id,
self.state.session_source.clone(),
self.state.originator.clone(),
self.state.model_verbosity,
self.state.enable_request_compression,
self.state.include_timing_metrics,
self.state.beta_features_header.clone(),
self.state.item_ids_enabled,
self.state.concurrent_reasoning_summaries_enabled,
self.state.attestation_provider.clone(),
self.http_client_factory.clone(),
)
.with_prompt_cache_key_override(self.prompt_cache_key_override.clone())
}
pub(crate) fn with_prompt_cache_key_override(
mut self,
prompt_cache_key_override: Option<String>,
+7
View File
@@ -45,6 +45,7 @@ use crate::mcp_tool_call::build_guardian_mcp_tool_review_request;
use crate::mcp_tool_call::is_mcp_tool_approval_question_id;
use crate::mcp_tool_call::lookup_mcp_tool_metadata;
use crate::mcp_tool_call::mcp_approvals_reviewer;
use crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager;
use crate::session::Codex;
use crate::session::CodexSpawnArgs;
use crate::session::CodexSpawnOk;
@@ -91,6 +92,11 @@ pub(crate) async fn run_codex_thread_interactive(
instructions: parent_session.user_instructions().await,
warnings: Vec::new(),
};
let model_provider_runtime = build_explicit_model_provider_runtime_with_models_manager(
&config,
Arc::clone(&auth_manager),
Arc::clone(&models_manager),
);
let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs {
config,
allow_provider_model_fallback: false,
@@ -98,6 +104,7 @@ pub(crate) async fn run_codex_thread_interactive(
installation_id: parent_session.installation_id.clone(),
auth_manager,
models_manager,
model_provider_runtime,
environment_manager: parent_session
.services
.turn_environments
+1 -1
View File
@@ -237,7 +237,7 @@ async fn run_compact_task_inner_impl(
let max_retries = turn_context.provider.info().stream_max_retries();
let mut retries = 0;
let mut client_session = sess.services.model_client.new_session();
let mut client_session = turn_context.model_client.new_session();
// Reuse one client session so turn-scoped state (sticky routing, websocket incremental
// request tracking)
// survives retries within this compact turn.
@@ -97,7 +97,7 @@ pub(super) async fn run_remote_compact_v2_attempt(
let mut owned_client_session = None;
let client_session = match client_session {
Some(client_session) => client_session,
None => owned_client_session.insert(sess.services.model_client.new_session()),
None => owned_client_session.insert(turn_context.model_client.new_session()),
};
let compaction_output_result = run_remote_compaction_request_v2(
sess,
+1 -1
View File
@@ -665,7 +665,7 @@ async fn spawn_guardian_review_session(
let codex = Box::pin(run_codex_thread_interactive(
spawn_config,
parent_session.services.auth_manager.clone(),
parent_session.services.models_manager.clone(),
Arc::clone(&parent_turn.models_manager),
Arc::clone(parent_session),
Arc::clone(parent_turn),
cancel_token.clone(),
+1
View File
@@ -116,6 +116,7 @@ mod thread_manager;
pub(crate) mod web_search;
pub(crate) mod windows_sandbox_read_grants;
pub use thread_manager::ForkSnapshot;
pub use thread_manager::ModelProviderSelection;
pub use thread_manager::NewThread;
pub use thread_manager::StartThreadOptions;
pub use thread_manager::ThreadManager;
+102 -24
View File
@@ -2,6 +2,7 @@ use crate::config::Config;
use arc_swap::ArcSwap;
use codex_login::AuthManager;
use codex_model_provider::ProviderCapabilities;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo;
use codex_models_manager::manager::SharedModelsManager;
@@ -17,35 +18,44 @@ struct ModelProviderRuntimeIdentity {
provider_info: ModelProviderInfo,
codex_home: PathBuf,
config_model_catalog: Option<ModelsResponse>,
auth_revision: u64,
}
impl ModelProviderRuntimeIdentity {
fn from_config(config: &Config) -> Self {
fn from_config(config: &Config, auth_manager: &AuthManager) -> Self {
Self {
provider_id: config.model_provider_id.clone(),
provider_info: config.model_provider.clone(),
codex_home: config.codex_home.to_path_buf(),
config_model_catalog: config.model_catalog.clone(),
auth_revision: auth_revision(auth_manager),
}
}
fn matches_config(&self, config: &Config) -> bool {
fn matches_config(&self, config: &Config, auth_manager: &AuthManager) -> bool {
self.provider_id == config.model_provider_id
&& self.provider_info == config.model_provider
&& self.codex_home == config.codex_home.to_path_buf()
&& self.config_model_catalog == config.model_catalog
&& self.auth_revision == auth_revision(auth_manager)
}
}
fn auth_revision(auth_manager: &AuthManager) -> u64 {
let receiver = auth_manager.auth_change_receiver();
*receiver.borrow()
}
/// Immutable process-default model-provider state published as one coherent value.
struct DefaultModelProviderSnapshot {
pub(crate) struct ModelProviderRuntimeSnapshot {
generation: u64,
identity: ModelProviderRuntimeIdentity,
provider: SharedModelProvider,
models_manager: SharedModelsManager,
capabilities: ProviderCapabilities,
}
impl DefaultModelProviderSnapshot {
impl ModelProviderRuntimeSnapshot {
fn new(
generation: u64,
identity: ModelProviderRuntimeIdentity,
@@ -55,34 +65,78 @@ impl DefaultModelProviderSnapshot {
identity.provider_info.clone(),
Some(Arc::clone(&auth_manager)),
);
let capabilities = provider.capabilities();
let models_manager = provider.models_manager(
identity.codex_home.clone(),
identity.config_model_catalog.clone(),
);
Self::from_provider(generation, identity, provider, models_manager)
}
fn from_provider(
generation: u64,
identity: ModelProviderRuntimeIdentity,
provider: SharedModelProvider,
models_manager: SharedModelsManager,
) -> Self {
let capabilities = provider.capabilities();
Self {
generation,
identity,
provider,
models_manager,
capabilities,
}
}
fn matches_config(&self, config: &Config) -> bool {
self.identity.matches_config(config)
fn matches_config(&self, config: &Config, auth_manager: &AuthManager) -> bool {
self.identity.matches_config(config, auth_manager)
}
pub(crate) fn generation(&self) -> u64 {
self.generation
}
pub(crate) fn provider_info(&self) -> &ModelProviderInfo {
&self.identity.provider_info
}
pub(crate) fn provider(&self) -> SharedModelProvider {
Arc::clone(&self.provider)
}
pub(crate) fn models_manager(&self) -> SharedModelsManager {
Arc::clone(&self.models_manager)
}
pub(crate) fn apply_to_config(&self, config: &mut Config) {
config.model_provider_id = self.identity.provider_id.clone();
config.model_provider = self.identity.provider_info.clone();
config.model_catalog = self.identity.config_model_catalog.clone();
}
}
#[derive(Clone)]
pub(crate) enum ModelProviderRuntimeSource {
RuntimeDefault(Arc<DefaultModelProviderRuntime>),
Explicit,
}
#[derive(Clone)]
pub(crate) struct InitialModelProviderRuntime {
pub(crate) source: ModelProviderRuntimeSource,
pub(crate) snapshot: Arc<ModelProviderRuntimeSnapshot>,
}
/// Owns the atomically replaceable process-default model-provider snapshot.
pub(crate) struct DefaultModelProviderRuntime {
snapshot: ArcSwap<DefaultModelProviderSnapshot>,
snapshot: ArcSwap<ModelProviderRuntimeSnapshot>,
refresh_lock: Mutex<()>,
}
impl DefaultModelProviderRuntime {
pub(crate) fn new(config: &Config, auth_manager: Arc<AuthManager>) -> Self {
Self::from_identity(
ModelProviderRuntimeIdentity::from_config(config),
ModelProviderRuntimeIdentity::from_config(config, &auth_manager),
auth_manager,
)
}
@@ -100,6 +154,7 @@ impl DefaultModelProviderRuntime {
provider_info,
codex_home,
config_model_catalog,
auth_revision: auth_revision(&auth_manager),
},
auth_manager,
)
@@ -110,7 +165,7 @@ impl DefaultModelProviderRuntime {
auth_manager: Arc<AuthManager>,
) -> Self {
Self {
snapshot: ArcSwap::from_pointee(DefaultModelProviderSnapshot::new(
snapshot: ArcSwap::from_pointee(ModelProviderRuntimeSnapshot::new(
/*generation*/ 0,
identity,
auth_manager,
@@ -125,13 +180,13 @@ impl DefaultModelProviderRuntime {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let current = self.snapshot.load_full();
if current.matches_config(config) {
if current.matches_config(config, &auth_manager) {
return false;
}
let next = DefaultModelProviderSnapshot::new(
let next = ModelProviderRuntimeSnapshot::new(
current.generation.saturating_add(1),
ModelProviderRuntimeIdentity::from_config(config),
ModelProviderRuntimeIdentity::from_config(config, &auth_manager),
auth_manager,
);
self.snapshot.store(Arc::new(next));
@@ -142,17 +197,8 @@ impl DefaultModelProviderRuntime {
Arc::clone(&self.snapshot.load().models_manager)
}
pub(crate) fn models_manager_for_config(
&self,
config: &Config,
auth_manager: Arc<AuthManager>,
) -> SharedModelsManager {
let current = self.snapshot.load();
if current.matches_config(config) {
return Arc::clone(&current.models_manager);
}
build_models_manager(config, auth_manager)
pub(crate) fn snapshot(&self) -> Arc<ModelProviderRuntimeSnapshot> {
self.snapshot.load_full()
}
pub(crate) fn provider_id(&self) -> String {
@@ -168,6 +214,38 @@ impl DefaultModelProviderRuntime {
}
}
pub(crate) fn build_explicit_model_provider_runtime(
config: &Config,
auth_manager: Arc<AuthManager>,
) -> InitialModelProviderRuntime {
InitialModelProviderRuntime {
source: ModelProviderRuntimeSource::Explicit,
snapshot: Arc::new(ModelProviderRuntimeSnapshot::new(
/*generation*/ 0,
ModelProviderRuntimeIdentity::from_config(config, &auth_manager),
auth_manager,
)),
}
}
pub(crate) fn build_explicit_model_provider_runtime_with_models_manager(
config: &Config,
auth_manager: Arc<AuthManager>,
models_manager: SharedModelsManager,
) -> InitialModelProviderRuntime {
let identity = ModelProviderRuntimeIdentity::from_config(config, &auth_manager);
let provider = create_model_provider(identity.provider_info.clone(), Some(auth_manager));
InitialModelProviderRuntime {
source: ModelProviderRuntimeSource::Explicit,
snapshot: Arc::new(ModelProviderRuntimeSnapshot::from_provider(
/*generation*/ 0,
identity,
provider,
models_manager,
)),
}
}
pub(crate) fn build_models_manager(
config: &Config,
auth_manager: Arc<AuthManager>,
+1 -1
View File
@@ -1012,7 +1012,7 @@ async fn handle_start_inner(
codex_response_handoff_prefix,
realtime_call_api_provider,
session_config,
model_client: sess.services.model_client.clone(),
model_client: sess.services.model_runtime.load().model_client.clone(),
sdp,
};
let start_output = sess.conversation.start(start).await?;
+1 -1
View File
@@ -60,7 +60,7 @@ pub(crate) async fn handle_retryable_response_stream_error(
// transient reconnect messages. In debug builds, keep full visibility for diagnosis.
let report_error = retry_count > 1
|| cfg!(debug_assertions)
|| !sess.services.model_client.responses_websocket_enabled();
|| !turn_context.model_client.responses_websocket_enabled();
if report_error {
// Surface retry information to any UI/front-end so the user understands what is
// happening instead of staring at a seemingly frozen screen.
+6
View File
@@ -302,6 +302,8 @@ use crate::SkillsService;
use crate::exec_policy::ExecPolicyUpdateError;
use crate::guardian::GuardianReviewSessionManager;
use crate::mcp::McpManager;
use crate::model_provider_runtime::InitialModelProviderRuntime;
use crate::model_provider_runtime::ModelProviderRuntimeSource;
use crate::network_policy_decision::execpolicy_network_rule_amendment;
use crate::rollout::map_session_init_error;
use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
@@ -311,6 +313,7 @@ use crate::skills::SkillLoadOutcome;
use crate::state::AutoCompactWindowIds;
use crate::state::AutoCompactWindowSnapshot;
use crate::state::PendingRequestPermissions;
use crate::state::SessionModelRuntime;
use crate::state::SessionServices;
use crate::state::SessionState;
#[cfg(test)]
@@ -412,6 +415,7 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) installation_id: String,
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) models_manager: SharedModelsManager,
pub(crate) model_provider_runtime: InitialModelProviderRuntime,
pub(crate) environment_manager: Arc<EnvironmentManager>,
pub(crate) skills_service: Arc<SkillsService>,
pub(crate) plugins_manager: Arc<PluginsManager>,
@@ -504,6 +508,7 @@ impl Codex {
installation_id,
auth_manager,
models_manager,
model_provider_runtime,
environment_manager,
skills_service,
plugins_manager,
@@ -686,6 +691,7 @@ impl Codex {
installation_id,
auth_manager.clone(),
models_manager.clone(),
model_provider_runtime,
exec_policy,
tx_event.clone(),
agent_status_tx.clone(),
+5 -4
View File
@@ -13,8 +13,7 @@ pub(super) async fn spawn_review_thread(
.review_model
.clone()
.unwrap_or_else(|| parent_turn_context.model_info.slug.clone());
let review_model_info = sess
.services
let review_model_info = parent_turn_context
.models_manager
.get_model_info(&model, &config.to_models_manager_config())
.await;
@@ -24,8 +23,7 @@ pub(super) async fn spawn_review_thread(
let _ = review_features.disable(Feature::WebSearchCached);
let _ = review_features.disable(Feature::Goals);
let review_web_search_mode = WebSearchMode::Disabled;
let available_models = sess
.services
let available_models = parent_turn_context
.models_manager
.list_models(
RefreshStrategy::OnlineIfUncached,
@@ -117,6 +115,9 @@ pub(super) async fn spawn_review_thread(
model_info: model_info.clone(),
session_telemetry: session_telemetry_for_context,
provider: provider_for_context,
models_manager: Arc::clone(&parent_turn_context.models_manager),
model_client: parent_turn_context.model_client.clone(),
model_provider_generation: parent_turn_context.model_provider_generation,
reasoning_effort,
reasoning_summary,
session_source,
+40 -32
View File
@@ -484,6 +484,7 @@ impl Session {
installation_id: String,
auth_manager: Arc<AuthManager>,
models_manager: SharedModelsManager,
model_provider_runtime: InitialModelProviderRuntime,
exec_policy: Arc<ExecPolicyManager>,
tx_event: Sender<Event>,
agent_status: watch::Sender<AgentStatus>,
@@ -1048,6 +1049,43 @@ impl Session {
}).await;
}
let model_client = ModelClient::new_with_provider(
model_provider_runtime.snapshot.provider(),
if config.features.enabled(Feature::UseAgentIdentity) {
AgentIdentityAuthPolicy::ChatGptAuth
} else {
AgentIdentityAuthPolicy::JwtOnly
},
thread_id,
session_configuration.session_source.clone(),
session_configuration.originator.clone(),
config.model_verbosity,
config.features.enabled(Feature::EnableRequestCompression),
config.features.enabled(Feature::RuntimeMetrics),
Self::build_model_client_beta_features_header(config.as_ref()),
/*item_ids_enabled*/ config.features.enabled(Feature::ItemIds)
|| matches!(
session_configuration.history_mode,
ThreadHistoryMode::Paginated
),
/*concurrent_reasoning_summaries_enabled*/ config
.features
.enabled(Feature::ConcurrentReasoningSummaries),
attestation_provider.clone(),
config.http_client_factory(),
)
.with_prompt_cache_key_override(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
session_configuration.parent_thread_id,
),
);
let model_runtime = SessionModelRuntime {
source: model_provider_runtime.source,
snapshot: model_provider_runtime.snapshot,
model_client: model_client.clone(),
};
let services = SessionServices {
// Initialize the MCP connection manager with an uninitialized
// instance. It will be replaced with one created via
@@ -1073,6 +1111,7 @@ impl Session {
show_raw_agent_reasoning: config.show_raw_agent_reasoning,
exec_policy,
auth_manager: Arc::clone(&auth_manager),
model_runtime: arc_swap::ArcSwap::from_pointee(model_runtime),
session_telemetry,
models_manager: Arc::clone(&models_manager),
tool_approvals: Mutex::new(ApprovalStore::default()),
@@ -1102,38 +1141,7 @@ impl Session {
thread_store: Arc::clone(&thread_store),
attestation_provider: attestation_provider.clone(),
time_provider,
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
if config.features.enabled(Feature::UseAgentIdentity) {
AgentIdentityAuthPolicy::ChatGptAuth
} else {
AgentIdentityAuthPolicy::JwtOnly
},
thread_id,
session_configuration.provider.clone(),
session_configuration.session_source.clone(),
session_configuration.originator.clone(),
config.model_verbosity,
config.features.enabled(Feature::EnableRequestCompression),
config.features.enabled(Feature::RuntimeMetrics),
Self::build_model_client_beta_features_header(config.as_ref()),
/*item_ids_enabled*/ config.features.enabled(Feature::ItemIds)
|| matches!(
session_configuration.history_mode,
ThreadHistoryMode::Paginated
),
/*concurrent_reasoning_summaries_enabled*/ config
.features
.enabled(Feature::ConcurrentReasoningSummaries),
attestation_provider,
config.http_client_factory(),
)
.with_prompt_cache_key_override(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
session_configuration.parent_thread_id,
),
),
model_client,
code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::clone(
&code_mode_session_provider,
)),
+77 -4
View File
@@ -520,6 +520,44 @@ fn test_model_client_session() -> crate::client::ModelClientSession {
.new_session()
}
fn session_model_runtime_for_tests(
config: &Config,
auth_manager: &Arc<AuthManager>,
models_manager: &SharedModelsManager,
thread_id: ThreadId,
session_configuration: &SessionConfiguration,
) -> SessionModelRuntime {
let initial =
crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager(
config,
Arc::clone(auth_manager),
Arc::clone(models_manager),
);
let model_client = ModelClient::new_with_provider(
initial.snapshot.provider(),
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
session_configuration.session_source.clone(),
session_configuration.originator.clone(),
config.model_verbosity,
config.features.enabled(Feature::EnableRequestCompression),
config.features.enabled(Feature::RuntimeMetrics),
Session::build_model_client_beta_features_header(config),
/*item_ids_enabled*/ config.features.enabled(Feature::ItemIds),
/*concurrent_reasoning_summaries_enabled*/
config
.features
.enabled(Feature::ConcurrentReasoningSummaries),
/*attestation_provider*/ None,
config.http_client_factory(),
);
SessionModelRuntime {
source: initial.source,
snapshot: initial.snapshot,
model_client,
}
}
fn developer_input_texts(items: &[ResponseItem]) -> Vec<&str> {
items
.iter()
@@ -5287,6 +5325,12 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
/*bundled_skills_enabled*/ true,
));
let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
let model_provider_runtime =
crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager(
config.as_ref(),
Arc::clone(&auth_manager),
Arc::clone(&models_manager),
);
let result = Session::new(
session_configuration,
Arc::clone(&config),
@@ -5294,6 +5338,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
"11111111-1111-4111-8111-111111111111".to_string(),
auth_manager,
models_manager,
model_provider_runtime,
Arc::new(ExecPolicyManager::default()),
tx_event,
agent_status_tx,
@@ -5473,6 +5518,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
show_raw_agent_reasoning: config.show_raw_agent_reasoning,
exec_policy,
auth_manager: auth_manager.clone(),
model_runtime: arc_swap::ArcSwap::from_pointee(session_model_runtime_for_tests(
config.as_ref(),
&auth_manager,
&models_manager,
thread_id,
&session_configuration,
)),
session_telemetry: session_telemetry.clone(),
models_manager: Arc::clone(&models_manager),
tool_approvals: Mutex::new(ApprovalStore::default()),
@@ -5547,12 +5599,13 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
.skills_service
.snapshot_for_config(&skills_input, Some(Arc::clone(&skill_fs)))
.await;
let model_runtime = services.model_runtime.load_full();
let turn_context = Session::make_turn_context(
thread_id,
SessionId::from(thread_id),
Some(Arc::clone(&auth_manager)),
&session_telemetry,
session_configuration.provider.clone(),
model_runtime.as_ref(),
&session_configuration,
config.multi_agent_version_from_features(),
services.user_shell.as_ref(),
@@ -5560,7 +5613,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
services.main_execve_wrapper_exe.as_ref(),
per_turn_config,
model_info,
&models_manager,
/*network*/ None,
resolved_turn_environments,
session_configuration.cwd().clone(),
@@ -5673,6 +5725,12 @@ async fn make_session_with_config_and_rx(
/*bundled_skills_enabled*/ true,
));
let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
let model_provider_runtime =
crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager(
config.as_ref(),
Arc::clone(&auth_manager),
Arc::clone(&models_manager),
);
let session = Session::new(
session_configuration,
@@ -5681,6 +5739,7 @@ async fn make_session_with_config_and_rx(
"11111111-1111-4111-8111-111111111111".to_string(),
auth_manager,
models_manager,
model_provider_runtime,
Arc::new(ExecPolicyManager::default()),
tx_event,
agent_status_tx,
@@ -5781,6 +5840,12 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
/*bundled_skills_enabled*/ true,
));
let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
let model_provider_runtime =
crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager(
config.as_ref(),
Arc::clone(&auth_manager),
Arc::clone(&models_manager),
);
let session = Session::new(
session_configuration,
@@ -5789,6 +5854,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
"11111111-1111-4111-8111-111111111111".to_string(),
auth_manager,
models_manager,
model_provider_runtime,
Arc::new(ExecPolicyManager::default()),
tx_event,
agent_status_tx,
@@ -7604,6 +7670,13 @@ where
show_raw_agent_reasoning: config.show_raw_agent_reasoning,
exec_policy,
auth_manager: Arc::clone(&auth_manager),
model_runtime: arc_swap::ArcSwap::from_pointee(session_model_runtime_for_tests(
config.as_ref(),
&auth_manager,
&models_manager,
thread_id,
&session_configuration,
)),
session_telemetry: session_telemetry.clone(),
models_manager: Arc::clone(&models_manager),
tool_approvals: Mutex::new(ApprovalStore::default()),
@@ -7678,12 +7751,13 @@ where
.skills_service
.snapshot_for_config(&skills_input, Some(Arc::clone(&skill_fs)))
.await;
let model_runtime = services.model_runtime.load_full();
let turn_context = Arc::new(Session::make_turn_context(
thread_id,
SessionId::from(thread_id),
Some(Arc::clone(&auth_manager)),
&session_telemetry,
session_configuration.provider.clone(),
model_runtime.as_ref(),
&session_configuration,
config.multi_agent_version_from_features(),
services.user_shell.as_ref(),
@@ -7691,7 +7765,6 @@ where
services.main_execve_wrapper_exe.as_ref(),
per_turn_config,
model_info,
&models_manager,
/*network*/ None,
resolved_turn_environments,
session_configuration.cwd().clone(),
@@ -715,6 +715,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
codex_thread_store::LocalThreadStoreConfig::from_config(&config),
/*state_db*/ None,
));
let model_provider_runtime =
crate::model_provider_runtime::build_explicit_model_provider_runtime_with_models_manager(
&config,
Arc::clone(&auth_manager),
Arc::clone(&models_manager),
);
let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
config,
@@ -723,6 +729,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
auth_manager,
models_manager,
model_provider_runtime,
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
skills_service,
plugins_manager,
+3 -3
View File
@@ -148,7 +148,7 @@ pub(crate) async fn run_turn(
cancellation_token: CancellationToken,
) -> CodexResult<Option<String>> {
let mut client_session =
prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session());
prewarmed_client_session.unwrap_or_else(|| turn_context.model_client.new_session());
// TODO(ccunningham): Pre-turn compaction runs before context updates and the
// new user message are recorded. Estimate pending incoming items (context
// diffs/full reinjection + user input) and trigger compaction preemptively
@@ -871,7 +871,7 @@ async fn maybe_run_previous_model_inline_compact(
let previous_model = previous_turn_settings.model;
let previous_model_turn_context = Arc::new(
turn_context
.with_model(previous_model.clone(), &sess.services.models_manager)
.with_model(previous_model.clone(), &turn_context.models_manager)
.await,
);
@@ -2271,7 +2271,7 @@ async fn try_run_sampling_request(
}
ResponseEvent::ModelsEtag(etag) => {
// Update internal state with latest models etag
sess.services
turn_context
.models_manager
.refresh_if_new_etag(etag, turn_context.config.http_client_factory())
.await;
+177 -33
View File
@@ -1,10 +1,10 @@
use super::*;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::model_provider_runtime::ModelProviderRuntimeSnapshot;
use crate::shell_snapshot::ShellSnapshotFile;
use codex_core_skills::HostSkillsSnapshot;
use codex_file_system::FileSystemSandboxContext;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_protocol::SessionId;
use codex_protocol::ThreadId;
use codex_protocol::models::AdditionalPermissionProfile;
@@ -109,6 +109,9 @@ pub struct TurnContext {
pub(crate) model_info: ModelInfo,
pub(crate) session_telemetry: SessionTelemetry,
pub(crate) provider: SharedModelProvider,
pub(crate) models_manager: SharedModelsManager,
pub(crate) model_client: ModelClient,
pub(crate) model_provider_generation: u64,
pub(crate) reasoning_effort: Option<ReasoningEffortConfig>,
pub(crate) reasoning_summary: ReasoningSummaryConfig,
pub(crate) session_source: SessionSource,
@@ -269,6 +272,9 @@ impl TurnContext {
.clone()
.with_model(model.as_str(), model_info.slug.as_str()),
provider: self.provider.clone(),
models_manager: Arc::clone(&self.models_manager),
model_client: self.model_client.clone(),
model_provider_generation: self.model_provider_generation,
reasoning_effort,
reasoning_summary: self.reasoning_summary,
session_source: self.session_source.clone(),
@@ -421,7 +427,121 @@ fn local_time_context() -> (String, String) {
}
}
struct PreparedModelProviderAdoption {
current: Arc<SessionModelRuntime>,
latest: Arc<ModelProviderRuntimeSnapshot>,
model: String,
model_info: ModelInfo,
}
impl Session {
async fn prepare_latest_default_model_provider_adoption(
&self,
) -> Option<PreparedModelProviderAdoption> {
let current = self.services.model_runtime.load_full();
let ModelProviderRuntimeSource::RuntimeDefault(runtime) = &current.source else {
return None;
};
let latest = runtime.snapshot();
if latest.generation() == current.snapshot.generation() {
return None;
}
let (refresh_strategy, mut config) = {
let state = self.state.lock().await;
let refresh_strategy = if state
.session_configuration
.session_source
.is_non_root_agent()
{
RefreshStrategy::Offline
} else {
RefreshStrategy::OnlineIfUncached
};
let config = (*state.session_configuration.original_config_do_not_use).clone();
(refresh_strategy, config)
};
let models_manager = latest.models_manager();
let _ = models_manager
.list_models(refresh_strategy, config.http_client_factory())
.await;
let model = models_manager
.get_default_model(
/*model*/ &None,
/*allow_provider_model_fallback*/ true,
refresh_strategy,
config.http_client_factory(),
)
.await;
latest.apply_to_config(&mut config);
config.model = Some(model.clone());
let model_info = models_manager
.get_model_info(model.as_str(), &config.to_models_manager_config())
.await;
Some(PreparedModelProviderAdoption {
current,
latest,
model,
model_info,
})
}
fn apply_default_model_provider_adoption(
&self,
state: &mut SessionState,
prepared: Option<PreparedModelProviderAdoption>,
) -> Arc<SessionModelRuntime> {
let installed = self.services.model_runtime.load_full();
let Some(prepared) = prepared else {
return installed;
};
if !Arc::ptr_eq(&installed.snapshot, &prepared.current.snapshot) {
return installed;
}
let PreparedModelProviderAdoption {
current,
latest,
model,
model_info,
} = prepared;
let supported_reasoning_levels = model_info
.supported_reasoning_levels
.iter()
.map(|preset| preset.effort.clone())
.collect::<Vec<_>>();
let current_reasoning_effort = state
.session_configuration
.collaboration_mode
.reasoning_effort();
let reasoning_effort = current_reasoning_effort
.filter(|effort| supported_reasoning_levels.contains(effort))
.or_else(|| model_info.default_reasoning_level.clone());
let mut config = (*state.session_configuration.original_config_do_not_use).clone();
latest.apply_to_config(&mut config);
config.model = Some(model.clone());
config.model_reasoning_effort = reasoning_effort.clone();
state.session_configuration.provider = latest.provider_info().clone();
state.session_configuration.collaboration_mode =
state.session_configuration.collaboration_mode.with_updates(
Some(model),
Some(reasoning_effort),
/*developer_instructions*/ None,
);
state.session_configuration.original_config_do_not_use = Arc::new(config);
let next = Arc::new(SessionModelRuntime {
source: current.source.clone(),
snapshot: latest.clone(),
model_client: current.model_client.rebind_provider(latest.provider()),
});
self.services.model_runtime.store(Arc::clone(&next));
// A prewarmed websocket is bound to the provider runtime that created it.
let _ = state.take_session_startup_prewarm();
next
}
/// Don't expand the number of mutated arguments on config. We are in the process of getting rid of it.
pub(crate) fn build_per_turn_config(
session_configuration: &SessionConfiguration,
@@ -482,7 +602,7 @@ impl Session {
session_id: SessionId,
auth_manager: Option<Arc<AuthManager>>,
session_telemetry: &SessionTelemetry,
provider: ModelProviderInfo,
model_runtime: &SessionModelRuntime,
session_configuration: &SessionConfiguration,
multi_agent_version: MultiAgentVersion,
user_shell: &shell::Shell,
@@ -490,7 +610,6 @@ impl Session {
main_execve_wrapper_exe: Option<&PathBuf>,
per_turn_config: Config,
model_info: ModelInfo,
models_manager: &SharedModelsManager,
network: Option<NetworkProxy>,
environments: TurnEnvironmentSnapshot,
cwd: AbsolutePathBuf,
@@ -506,8 +625,9 @@ impl Session {
model_info.slug.as_str(),
);
let session_source = session_configuration.session_source.clone();
let auth_manager_for_context = auth_manager.clone();
let provider_for_context = create_model_provider(provider, auth_manager);
let auth_manager_for_context = auth_manager;
let provider_for_context = model_runtime.snapshot.provider();
let models_manager = model_runtime.models_manager();
let session_telemetry_for_context = session_telemetry;
let available_models = models_manager.try_list_models().unwrap_or_default();
let unified_exec_shell_mode = UnifiedExecShellMode::for_session(
@@ -549,6 +669,9 @@ impl Session {
model_info,
session_telemetry: session_telemetry_for_context,
provider: provider_for_context,
models_manager,
model_client: model_runtime.model_client.clone(),
model_provider_generation: model_runtime.snapshot.generation(),
reasoning_effort,
reasoning_summary,
session_source,
@@ -589,8 +712,11 @@ impl Session {
updates: SessionSettingsUpdate,
) -> CodexResult<Arc<TurnContext>> {
let notify_config_contributors = !self.services.extensions.config_contributors().is_empty();
let prepared_model_provider = self.prepare_latest_default_model_provider_adoption().await;
let update_result: CodexResult<_> = {
let mut state = self.state.lock().await;
let model_runtime =
self.apply_default_model_provider_adoption(&mut state, prepared_model_provider);
match state.session_configuration.clone().apply(&updates) {
Ok(next) => {
let previous_permission_profile =
@@ -611,6 +737,7 @@ impl Session {
state.session_configuration = next.clone();
Ok((
next,
model_runtime,
permission_profile_changed,
previous_config,
new_config,
@@ -620,22 +747,27 @@ impl Session {
}
};
let (session_configuration, permission_profile_changed, previous_config, new_config) =
match update_result {
Ok(update) => update,
Err(err) => {
let message = err.to_string();
self.send_event_raw(Event {
id: sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: message.clone(),
codex_error_info: Some(CodexErrorInfo::BadRequest),
}),
})
.await;
return Err(CodexErr::InvalidRequest(message));
}
};
let (
session_configuration,
model_runtime,
permission_profile_changed,
previous_config,
new_config,
) = match update_result {
Ok(update) => update,
Err(err) => {
let message = err.to_string();
self.send_event_raw(Event {
id: sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: message.clone(),
codex_error_info: Some(CodexErrorInfo::BadRequest),
}),
})
.await;
return Err(CodexErr::InvalidRequest(message));
}
};
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
if permission_profile_changed {
@@ -647,6 +779,7 @@ impl Session {
.new_turn_from_configuration(
sub_id,
session_configuration,
model_runtime,
updates.final_output_json_schema,
)
.await)
@@ -656,11 +789,13 @@ impl Session {
&self,
sub_id: String,
session_configuration: SessionConfiguration,
model_runtime: Arc<SessionModelRuntime>,
final_output_json_schema: Option<Option<Value>>,
) -> Arc<TurnContext> {
self.new_turn_context_from_configuration(
sub_id,
session_configuration,
model_runtime,
final_output_json_schema,
TurnMultiAgentRuntime::ResolveAndStore,
)
@@ -671,10 +806,12 @@ impl Session {
&self,
sub_id: String,
session_configuration: SessionConfiguration,
model_runtime: Arc<SessionModelRuntime>,
) -> Arc<TurnContext> {
self.new_turn_context_from_configuration(
sub_id,
session_configuration,
model_runtime,
/*final_output_json_schema*/ None,
TurnMultiAgentRuntime::Preview,
)
@@ -686,6 +823,7 @@ impl Session {
&self,
sub_id: String,
session_configuration: SessionConfiguration,
model_runtime: Arc<SessionModelRuntime>,
final_output_json_schema: Option<Option<Value>>,
multi_agent_runtime: TurnMultiAgentRuntime,
) -> Arc<TurnContext> {
@@ -706,9 +844,8 @@ impl Session {
.set_permission_profile(session_configuration.permission_profile());
}
let model_info = self
.services
.models_manager
let model_info = model_runtime
.models_manager()
.get_model_info(
session_configuration.collaboration_mode.model(),
&per_turn_config.to_models_manager_config(),
@@ -752,7 +889,7 @@ impl Session {
self.session_id(),
Some(Arc::clone(&self.services.auth_manager)),
&self.services.session_telemetry,
session_configuration.provider.clone(),
model_runtime.as_ref(),
&session_configuration,
multi_agent_version,
self.services.user_shell.as_ref(),
@@ -760,7 +897,6 @@ impl Session {
self.services.main_execve_wrapper_exe.as_ref(),
per_turn_config,
model_info,
&self.services.models_manager,
self.services
.network_proxy
.load_full()
@@ -820,10 +956,11 @@ impl Session {
}
pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc<TurnContext> {
let session_configuration = self.default_turn_configuration().await;
let (session_configuration, model_runtime) = self.default_turn_configuration().await;
self.new_turn_from_configuration(
sub_id,
session_configuration,
model_runtime,
/*final_output_json_schema*/ None,
)
.await
@@ -833,13 +970,20 @@ impl Session {
&self,
sub_id: String,
) -> Arc<TurnContext> {
let session_configuration = self.default_turn_configuration().await;
self.new_startup_prewarm_turn_from_configuration(sub_id, session_configuration)
.await
let (session_configuration, model_runtime) = self.default_turn_configuration().await;
self.new_startup_prewarm_turn_from_configuration(
sub_id,
session_configuration,
model_runtime,
)
.await
}
async fn default_turn_configuration(&self) -> SessionConfiguration {
let state = self.state.lock().await;
state.session_configuration.clone()
async fn default_turn_configuration(&self) -> (SessionConfiguration, Arc<SessionModelRuntime>) {
let prepared_model_provider = self.prepare_latest_default_model_provider_adoption().await;
let mut state = self.state.lock().await;
let model_runtime =
self.apply_default_model_provider_adoption(&mut state, prepared_model_provider);
(state.session_configuration.clone(), model_runtime)
}
}
+3 -3
View File
@@ -182,10 +182,10 @@ impl SessionStartupPrewarmHandle {
impl Session {
pub(crate) async fn schedule_startup_prewarm(self: &Arc<Self>, base_instructions: String) {
if !self.services.model_client.responses_websocket_enabled() {
let model_client = self.services.model_runtime.load().model_client.clone();
if !model_client.responses_websocket_enabled() {
// Without websocket prewarm, resolve auth once so Agent Identity bootstrap can
// register or engage this session's bearer fallback before the first user request.
let model_client = self.services.model_client.clone();
tokio::spawn(async move {
if let Err(err) = model_client.prewarm_auth().await {
warn!("startup auth prewarm failed: {err:#}");
@@ -303,7 +303,7 @@ async fn schedule_startup_prewarm_inner(
window_id,
CodexResponsesRequestKind::Prewarm,
);
let mut client_session = session.services.model_client.new_session();
let mut client_session = startup_turn_context.model_client.new_session();
let websocket_warmup_started_at = Instant::now();
client_session
.prewarm_websocket(
+1
View File
@@ -7,6 +7,7 @@ mod turn;
pub(crate) use additional_context::AdditionalContextStore;
pub(crate) use auto_compact_window::AutoCompactWindowIds;
pub(crate) use auto_compact_window::AutoCompactWindowSnapshot;
pub(crate) use service::SessionModelRuntime;
pub(crate) use service::SessionServices;
pub(crate) use session::SessionState;
pub(crate) use turn::ActiveTurn;
+18
View File
@@ -16,6 +16,8 @@ use crate::exec_policy::ExecPolicyManager;
use crate::guardian::GuardianRejection;
use crate::guardian::GuardianRejectionCircuitBreaker;
use crate::mcp::McpManager;
use crate::model_provider_runtime::ModelProviderRuntimeSnapshot;
use crate::model_provider_runtime::ModelProviderRuntimeSource;
use crate::session::McpRuntimeSnapshot;
use crate::tools::code_mode::CodeModeService;
use crate::tools::handlers::ToolSearchHandlerCache;
@@ -68,6 +70,9 @@ pub(crate) struct SessionServices {
pub(crate) show_raw_agent_reasoning: bool,
pub(crate) exec_policy: Arc<ExecPolicyManager>,
pub(crate) auth_manager: Arc<AuthManager>,
/// Atomically replaceable provider runtime used to build future turns. Each turn clones one
/// immutable runtime so an in-flight turn cannot observe a provider refresh.
pub(crate) model_runtime: ArcSwap<SessionModelRuntime>,
pub(crate) models_manager: SharedModelsManager,
pub(crate) session_telemetry: SessionTelemetry,
pub(crate) tool_approvals: Mutex<ApprovalStore>,
@@ -103,6 +108,19 @@ pub(crate) struct SessionServices {
pub(crate) turn_environments: Arc<ThreadEnvironments>,
}
#[derive(Clone)]
pub(crate) struct SessionModelRuntime {
pub(crate) source: ModelProviderRuntimeSource,
pub(crate) snapshot: Arc<ModelProviderRuntimeSnapshot>,
pub(crate) model_client: ModelClient,
}
impl SessionModelRuntime {
pub(crate) fn models_manager(&self) -> SharedModelsManager {
self.snapshot.models_manager()
}
}
impl SessionServices {
/// Installs the manager before validating required servers so startup-time elicitation can
/// resolve through the session's manager while validation waits.
-5
View File
@@ -37,7 +37,6 @@ use crate::state::TaskKind;
use codex_analytics::TurnProfileFact;
use codex_analytics::TurnTokenUsageFact;
use codex_login::AuthManager;
use codex_models_manager::manager::SharedModelsManager;
use codex_otel::SessionTelemetry;
use codex_otel::TURN_E2E_DURATION_METRIC;
use codex_otel::TURN_MEMORY_METRIC;
@@ -197,10 +196,6 @@ impl SessionTaskContext {
pub(crate) fn auth_manager(&self) -> Arc<AuthManager> {
Arc::clone(&self.session.services.auth_manager)
}
pub(crate) fn models_manager(&self) -> SharedModelsManager {
Arc::clone(&self.session.services.models_manager)
}
}
/// Async task that drives a [`Session`] turn.
+1 -1
View File
@@ -125,7 +125,7 @@ async fn start_review_conversation(
(run_codex_thread_one_shot(
sub_agent_config,
session.auth_manager(),
session.models_manager(),
Arc::clone(&ctx.models_manager),
input,
session.clone_session(),
ctx.clone(),
+67 -13
View File
@@ -9,6 +9,9 @@ use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::environment_selection::default_thread_environment_selections;
use crate::mcp::McpManager;
use crate::model_provider_runtime::DefaultModelProviderRuntime;
use crate::model_provider_runtime::InitialModelProviderRuntime;
use crate::model_provider_runtime::ModelProviderRuntimeSource;
use crate::model_provider_runtime::build_explicit_model_provider_runtime;
use crate::rollout::truncation;
use crate::session::Codex;
use crate::session::CodexSpawnArgs;
@@ -200,6 +203,14 @@ pub struct StartThreadOptions {
pub supports_openai_form_elicitation: bool,
}
/// Determines whether a thread follows the process-default model provider or keeps its initial
/// provider pinned for the lifetime of the loaded session.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ModelProviderSelection {
RuntimeDefault,
Explicit,
}
fn originator_from_service_name(service_name: Option<&str>) -> Option<String> {
let service_name = service_name?.trim();
for originator in ["codex_work_desktop", "codex_work_web", "codex_work_mobile"] {
@@ -241,7 +252,7 @@ pub(crate) struct ThreadManagerState {
threads: Arc<RwLock<HashMap<ThreadId, Arc<CodexThread>>>>,
thread_created_tx: broadcast::Sender<ThreadId>,
auth_manager: Arc<AuthManager>,
default_model_provider_runtime: DefaultModelProviderRuntime,
default_model_provider_runtime: Arc<DefaultModelProviderRuntime>,
environment_manager: Arc<EnvironmentManager>,
skills_service: Arc<SkillsService>,
plugins_manager: Arc<PluginsManager>,
@@ -334,10 +345,10 @@ impl ThreadManager {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
thread_created_tx,
default_model_provider_runtime: DefaultModelProviderRuntime::new(
default_model_provider_runtime: Arc::new(DefaultModelProviderRuntime::new(
config,
auth_manager.clone(),
),
)),
environment_manager,
skills_service,
plugins_manager,
@@ -456,13 +467,13 @@ impl ThreadManager {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
thread_created_tx,
default_model_provider_runtime: DefaultModelProviderRuntime::from_parts(
default_model_provider_runtime: Arc::new(DefaultModelProviderRuntime::from_parts(
OPENAI_PROVIDER_ID.to_string(),
provider,
codex_home,
/*config_model_catalog*/ None,
auth_manager.clone(),
),
)),
environment_manager,
skills_service,
plugins_manager,
@@ -706,14 +717,31 @@ impl ThreadManager {
&self,
options: StartThreadOptions,
) -> CodexResult<NewThread> {
self.start_thread_with_options_and_fork_source(options, /*forked_from_thread_id*/ None)
.await
self.start_thread_with_options_and_provider_selection(
options,
ModelProviderSelection::Explicit,
)
.await
}
pub async fn start_thread_with_options_and_provider_selection(
&self,
options: StartThreadOptions,
model_provider_selection: ModelProviderSelection,
) -> CodexResult<NewThread> {
self.start_thread_with_options_and_fork_source(
options,
/*forked_from_thread_id*/ None,
model_provider_selection,
)
.await
}
async fn start_thread_with_options_and_fork_source(
&self,
options: StartThreadOptions,
forked_from_thread_id: Option<ThreadId>,
model_provider_selection: ModelProviderSelection,
) -> CodexResult<NewThread> {
let agent_control = self.agent_control_for_config(&options.config);
let (resumed_session_source, resumed_thread_source) = options
@@ -742,6 +770,7 @@ impl ThreadManager {
options.thread_extension_init,
options.supports_openai_form_elicitation,
/*user_shell_override*/ None,
model_provider_selection,
))
.await
}
@@ -779,8 +808,12 @@ impl ThreadManager {
inherited_multi_agent_version,
),
);
self.start_thread_with_options_and_fork_source(options, Some(forked_from_thread_id))
.await
self.start_thread_with_options_and_fork_source(
options,
Some(forked_from_thread_id),
ModelProviderSelection::Explicit,
)
.await
}
pub async fn resume_thread_from_rollout(
@@ -839,6 +872,7 @@ impl ThreadManager {
/*thread_extension_init*/ ExtensionDataInit::default(),
supports_openai_form_elicitation,
/*user_shell_override*/ None,
ModelProviderSelection::Explicit,
))
.await
}
@@ -910,6 +944,7 @@ impl ThreadManager {
/*thread_extension_init*/ ExtensionDataInit::default(),
supports_openai_form_elicitation,
/*user_shell_override*/ Some(user_shell_override),
ModelProviderSelection::Explicit,
))
.await
}
@@ -1412,6 +1447,7 @@ impl ThreadManagerState {
/*thread_extension_init*/ ExtensionDataInit::default(),
/*supports_openai_form_elicitation*/ false,
/*user_shell_override*/ None,
ModelProviderSelection::Explicit,
))
.await
}
@@ -1452,6 +1488,7 @@ impl ThreadManagerState {
/*thread_extension_init*/ ExtensionDataInit::default(),
/*supports_openai_form_elicitation*/ false,
/*user_shell_override*/ None,
ModelProviderSelection::Explicit,
))
.await
}
@@ -1494,6 +1531,7 @@ impl ThreadManagerState {
thread_extension_init,
/*supports_openai_form_elicitation*/ false,
/*user_shell_override*/ None,
ModelProviderSelection::Explicit,
))
.await
}
@@ -1537,6 +1575,7 @@ impl ThreadManagerState {
thread_extension_init,
supports_openai_form_elicitation,
user_shell_override,
ModelProviderSelection::Explicit,
))
.await
}
@@ -1544,7 +1583,7 @@ impl ThreadManagerState {
#[allow(clippy::too_many_arguments)]
pub(crate) async fn spawn_thread_with_source(
&self,
config: Config,
mut config: Config,
initial_history: InitialHistory,
history_mode: Option<ThreadHistoryMode>,
allow_provider_model_fallback: bool,
@@ -1563,6 +1602,7 @@ impl ThreadManagerState {
thread_extension_init: ExtensionDataInit,
supports_openai_form_elicitation: bool,
user_shell_override: Option<crate::shell::Shell>,
model_provider_selection: ModelProviderSelection,
) -> CodexResult<NewThread> {
let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
if let InitialHistory::Resumed(resumed) = &initial_history {
@@ -1610,9 +1650,22 @@ impl ThreadManagerState {
forked_from_thread_id,
)
.await;
let models_manager = self
.default_model_provider_runtime
.models_manager_for_config(&config, Arc::clone(&auth_manager));
let model_provider_runtime = match model_provider_selection {
ModelProviderSelection::RuntimeDefault => {
let snapshot = self.default_model_provider_runtime.snapshot();
snapshot.apply_to_config(&mut config);
InitialModelProviderRuntime {
source: ModelProviderRuntimeSource::RuntimeDefault(Arc::clone(
&self.default_model_provider_runtime,
)),
snapshot,
}
}
ModelProviderSelection::Explicit => {
build_explicit_model_provider_runtime(&config, Arc::clone(&auth_manager))
}
};
let models_manager = model_provider_runtime.snapshot.models_manager();
let CodexSpawnOk {
codex, thread_id, ..
} = Box::pin(Codex::spawn(CodexSpawnArgs {
@@ -1622,6 +1675,7 @@ impl ThreadManagerState {
installation_id: self.installation_id.clone(),
auth_manager,
models_manager,
model_provider_runtime,
environment_manager: Arc::clone(&self.environment_manager),
skills_service: Arc::clone(&self.skills_service),
plugins_manager: Arc::clone(&self.plugins_manager),
+134 -4
View File
@@ -64,6 +64,34 @@ async fn listed_model_ids(models_manager: &SharedModelsManager) -> Vec<String> {
.collect()
}
async fn start_thread_with_provider_selection(
manager: &ThreadManager,
config: Config,
model_provider_selection: ModelProviderSelection,
) -> NewThread {
let environments = manager.default_environment_selections(&config.cwd);
manager
.start_thread_with_options_and_provider_selection(
StartThreadOptions {
config,
allow_provider_model_fallback: false,
initial_history: InitialHistory::New,
history_mode: None,
session_source: None,
thread_source: None,
dynamic_tools: Vec::new(),
metrics_service_name: None,
parent_trace: None,
environments,
thread_extension_init: ExtensionDataInit::default(),
supports_openai_form_elicitation: false,
},
model_provider_selection,
)
.await
.expect("start thread")
}
struct FakeAgentGraphStore {
root_thread_id: ThreadId,
descendant_thread_ids: Vec<ThreadId>,
@@ -1492,10 +1520,12 @@ async fn new_threads_select_models_manager_for_their_effective_config() {
vec!["model-b".to_string()]
);
let default_thread = manager
.start_thread(provider_b_config)
.await
.expect("start default-provider thread");
let default_thread = start_thread_with_provider_selection(
&manager,
provider_b_config,
ModelProviderSelection::RuntimeDefault,
)
.await;
assert!(Arc::ptr_eq(
&default_thread.thread.codex.session.services.models_manager,
&provider_b_models_manager,
@@ -1529,6 +1559,106 @@ async fn new_threads_select_models_manager_for_their_effective_config() {
);
}
#[tokio::test]
async fn loaded_default_provider_session_adopts_refresh_between_turns() {
let temp_dir = tempdir().expect("tempdir");
let mut provider_a_config = test_config().await;
provider_a_config.codex_home = temp_dir.path().join("codex-home").abs();
provider_a_config.cwd = provider_a_config.codex_home.abs();
std::fs::create_dir_all(&provider_a_config.codex_home).expect("create codex home");
provider_a_config.model = None;
provider_a_config.model_provider_id = "provider-a".to_string();
provider_a_config.model_provider.name = "Provider A".to_string();
provider_a_config.model_catalog = Some(static_model_catalog("model-a"));
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy"));
let manager = ThreadManager::new(
&provider_a_config,
auth_manager,
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
empty_extension_registry(),
Arc::new(crate::test_support::EmptyUserInstructionsProvider),
/*analytics_events_client*/ None,
thread_store_from_config(&provider_a_config, /*state_db*/ None),
/*agent_graph_store*/ None,
TEST_INSTALLATION_ID.to_string(),
/*attestation_provider*/ None,
/*external_time_provider*/ None,
);
let default_thread = start_thread_with_provider_selection(
&manager,
provider_a_config.clone(),
ModelProviderSelection::RuntimeDefault,
)
.await;
let explicit_thread = start_thread_with_provider_selection(
&manager,
provider_a_config.clone(),
ModelProviderSelection::Explicit,
)
.await;
let captured_a_turn = default_thread.thread.codex.session.new_default_turn().await;
assert_eq!(captured_a_turn.config.model_provider_id, "provider-a");
assert_eq!(captured_a_turn.model_info.slug, "model-a");
assert_eq!(captured_a_turn.model_provider_generation, 0);
let mut provider_b_config = provider_a_config;
provider_b_config.model_provider_id = "provider-b".to_string();
provider_b_config.model_provider.name = "Provider B".to_string();
provider_b_config.model_catalog = Some(static_model_catalog("model-b"));
assert!(manager.refresh_default_model_provider(&provider_b_config));
// An already-created turn retains the exact provider runtime it captured.
assert_eq!(captured_a_turn.config.model_provider_id, "provider-a");
assert_eq!(captured_a_turn.model_info.slug, "model-a");
assert_eq!(
listed_model_ids(&captured_a_turn.models_manager).await,
vec!["model-a".to_string()]
);
let adopted_turn = default_thread.thread.codex.session.new_default_turn().await;
assert_eq!(adopted_turn.config.model_provider_id, "provider-b");
assert_eq!(adopted_turn.model_info.slug, "model-b");
assert_eq!(adopted_turn.model_provider_generation, 1);
assert_eq!(
listed_model_ids(&adopted_turn.models_manager).await,
vec!["model-b".to_string()]
);
assert_eq!(
default_thread
.thread
.config_snapshot()
.await
.model_provider_id,
"provider-b"
);
let pinned_turn = explicit_thread
.thread
.codex
.session
.new_default_turn()
.await;
assert_eq!(pinned_turn.config.model_provider_id, "provider-a");
assert_eq!(pinned_turn.model_info.slug, "model-a");
assert_eq!(pinned_turn.model_provider_generation, 0);
let mut expected_completed = vec![default_thread.thread_id, explicit_thread.thread_id];
expected_completed.sort_by_key(std::string::ToString::to_string);
assert_eq!(
manager
.shutdown_all_threads_bounded(Duration::from_secs(10))
.await,
ThreadShutdownReport {
completed: expected_completed,
submit_failed: Vec::new(),
timed_out: Vec::new(),
}
);
}
#[test]
fn interrupted_fork_snapshot_appends_interrupt_boundary() {
let committed_history =
+91 -1
View File
@@ -1,21 +1,34 @@
use std::fmt;
use std::path::Path;
use codex_config::types::AuthCredentialsStoreMode;
use serde::Deserialize;
use serde::Serialize;
use super::manager::AuthManager;
use super::manager::load_auth_dot_json;
use super::manager::logout;
use super::manager::save_auth;
use super::storage::AuthDotJson;
use super::storage::AuthKeyringBackendKind;
use codex_protocol::auth::AuthMode;
/// Managed Amazon Bedrock API key persisted in `auth.json`.
#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
#[derive(Deserialize, Serialize, Clone, PartialEq, Eq)]
pub struct BedrockApiKeyAuth {
pub api_key: String,
pub region: String,
}
impl fmt::Debug for BedrockApiKeyAuth {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BedrockApiKeyAuth")
.field("api_key", &"<redacted>")
.field("region", &self.region)
.finish()
}
}
/// Writes an `auth.json` that contains only the Amazon Bedrock API key auth.
pub fn login_with_bedrock_api_key(
codex_home: &Path,
@@ -44,6 +57,83 @@ pub fn login_with_bedrock_api_key(
)
}
pub(super) fn load_stored_bedrock_api_key(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> std::io::Result<Option<BedrockApiKeyAuth>> {
load_stored_bedrock_api_key_auth(
codex_home,
auth_credentials_store_mode,
keyring_backend_kind,
)
.map(|auth| auth.and_then(|auth| auth.bedrock_api_key))
}
fn load_stored_bedrock_api_key_auth(
codex_home: &Path,
auth_credentials_store_mode: AuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
) -> std::io::Result<Option<AuthDotJson>> {
load_auth_dot_json(
codex_home,
auth_credentials_store_mode,
keyring_backend_kind,
)
.map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key))
}
impl AuthManager {
pub fn has_stored_bedrock_api_key_auth(&self) -> std::io::Result<bool> {
load_stored_bedrock_api_key_auth(
&self.codex_home,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
)
.map(|auth| auth.is_some())
}
/// Persists managed Bedrock auth without changing the in-memory auth snapshot.
///
/// The account coordinator publishes the provider-scoped cache update only
/// after the associated config mutation succeeds.
pub fn persist_bedrock_api_key_auth(&self, api_key: &str, region: &str) -> std::io::Result<()> {
login_with_bedrock_api_key(
&self.codex_home,
api_key,
region,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
)
}
/// Reloads only managed Bedrock auth, preserving cached general Codex auth.
pub fn reload_bedrock_api_key_auth(&self) -> std::io::Result<bool> {
let bedrock_api_key = load_stored_bedrock_api_key(
&self.codex_home,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
)?;
Ok(self.set_cached_bedrock_api_key_auth(bedrock_api_key))
}
/// Removes only managed Bedrock auth, preserving cached general Codex auth.
pub fn logout_bedrock_api_key_auth(&self) -> std::io::Result<bool> {
let stored_bedrock_auth = self.has_stored_bedrock_api_key_auth()?;
let removed = if stored_bedrock_auth {
logout(
&self.codex_home,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
)?
} else {
false
};
self.set_cached_bedrock_api_key_auth(/*new_bedrock_api_key*/ None);
Ok(removed)
}
}
#[cfg(test)]
#[path = "bedrock_api_key_tests.rs"]
mod tests;
+106 -32
View File
@@ -1,13 +1,15 @@
use codex_config::types::AuthCredentialsStoreMode;
use codex_protocol::auth::AuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
use pretty_assertions::assert_eq;
use serial_test::serial;
use tempfile::tempdir;
use super::*;
use crate::auth::AuthConfig;
use crate::auth::AuthKeyringBackendKind;
use crate::auth::AuthManager;
use crate::auth::CodexAuth;
use crate::auth::enforce_login_restrictions;
use crate::auth::storage::AuthStorageBackend;
use crate::auth::storage::FileAuthStorage;
@@ -42,20 +44,24 @@ fn bedrock_auth() -> BedrockApiKeyAuth {
}
}
fn auth_config(codex_home: &std::path::Path, forced_login_method: ForcedLoginMethod) -> AuthConfig {
AuthConfig {
codex_home: codex_home.to_path_buf(),
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
keyring_backend_kind: AuthKeyringBackendKind::default(),
forced_login_method: Some(forced_login_method),
chatgpt_base_url: None,
forced_chatgpt_workspace_id: None,
auth_route_config: None,
}
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> {
async fn managed_bedrock_persistence_preserves_cached_openai_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
storage.save(&api_key_auth())?;
login_with_bedrock_api_key(
codex_home.path(),
"bedrock-api-key-test",
"us-east-1",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
let auth_manager = AuthManager::new(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
@@ -67,6 +73,8 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()>
)
.await;
auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?;
let loaded = storage.load()?.expect("auth should be stored");
let expected = AuthDotJson {
auth_mode: Some(AuthMode::BedrockApiKey),
@@ -78,25 +86,33 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()>
bedrock_api_key: Some(bedrock_auth()),
};
assert_eq!(loaded, expected);
assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey));
assert_eq!(auth_manager.auth_mode(), Some(AuthMode::ApiKey));
assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None);
assert!(auth_manager.reload_bedrock_api_key_auth()?);
assert_eq!(
auth_manager.auth_cached().and_then(|auth| match auth {
CodexAuth::BedrockApiKey(auth) => Some(auth),
CodexAuth::ApiKey(_)
| CodexAuth::Chatgpt(_)
| CodexAuth::ChatgptAuthTokens(_)
| CodexAuth::Headers(_)
| CodexAuth::AgentIdentity(_)
| CodexAuth::PersonalAccessToken(_) => None,
}),
auth_manager
.auth_cached()
.and_then(|auth| auth.api_key().map(str::to_string)),
Some("sk-test-key".to_string())
);
assert_eq!(
auth_manager.bedrock_api_key_auth_cached(),
Some(bedrock_auth())
);
Ok(())
}
#[test]
fn bedrock_auth_debug_redacts_api_key() {
assert_eq!(
format!("{:?}", bedrock_auth()),
"BedrockApiKeyAuth { api_key: \"<redacted>\", region: \"us-east-1\" }"
);
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn logout_removes_bedrock_auth() -> anyhow::Result<()> {
async fn auth_manager_logout_removes_bedrock_api_key_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
login_with_bedrock_api_key(
@@ -121,12 +137,13 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> {
assert_eq!(storage.load()?, None);
assert_eq!(auth_manager.auth_cached(), None);
assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None);
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> {
async fn bedrock_only_auth_storage_does_not_create_codex_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
storage.save(&bedrock_only_auth())?;
@@ -142,19 +159,44 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()>
)
.await;
assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey));
assert_eq!(auth_manager.auth_mode(), None);
assert_eq!(
auth_manager.auth_cached().and_then(|auth| match auth {
CodexAuth::BedrockApiKey(auth) => Some(auth),
CodexAuth::ApiKey(_)
| CodexAuth::Chatgpt(_)
| CodexAuth::ChatgptAuthTokens(_)
| CodexAuth::Headers(_)
| CodexAuth::AgentIdentity(_)
| CodexAuth::PersonalAccessToken(_) => None,
}),
auth_manager.bedrock_api_key_auth_cached(),
Some(bedrock_auth())
);
assert_eq!(auth_manager.auth_cached(), None);
Ok(())
}
#[tokio::test]
#[serial(codex_auth_env)]
async fn bedrock_logout_preserves_cached_openai_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
storage.save(&api_key_auth())?;
let auth_manager = AuthManager::new(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
/*forced_chatgpt_workspace_id*/ None,
/*chatgpt_base_url*/ None,
AuthKeyringBackendKind::default(),
/*auth_route_config*/ None,
)
.await;
auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?;
auth_manager.reload_bedrock_api_key_auth()?;
assert!(auth_manager.logout_bedrock_api_key_auth()?);
assert_eq!(storage.load()?, None);
assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None);
assert_eq!(
auth_manager
.auth_cached()
.and_then(|auth| auth.api_key().map(str::to_string)),
Some("sk-test-key".to_string())
);
Ok(())
}
@@ -180,3 +222,35 @@ async fn login_with_api_key_clears_bedrock_api_key() -> anyhow::Result<()> {
assert_eq!(storage.load()?, Some(api_key_auth()));
Ok(())
}
#[tokio::test]
async fn forced_chatgpt_login_removes_bedrock_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
storage.save(&bedrock_only_auth())?;
let error =
enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Chatgpt))
.await
.expect_err("Bedrock auth should violate forced ChatGPT login");
assert_eq!(
error.to_string(),
"ChatGPT login is required, but an API key is currently being used. Logging out."
);
assert_eq!(storage.load()?, None);
Ok(())
}
#[tokio::test]
async fn forced_api_login_allows_bedrock_auth() -> anyhow::Result<()> {
let codex_home = tempdir()?;
let storage = FileAuthStorage::new(codex_home.path().to_path_buf());
let auth = bedrock_only_auth();
storage.save(&auth)?;
enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Api)).await?;
assert_eq!(storage.load()?, Some(auth));
Ok(())
}
+111 -11
View File
@@ -36,6 +36,7 @@ use super::agent_identity::record_needs_task_registration;
use super::agent_identity::register_managed_chatgpt_agent_identity;
use super::agent_identity::require_agent_identity_authapi_base_url;
use super::agent_identity::verified_record_from_jwt;
use super::bedrock_api_key::load_stored_bedrock_api_key;
use super::external_bearer::BearerTokenRefresher;
use super::revoke::revoke_auth_tokens;
use crate::auth::AuthHeaders;
@@ -1063,7 +1064,7 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url(
config: &AuthConfig,
agent_identity_authapi_base_url: Option<&str>,
) -> std::io::Result<()> {
let Some(auth) = load_auth(
let auth = load_auth(
&config.codex_home,
/*enable_codex_api_key_env*/ true,
config.auth_credentials_store_mode,
@@ -1073,13 +1074,23 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url(
agent_identity_authapi_base_url,
config.auth_route_config.as_ref(),
)
.await?
else {
.await?;
let stored_bedrock_api_key = load_stored_bedrock_api_key(
&config.codex_home,
config.auth_credentials_store_mode,
config.keyring_backend_kind,
)?;
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode).or_else(|| {
stored_bedrock_api_key
.as_ref()
.map(|_| AuthMode::BedrockApiKey)
});
let Some(auth_mode) = auth_mode else {
return Ok(());
};
if let Some(required_method) = config.forced_login_method {
let method_violation = match (required_method, auth.auth_mode()) {
let method_violation = match (required_method, auth_mode) {
(ForcedLoginMethod::Api, AuthMode::ApiKey)
| (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) => None,
(ForcedLoginMethod::Chatgpt, AuthMode::Chatgpt)
@@ -1113,7 +1124,10 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url(
}
if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() {
let chatgpt_account_id = match &auth {
let Some(auth) = auth.as_ref() else {
return Ok(());
};
let chatgpt_account_id = match auth {
CodexAuth::ApiKey(_) | CodexAuth::Headers(_) | CodexAuth::BedrockApiKey(_) => {
return Ok(());
}
@@ -1234,7 +1248,10 @@ async fn load_auth(
AuthCredentialsStoreMode::Ephemeral,
AuthKeyringBackendKind::default(),
);
if let Some(auth_dot_json) = ephemeral_storage.load()? {
if let Some(auth_dot_json) = ephemeral_storage
.load()?
.filter(|auth| !auth.is_bedrock_api_key())
{
let auth = CodexAuth::from_auth_dot_json(
codex_home,
auth_dot_json,
@@ -1283,7 +1300,8 @@ async fn load_auth(
keyring_backend_kind,
);
let auth_dot_json = match storage.load()? {
Some(auth) => auth,
Some(auth) if !auth.is_bedrock_api_key() => auth,
Some(_) => return Ok(None),
None => return Ok(None),
};
@@ -1458,6 +1476,11 @@ fn refresh_token_endpoint() -> String {
}
impl AuthDotJson {
/// Returns whether this stored payload resolves to managed Amazon Bedrock API key auth.
pub fn is_bedrock_api_key(&self) -> bool {
self.resolved_mode() == AuthMode::BedrockApiKey
}
fn from_external_access_token(
access_token: &str,
chatgpt_account_id: &str,
@@ -1520,6 +1543,7 @@ impl AuthDotJson {
#[derive(Clone)]
struct CachedAuth {
auth: Option<CodexAuth>,
bedrock_api_key: Option<BedrockApiKeyAuth>,
/// Permanent refresh failure cached for the current auth snapshot so
/// later refresh attempts for the same credentials fail fast without network.
permanent_refresh_failure: Option<AuthScopedRefreshFailure>,
@@ -1538,6 +1562,7 @@ impl Debug for CachedAuth {
"auth_mode",
&self.auth.as_ref().map(CodexAuth::api_auth_mode),
)
.field("has_bedrock_api_key", &self.bedrock_api_key.is_some())
.field(
"permanent_refresh_failure",
&self
@@ -1765,12 +1790,12 @@ impl UnauthorizedRecovery {
/// `reload()` is called explicitly. This matches the design goal of avoiding
/// different parts of the program seeing inconsistent auth data mid‑run.
pub struct AuthManager {
codex_home: PathBuf,
pub(super) codex_home: PathBuf,
inner: RwLock<CachedAuth>,
auth_change_tx: watch::Sender<u64>,
enable_codex_api_key_env: bool,
auth_credentials_store_mode: AuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
pub(super) auth_credentials_store_mode: AuthCredentialsStoreMode,
pub(super) keyring_backend_kind: AuthKeyringBackendKind,
forced_chatgpt_workspace_id: RwLock<Option<Vec<String>>>,
chatgpt_base_url: Option<String>,
agent_identity_authapi_base_url: Option<String>,
@@ -1849,6 +1874,13 @@ impl AuthManager {
) -> Self {
let agent_identity_authapi_base_url =
agent_identity_authapi_base_url(chatgpt_base_url.as_deref()).ok();
let bedrock_api_key = load_stored_bedrock_api_key(
&codex_home,
auth_credentials_store_mode,
keyring_backend_kind,
)
.ok()
.flatten();
let managed_auth = load_auth(
&codex_home,
enable_codex_api_key_env,
@@ -1867,6 +1899,7 @@ impl AuthManager {
codex_home,
inner: RwLock::new(CachedAuth {
auth: managed_auth,
bedrock_api_key,
permanent_refresh_failure: None,
}),
auth_change_tx,
@@ -1888,6 +1921,7 @@ impl AuthManager {
pub fn from_auth_for_testing(auth: CodexAuth) -> Arc<Self> {
let cached = CachedAuth {
auth: Some(auth),
bedrock_api_key: None,
permanent_refresh_failure: None,
};
let (auth_change_tx, _auth_change_rx) = watch::channel(0);
@@ -1914,6 +1948,7 @@ impl AuthManager {
pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc<Self> {
let cached = CachedAuth {
auth: Some(auth),
bedrock_api_key: None,
permanent_refresh_failure: None,
};
let (auth_change_tx, _auth_change_rx) = watch::channel(0);
@@ -1943,6 +1978,7 @@ impl AuthManager {
) -> Arc<Self> {
let cached = CachedAuth {
auth: Some(auth),
bedrock_api_key: None,
permanent_refresh_failure: None,
};
let (auth_change_tx, _auth_change_rx) = watch::channel(0);
@@ -1974,6 +2010,7 @@ impl AuthManager {
codex_home: PathBuf::from("non-existent"),
inner: RwLock::new(CachedAuth {
auth: None,
bedrock_api_key: None,
permanent_refresh_failure: None,
}),
auth_change_tx,
@@ -2001,6 +2038,14 @@ impl AuthManager {
.and_then(|cached| cached.auth.clone())
}
/// Current managed Amazon Bedrock API key without applying Codex auth precedence.
pub fn bedrock_api_key_auth_cached(&self) -> Option<BedrockApiKeyAuth> {
self.inner
.read()
.ok()
.and_then(|cached| cached.bedrock_api_key.clone())
}
/// Subscribes to cached auth changes that can affect request recovery.
pub fn auth_change_receiver(&self) -> watch::Receiver<u64> {
self.auth_change_tx.subscribe()
@@ -2103,7 +2148,8 @@ impl AuthManager {
pub async fn reload(&self) -> bool {
tracing::info!("Reloading auth");
let new_auth = self.load_auth().await;
self.set_cached_auth(new_auth)
let new_bedrock_api_key = self.load_bedrock_api_key_from_storage();
self.set_cached_auth_state(new_auth, new_bedrock_api_key)
}
async fn reload_if_account_id_matches(
@@ -2219,6 +2265,60 @@ impl AuthManager {
.flatten()
}
fn load_bedrock_api_key_from_storage(&self) -> Option<BedrockApiKeyAuth> {
load_stored_bedrock_api_key(
&self.codex_home,
self.auth_credentials_store_mode,
self.keyring_backend_kind,
)
.ok()
.flatten()
}
fn set_cached_auth_state(
&self,
new_auth: Option<CodexAuth>,
new_bedrock_api_key: Option<BedrockApiKeyAuth>,
) -> bool {
if let Ok(mut guard) = self.inner.write() {
let previous = guard.auth.as_ref();
let auth_changed = !Self::auths_equal(previous, new_auth.as_ref());
let auth_changed_for_refresh =
!Self::auths_equal_for_refresh(previous, new_auth.as_ref());
let bedrock_auth_changed = guard.bedrock_api_key != new_bedrock_api_key;
if auth_changed_for_refresh {
guard.permanent_refresh_failure = None;
}
let changed = auth_changed || bedrock_auth_changed;
tracing::info!("Reloaded auth, changed: {changed}");
guard.auth = new_auth;
guard.bedrock_api_key = new_bedrock_api_key;
if auth_changed_for_refresh || bedrock_auth_changed {
self.auth_change_tx.send_modify(|revision| *revision += 1);
}
changed
} else {
false
}
}
pub(super) fn set_cached_bedrock_api_key_auth(
&self,
new_bedrock_api_key: Option<BedrockApiKeyAuth>,
) -> bool {
if let Ok(mut guard) = self.inner.write() {
let changed = guard.bedrock_api_key != new_bedrock_api_key;
tracing::info!("Reloaded managed Bedrock auth, changed: {changed}");
guard.bedrock_api_key = new_bedrock_api_key;
if changed {
self.auth_change_tx.send_modify(|revision| *revision += 1);
}
changed
} else {
false
}
}
fn set_cached_auth(&self, new_auth: Option<CodexAuth>) -> bool {
if let Ok(mut guard) = self.inner.write() {
let previous = guard.auth.as_ref();
@@ -39,13 +39,13 @@ use mantle::runtime_base_url;
pub(crate) struct AmazonBedrockModelProvider {
pub(crate) info: ModelProviderInfo,
pub(crate) aws: ModelProviderAwsAuthInfo,
auth_manager: Option<Arc<AuthManager>>,
managed_auth: Option<BedrockApiKeyAuth>,
}
impl AmazonBedrockModelProvider {
pub(crate) fn new(
provider_info: ModelProviderInfo,
auth_manager: Option<Arc<AuthManager>>,
managed_auth: Option<BedrockApiKeyAuth>,
) -> Self {
let aws = provider_info
.aws
@@ -57,23 +57,12 @@ impl AmazonBedrockModelProvider {
Self {
info: provider_info,
aws,
auth_manager,
managed_auth,
}
}
fn managed_auth(&self) -> Option<BedrockApiKeyAuth> {
self.auth_manager
.as_ref()
.and_then(|auth_manager| auth_manager.auth_cached())
.and_then(|auth| match auth {
CodexAuth::BedrockApiKey(auth) => Some(auth),
CodexAuth::ApiKey(_)
| CodexAuth::Chatgpt(_)
| CodexAuth::ChatgptAuthTokens(_)
| CodexAuth::Headers(_)
| CodexAuth::AgentIdentity(_)
| CodexAuth::PersonalAccessToken(_) => None,
})
self.managed_auth.clone()
}
async fn auth(&self) -> Option<CodexAuth> {
@@ -127,8 +116,7 @@ impl ModelProvider for AmazonBedrockModelProvider {
}
fn auth_manager(&self) -> Option<Arc<AuthManager>> {
self.managed_auth()
.and_then(|_| self.auth_manager.as_ref().cloned())
None
}
fn auth(&self) -> ModelProviderFuture<'_, Option<CodexAuth>> {
@@ -208,22 +196,15 @@ mod tests {
api_key: "managed-bedrock-api-key".to_string(),
region: "us-east-1".to_string(),
};
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(managed_auth.clone()));
let provider = AmazonBedrockModelProvider::new(
ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo {
profile: Some("aws-profile-that-should-not-be-loaded".to_string()),
region: Some("us-west-2".to_string()),
})),
Some(auth_manager.clone()),
Some(managed_auth.clone()),
);
assert!(Arc::ptr_eq(
&provider
.auth_manager()
.expect("managed Bedrock auth manager should be exposed"),
&auth_manager,
));
assert!(provider.auth_manager().is_none());
assert_eq!(
provider.auth().await,
Some(CodexAuth::BedrockApiKey(managed_auth))
@@ -256,12 +237,10 @@ mod tests {
}
#[tokio::test]
async fn openai_auth_is_not_exposed_to_bedrock() {
async fn bedrock_without_managed_auth_uses_aws_account() {
let provider = AmazonBedrockModelProvider::new(
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
"openai-api-key",
))),
/*managed_auth*/ None,
);
assert!(provider.auth_manager().is_none());
@@ -281,7 +260,7 @@ mod tests {
fn capabilities_disable_unsupported_hosted_tools() {
let provider = AmazonBedrockModelProvider::new(
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
/*auth_manager*/ None,
/*managed_auth*/ None,
);
assert_eq!(
@@ -298,7 +277,7 @@ mod tests {
fn approval_review_preferred_model_uses_bedrock_gpt_5_4() {
let provider = AmazonBedrockModelProvider::new(
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
/*auth_manager*/ None,
/*managed_auth*/ None,
);
assert_eq!(
+4 -12
View File
@@ -220,7 +220,10 @@ pub fn create_model_provider(
auth_manager: Option<Arc<AuthManager>>,
) -> SharedModelProvider {
if provider_info.is_amazon_bedrock() {
Arc::new(AmazonBedrockModelProvider::new(provider_info, auth_manager))
let managed_auth = auth_manager
.as_ref()
.and_then(|auth_manager| auth_manager.bedrock_api_key_auth_cached());
Arc::new(AmazonBedrockModelProvider::new(provider_info, managed_auth))
} else {
Arc::new(ConfiguredModelProvider::new(provider_info, auth_manager))
}
@@ -529,17 +532,6 @@ mod tests {
assert!(provider.auth_manager().is_none());
}
#[tokio::test]
async fn create_model_provider_uses_managed_auth_for_amazon_bedrock_provider() {
let auth = bedrock_api_key_auth();
let provider = create_model_provider(
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
Some(AuthManager::from_auth_for_testing(auth.clone())),
);
assert_eq!(provider.auth().await, Some(auth));
}
#[test]
fn openai_provider_returns_unauthenticated_openai_account_state() {
let provider = create_model_provider(