mirror of
https://github.com/openai/codex.git
synced 2026-10-30 19:18:49 +00:00
Refresh loaded model provider sessions
This commit is contained in:
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(¤t.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>,
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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) = ¤t.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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user