From ffc79ad48e1bf829972390d669df7cd65f7f69d0 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Fri, 29 May 2026 19:04:41 -0700 Subject: [PATCH] changes --- codex-rs/Cargo.lock | 1 + .../schema/json/ClientRequest.json | 24 ++++ .../codex_app_server_protocol.schemas.json | 40 ++++++ .../codex_app_server_protocol.v2.schemas.json | 40 ++++++ .../schema/json/v2/LoginAccountParams.json | 24 ++++ .../schema/json/v2/LoginAccountResponse.json | 16 +++ .../typescript/v2/LoginAccountParams.ts | 2 +- .../typescript/v2/LoginAccountResponse.ts | 2 +- .../src/protocol/v2/account.rs | 6 + codex-rs/app-server/src/request_processors.rs | 7 + .../request_processors/account_processor.rs | 131 +++++++++++++++++- .../request_processors/config_processor.rs | 6 +- codex-rs/cli/src/doctor.rs | 6 +- codex-rs/core/src/client.rs | 4 +- codex-rs/core/src/client_tests.rs | 2 + codex-rs/core/src/guardian/tests.rs | 19 ++- codex-rs/core/src/mcp_tool_call_tests.rs | 3 + codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 3 + .../core/src/session/tests/guardian_tests.rs | 4 + codex-rs/core/src/session/turn_context.rs | 6 +- codex-rs/core/src/test_support.rs | 2 +- codex-rs/core/src/thread_manager.rs | 14 +- .../src/tools/handlers/multi_agents_tests.rs | 6 +- codex-rs/core/src/tools/spec_plan_tests.rs | 7 +- codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 2 + .../core/tests/suite/client_websockets.rs | 1 + .../ext/image-generation/src/extension.rs | 9 +- codex-rs/ext/web-search/src/extension.rs | 4 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/model-provider/Cargo.toml | 3 +- .../model-provider/src/amazon_bedrock/auth.rs | 46 +++++- .../src/amazon_bedrock/mantle.rs | 24 +++- .../model-provider/src/amazon_bedrock/mod.rs | 34 ++++- .../src/amazon_bedrock/provider_auth.rs | 108 +++++++++++++++ codex-rs/model-provider/src/lib.rs | 5 + codex-rs/model-provider/src/provider.rs | 15 +- codex-rs/tui/src/app.rs | 14 +- 39 files changed, 610 insertions(+), 35 deletions(-) create mode 100644 codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e7bb18b5f6..644c08c166 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3307,6 +3307,7 @@ dependencies = [ "codex-response-debug-context", "http 1.4.0", "pretty_assertions", + "serde", "serde_json", "tokio", "tracing", diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index dc5dd6ad60..8aeae5a1fc 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1331,6 +1331,30 @@ ], "title": "ChatgptAuthTokensLoginAccountParams", "type": "object" + }, + { + "properties": { + "key": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockLoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "key", + "region", + "type" + ], + "title": "AmazonBedrockLoginAccountParams", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 6aa4c5f35e..6fcec76b1d 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -10635,6 +10635,30 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "properties": { + "key": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "key", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" @@ -10728,6 +10752,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 6d859ee5f3..4d9beb83ba 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -7157,6 +7157,30 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "properties": { + "key": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "key", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" @@ -7250,6 +7274,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json index ab7b852c91..29fcc292da 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json @@ -89,6 +89,30 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "properties": { + "key": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "key", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" diff --git a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json index a800bffccd..802440d637 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json @@ -87,6 +87,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts index e6f1e2ed43..f8e7d306a8 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts @@ -18,4 +18,4 @@ chatgptAccountId: string, * When `null`, Codex attempts to derive the plan type from access-token * claims. If unavailable, the plan defaults to `unknown`. */ -chatgptPlanType?: string | null, }; +chatgptPlanType?: string | null, } | { "type": "amazonBedrock", key: string, region: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts index 34bccd6578..5a9f34ead3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts @@ -14,4 +14,4 @@ verificationUrl: string, /** * One-time code the user must enter after signing in. */ -userCode: string, } | { "type": "chatgptAuthTokens", }; +userCode: string, } | { "type": "chatgptAuthTokens", } | { "type": "amazonBedrock", }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/account.rs b/codex-rs/app-server-protocol/src/protocol/v2/account.rs index c2a538982c..e7bf12041a 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -80,6 +80,9 @@ pub enum LoginAccountParams { #[ts(optional = nullable)] chatgpt_plan_type: Option, }, + #[serde(rename = "amazonBedrock", rename_all = "camelCase")] + #[ts(rename = "amazonBedrock", rename_all = "camelCase")] + AmazonBedrock { key: String, region: String }, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -113,6 +116,9 @@ pub enum LoginAccountResponse { #[serde(rename = "chatgptAuthTokens", rename_all = "camelCase")] #[ts(rename = "chatgptAuthTokens", rename_all = "camelCase")] ChatgptAuthTokens {}, + #[serde(rename = "amazonBedrock", rename_all = "camelCase")] + #[ts(rename = "amazonBedrock", rename_all = "camelCase")] + AmazonBedrock {}, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 36008ce0b9..01b1385000 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -47,6 +47,8 @@ use codex_app_server_protocol::CommandExecParams; use codex_app_server_protocol::CommandExecResizeParams; use codex_app_server_protocol::CommandExecTerminateParams; use codex_app_server_protocol::CommandExecWriteParams; +use codex_app_server_protocol::ConfigBatchWriteParams; +use codex_app_server_protocol::ConfigEdit as ConfigWriteEdit; use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::ConversationGitInfo; use codex_app_server_protocol::ConversationSummary; @@ -100,6 +102,7 @@ use codex_app_server_protocol::McpServerStatusDetail; use codex_app_server_protocol::McpServerToolCallParams; use codex_app_server_protocol::McpServerToolCallResponse; use codex_app_server_protocol::MemoryResetResponse; +use codex_app_server_protocol::MergeStrategy; use codex_app_server_protocol::MockExperimentalMethodParams; use codex_app_server_protocol::MockExperimentalMethodResponse; use codex_app_server_protocol::ModelListParams; @@ -354,8 +357,12 @@ use codex_mcp::discover_supported_scopes; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; use codex_memories_write::clear_memory_roots_contents; +use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; use codex_model_provider::ProviderAccountError; use codex_model_provider::create_model_provider; +use codex_model_provider::delete_amazon_bedrock_auth; +use codex_model_provider::is_supported_amazon_bedrock_region; +use codex_model_provider::save_amazon_bedrock_auth; use codex_models_manager::collaboration_mode_presets::builtin_collaboration_mode_presets; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 8ebfc037e1..58a1c8e4ff 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -238,6 +238,9 @@ impl AccountRequestProcessor { ) .await; } + LoginAccountParams::AmazonBedrock { key, region } => { + self.login_amazon_bedrock_v2(request_id, key, region).await; + } } Ok(()) } @@ -300,6 +303,123 @@ impl AccountRequestProcessor { } } + async fn login_amazon_bedrock_common( + &self, + key: &str, + region: &str, + ) -> std::result::Result<(), JSONRPCErrorError> { + let key = key.trim(); + if key.is_empty() { + return Err(invalid_request("Amazon Bedrock key must not be empty.")); + } + + let region = region.trim(); + if !is_supported_amazon_bedrock_region(region) { + return Err(invalid_request(format!( + "Amazon Bedrock Mantle does not support region `{region}`" + ))); + } + + // Cancel any active login attempt. + { + let mut guard = self.active_login.lock().await; + if let Some(active) = guard.take() { + drop(active); + } + } + + save_amazon_bedrock_auth(&self.config.codex_home, key, region) + .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; + + self.config_manager + .batch_write(ConfigBatchWriteParams { + edits: vec![ + ConfigWriteEdit { + key_path: "model_provider".to_string(), + value: serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), + merge_strategy: MergeStrategy::Upsert, + }, + ConfigWriteEdit { + key_path: format!( + "model_providers.{AMAZON_BEDROCK_PROVIDER_ID}.aws.region" + ), + value: serde_json::json!(region), + merge_strategy: MergeStrategy::Upsert, + }, + ], + file_path: None, + expected_version: None, + reload_user_config: false, + }) + .await + .map_err(|err| { + internal_error(format!("failed to update Amazon Bedrock config: {err}")) + })?; + + let next_config = self + .config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + for thread_id in self.thread_manager.list_thread_ids().await { + let Ok(thread) = self.thread_manager.get_thread(thread_id).await else { + continue; + }; + thread.refresh_runtime_config(next_config.clone()).await; + } + Self::spawn_effective_plugins_changed_task( + Arc::clone(&self.thread_manager), + self.config_manager.clone(), + ); + + Ok(()) + } + + async fn login_amazon_bedrock_v2( + &self, + request_id: ConnectionRequestId, + key: String, + region: String, + ) { + let result = self + .login_amazon_bedrock_common(&key, ®ion) + .await + .map(|()| LoginAccountResponse::AmazonBedrock {}); + let logged_in = result.is_ok(); + self.outgoing.send_result(request_id, result).await; + + if logged_in { + self.send_amazon_bedrock_login_success_notifications().await; + } + } + + async fn send_amazon_bedrock_login_success_notifications(&self) { + Self::maybe_refresh_remote_installed_plugins_cache_for_current_config( + &self.config_manager, + &self.thread_manager, + self.auth_manager.auth_cached(), + ) + .await; + + self.outgoing + .send_server_notification(ServerNotification::AccountLoginCompleted( + AccountLoginCompletedNotification { + login_id: None, + success: true, + error: None, + }, + )) + .await; + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated( + AccountUpdatedNotification { + auth_mode: None, + plan_type: None, + }, + )) + .await; + } + // Build options for a ChatGPT login attempt; performs validation. async fn login_chatgpt_common( &self, @@ -682,6 +802,9 @@ impl AccountRequestProcessor { return Err(internal_error(format!("logout failed: {err}"))); } } + delete_amazon_bedrock_auth(&self.config.codex_home).map_err(|err| { + internal_error(format!("failed to remove Amazon Bedrock auth: {err}")) + })?; Self::maybe_refresh_remote_installed_plugins_cache_for_current_config( &self.config_manager, @@ -810,9 +933,15 @@ impl AccountRequestProcessor { self.refresh_token_if_requested(do_refresh).await; + let config = self + .config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; let provider = create_model_provider( - self.config.model_provider.clone(), + config.model_provider.clone(), Some(self.auth_manager.clone()), + Some(config.codex_home.to_path_buf()), ); let account_state = match provider.account_state() { Ok(account_state) => account_state, diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 3571026193..da75cedad9 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -172,7 +172,11 @@ impl ConfigRequestProcessor { &self, ) -> Result { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let provider = create_model_provider(config.model_provider, /*auth_manager*/ None); + let provider = create_model_provider( + config.model_provider, + /*auth_manager*/ None, + Some(config.codex_home.to_path_buf()), + ); let capabilities = provider.capabilities(); Ok(ModelProviderCapabilitiesReadResponse { namespace_tools: capabilities.namespace_tools, diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index 2e53d5d51f..354b2a0a15 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -2260,7 +2260,11 @@ async fn websocket_reachability_check( provider.websocket_connect_timeout().as_millis() )); - let runtime_provider = create_model_provider(provider.clone(), auth_manager); + let runtime_provider = create_model_provider( + provider.clone(), + auth_manager, + Some(config.codex_home.to_path_buf()), + ); let auth = runtime_provider.auth().await; details.push(format!( "auth mode: {}", diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 124dfb18c0..8db66f75b2 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -24,6 +24,7 @@ //! fails, normal stream retry/fallback logic handles recovery on the same turn. use std::collections::HashMap; +use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; @@ -321,6 +322,7 @@ impl ModelClient { thread_id: ThreadId, installation_id: String, provider_info: ModelProviderInfo, + codex_home: PathBuf, session_source: SessionSource, parent_thread_id: Option, model_verbosity: Option, @@ -329,7 +331,7 @@ impl ModelClient { beta_features_header: Option, attestation_provider: Option>, ) -> Self { - let model_provider = create_model_provider(provider_info, auth_manager); + let model_provider = create_model_provider(provider_info, auth_manager, Some(codex_home)); let codex_api_key_env_enabled = model_provider .auth_manager() .as_ref() diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index f8d904036a..4bef86151c 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -76,6 +76,7 @@ fn test_model_client_with_parent( thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), provider, + std::env::temp_dir(), session_source, parent_thread_id, /*model_verbosity*/ None, @@ -530,6 +531,7 @@ fn model_client_with_counting_attestation( ThreadId::new(), /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), provider, + std::env::temp_dir(), SessionSource::Exec, /*parent_thread_id*/ None, /*model_verbosity*/ None, diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index aba67d8500..c4f94000ae 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -188,7 +188,11 @@ async fn guardian_test_session_and_turn_with_base_url( ); session.services.models_manager = models_manager; turn.config = Arc::clone(&config); - turn.provider = create_model_provider(config.model_provider.clone(), turn.auth_manager.clone()); + turn.provider = create_model_provider( + config.model_provider.clone(), + turn.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), + ); turn.user_instructions = None; (Arc::new(session), Arc::new(turn)) @@ -1428,7 +1432,11 @@ async fn guardian_review_request_layout_matches_model_visible_request_snapshot() ); session.services.models_manager = models_manager; turn.config = Arc::clone(&config); - turn.provider = create_model_provider(config.model_provider.clone(), turn.auth_manager.clone()); + turn.provider = create_model_provider( + config.model_provider.clone(), + turn.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), + ); let session = Arc::new(session); let turn = Arc::new(turn); seed_guardian_parent_history(&session, &turn).await; @@ -2007,8 +2015,11 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> .models_manager = models_manager; let turn_mut = Arc::get_mut(&mut turn).expect("turn should be uniquely owned"); turn_mut.config = Arc::clone(&config); - turn_mut.provider = - create_model_provider(config.model_provider.clone(), turn_mut.auth_manager.clone()); + turn_mut.provider = create_model_provider( + config.model_provider.clone(), + turn_mut.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), + ); turn_mut.user_instructions = None; seed_guardian_parent_history(&session, &turn).await; diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index c9259b8f2d..62241a88eb 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -2339,6 +2339,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { turn_context.provider = create_model_provider( config.model_provider.clone(), turn_context.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); @@ -2623,6 +2624,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { turn_context.provider = create_model_provider( config.model_provider.clone(), turn_context.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); @@ -2843,6 +2845,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { turn_context.provider = create_model_provider( config.model_provider.clone(), turn_context.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 05773abb8c..81e3dec144 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1037,6 +1037,7 @@ impl Session { thread_id, installation_id.clone(), session_configuration.provider.clone(), + config.codex_home.to_path_buf(), session_configuration.session_source.clone(), session_configuration.parent_thread_id, config.model_verbosity, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index a3e5558971..234260c467 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -437,6 +437,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession { thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), ModelProviderInfo::create_openai_provider(/* base_url */ /*base_url*/ None), + std::env::temp_dir(), codex_protocol::protocol::SessionSource::Exec, /*parent_thread_id*/ None, /*model_verbosity*/ None, @@ -4770,6 +4771,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), session_configuration.provider.clone(), + config.codex_home.to_path_buf(), session_configuration.session_source.clone(), session_configuration.parent_thread_id, config.model_verbosity, @@ -6860,6 +6862,7 @@ where thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), session_configuration.provider.clone(), + config.codex_home.to_path_buf(), session_configuration.session_source.clone(), session_configuration.parent_thread_id, config.model_verbosity, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d8b46c9999..eeee4f8958 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -114,6 +114,7 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { turn_context_raw.provider = create_model_provider( config.model_provider.clone(), turn_context_raw.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); @@ -205,6 +206,7 @@ async fn request_permissions_guardian_review_stops_when_cancelled() { turn_context_raw.provider = create_model_provider( config.model_provider.clone(), turn_context_raw.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let requested_permissions = RequestPermissionProfile { @@ -320,6 +322,7 @@ async fn guardian_allows_shell_command_additional_permissions_requests_past_poli turn_context_raw.provider = create_model_provider( config.model_provider.clone(), turn_context_raw.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); @@ -425,6 +428,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_ turn_context_raw.provider = create_model_provider( config.model_provider.clone(), turn_context_raw.auth_manager.clone(), + Some(config.codex_home.to_path_buf()), ); let session = Arc::new(session); let turn_context = Arc::new(turn_context_raw); diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index d600368137..0139ea9036 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -489,7 +489,11 @@ impl Session { ); 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 provider_for_context = create_model_provider( + provider, + auth_manager, + Some(per_turn_config.codex_home.to_path_buf()), + ); 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( diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index e491ca7413..33cc82af41 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -121,7 +121,7 @@ pub fn models_manager_with_provider( auth_manager: Arc, provider: ModelProviderInfo, ) -> SharedModelsManager { - let provider = create_model_provider(provider, Some(auth_manager)); + let provider = create_model_provider(provider, Some(auth_manager), Some(codex_home.clone())); provider.models_manager(codex_home, /*config_model_catalog*/ None) } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d130a35815..2ace991e33 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -221,7 +221,11 @@ pub fn build_models_manager( config: &Config, auth_manager: Arc, ) -> SharedModelsManager { - let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); + let provider = create_model_provider( + config.model_provider.clone(), + Some(auth_manager), + Some(config.codex_home.to_path_buf()), + ); provider.models_manager( config.codex_home.to_path_buf(), config.model_catalog.clone(), @@ -380,8 +384,12 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, - models_manager: create_model_provider(provider, Some(auth_manager.clone())) - .models_manager(codex_home, /*config_model_catalog*/ None), + models_manager: create_model_provider( + provider, + Some(auth_manager.clone()), + Some(codex_home.clone()), + ) + .models_manager(codex_home, /*config_model_catalog*/ None), environment_manager, skills_manager, plugins_manager, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index d4ce84ad55..1530e8f8c4 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -264,8 +264,12 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { turn.approval_policy .set(AskForApproval::OnRequest) .expect("approval policy should be set"); - turn.provider = create_model_provider(provider_info, turn.auth_manager.clone()); turn.config = Arc::new(config); + turn.provider = create_model_provider( + provider_info, + turn.auth_manager.clone(), + Some(turn.config.codex_home.to_path_buf()), + ); let invocation = invocation( Arc::new(session), diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 1cb36e474d..000fb96755 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -271,6 +271,7 @@ fn use_chatgpt_auth(turn: &mut TurnContext) { turn.provider = create_model_provider( turn.config.model_provider.clone(), turn.auth_manager.clone(), + Some(turn.config.codex_home.to_path_buf()), ); } @@ -280,7 +281,11 @@ fn use_bedrock_provider(turn: &mut TurnContext) { config.model_provider_id = AMAZON_BEDROCK_PROVIDER_ID.to_string(); config.model_provider = provider_info.clone(); }); - turn.provider = create_model_provider(provider_info, turn.auth_manager.clone()); + turn.provider = create_model_provider( + provider_info, + turn.auth_manager.clone(), + Some(turn.config.codex_home.to_path_buf()), + ); } struct WebRunExtensionTool; diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index eda48a2c0e..14aaa3893a 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -104,6 +104,7 @@ async fn responses_stream_includes_subagent_header_on_review() { thread_id, /*installation_id*/ TEST_INSTALLATION_ID.to_string(), provider.clone(), + config.codex_home.to_path_buf(), session_source, /*parent_thread_id*/ None, config.model_verbosity, @@ -233,6 +234,7 @@ async fn responses_stream_includes_subagent_header_on_other() { thread_id, /*installation_id*/ TEST_INSTALLATION_ID.to_string(), provider.clone(), + config.codex_home.to_path_buf(), session_source, /*parent_thread_id*/ None, config.model_verbosity, @@ -351,6 +353,7 @@ async fn responses_respects_model_info_overrides_from_config() { thread_id, /*installation_id*/ TEST_INSTALLATION_ID.to_string(), provider.clone(), + config.codex_home.to_path_buf(), session_source, /*parent_thread_id*/ None, config.model_verbosity, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c954a954fb..0ea5e97e99 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -903,6 +903,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), provider, + config.codex_home.to_path_buf(), SessionSource::Exec, /*parent_thread_id*/ None, config.model_verbosity, @@ -2360,6 +2361,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), provider.clone(), + config.codex_home.to_path_buf(), SessionSource::Exec, /*parent_thread_id*/ None, config.model_verbosity, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 41db88f1ad..b54af829b5 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2155,6 +2155,7 @@ async fn websocket_harness_with_provider_options( thread_id, /*installation_id*/ TEST_INSTALLATION_ID.to_string(), provider.clone(), + config.codex_home.to_path_buf(), SessionSource::Exec, /*parent_thread_id*/ None, config.model_verbosity, diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index f7a9f2406f..567801fdef 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use codex_core::config::Config; @@ -26,6 +27,7 @@ struct ImageGenerationExtension { struct ImageGenerationExtensionConfig { enabled: bool, provider: ModelProviderInfo, + codex_home: PathBuf, } impl From<&Config> for ImageGenerationExtensionConfig { @@ -35,6 +37,7 @@ impl From<&Config> for ImageGenerationExtensionConfig { enabled: config.features.enabled(Feature::ImageGenExt) && config.model_provider.is_openai(), provider: config.model_provider.clone(), + codex_home: config.codex_home.to_path_buf(), } } } @@ -77,7 +80,11 @@ impl ToolContributor for ImageGenerationExtension { } vec![Arc::new(ImageGenerationTool::new(CodexImagesBackend::new( - create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())), + create_model_provider( + config.provider.clone(), + Some(self.auth_manager.clone()), + Some(config.codex_home.clone()), + ), )))] } } diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index 30fc48961e..0ec1411674 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use codex_api::AllowedCaller; @@ -31,6 +32,7 @@ struct WebSearchExtension { struct WebSearchExtensionConfig { enabled: bool, provider: ModelProviderInfo, + codex_home: PathBuf, settings: SearchSettings, } @@ -42,6 +44,7 @@ impl From<&Config> for WebSearchExtensionConfig { && config.model_provider.is_openai() && web_search_mode != WebSearchMode::Disabled, provider: config.model_provider.clone(), + codex_home: config.codex_home.to_path_buf(), settings: search_settings(config, web_search_mode), } } @@ -120,6 +123,7 @@ impl ToolContributor for WebSearchExtension { provider: create_model_provider( config.provider.clone(), Some(self.auth_manager.clone()), + Some(config.codex_home.clone()), ), settings: config.settings.clone(), })] diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 36099fe2e8..4ef07a18ae 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -179,6 +179,7 @@ impl MemoryStartupContext { self.thread_id, installation_id, config.model_provider.clone(), + config.codex_home.to_path_buf(), session_source, config_snapshot.parent_thread_id, config.model_verbosity, diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index 58235ab24d..3c6d940616 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -26,11 +26,12 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-response-debug-context = { workspace = true } http = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } tokio = { workspace = true, features = ["sync", "time"] } tracing = { workspace = true, features = ["log"] } [dev-dependencies] pretty_assertions = { workspace = true } -serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } wiremock = { workspace = true } diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index ecfd2dd533..11e20a1cbe 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -18,19 +18,29 @@ use crate::BearerAuthProvider; use super::mantle::aws_auth_config; use super::mantle::region_from_config; +use super::provider_auth::AmazonBedrockAuth; const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; const AWS_REGION_ENV_VAR: &str = "AWS_REGION"; const AWS_DEFAULT_REGION_ENV_VAR: &str = "AWS_DEFAULT_REGION"; pub(super) enum BedrockAuthMethod { + StoredBearerToken { token: String, region: String }, EnvBearerToken { token: String, region: String }, AwsSdkAuth { context: AwsAuthContext }, } pub(super) async fn resolve_auth_method( + stored_auth: Option<&AmazonBedrockAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { + if let Some(auth) = stored_auth { + return Ok(BedrockAuthMethod::StoredBearerToken { + token: auth.bearer_token.clone(), + region: auth.region.clone(), + }); + } + if let Some(token) = non_empty_env_var_from(AWS_BEARER_TOKEN_BEDROCK_ENV_VAR, std::env::var) { let region = bearer_token_region(aws, std::env::var)?; return Ok(BedrockAuthMethod::EnvBearerToken { token, region }); @@ -44,10 +54,12 @@ pub(super) async fn resolve_auth_method( } pub(super) async fn resolve_provider_auth( + stored_auth: Option<&AmazonBedrockAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { - match resolve_auth_method(aws).await? { - BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { + match resolve_auth_method(stored_auth, aws).await? { + BedrockAuthMethod::StoredBearerToken { token, .. } + | BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { token: Some(token), account_id: None, is_fedramp_account: false, @@ -194,6 +206,36 @@ mod tests { ); } + #[tokio::test] + async fn stored_bedrock_bearer_auth_takes_precedence() { + let auth = AmazonBedrockAuth { + bearer_token: "stored-bedrock-key".to_string(), + region: "eu-west-1".to_string(), + }; + + let method = resolve_auth_method( + Some(&auth), + &ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-east-1".to_string()), + }, + ) + .await + .expect("stored auth should resolve"); + + match method { + BedrockAuthMethod::StoredBearerToken { token, region } => { + assert_eq!( + (token, region), + ("stored-bedrock-key".to_string(), "eu-west-1".to_string()) + ); + } + BedrockAuthMethod::EnvBearerToken { .. } | BedrockAuthMethod::AwsSdkAuth { .. } => { + panic!("stored auth should take precedence") + } + } + } + #[test] fn bedrock_bearer_auth_uses_aws_region_env() { let region = bearer_token_region( diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 7881845e45..1e22eab30e 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -5,6 +5,7 @@ use codex_protocol::error::Result; use super::auth::BedrockAuthMethod; use super::auth::resolve_auth_method; +use super::provider_auth::AmazonBedrockAuth; const BEDROCK_MANTLE_SERVICE_NAME: &str = "bedrock-mantle"; const BEDROCK_MANTLE_SUPPORTED_REGIONS: [&str; 12] = [ @@ -39,7 +40,7 @@ pub(super) fn region_from_config(aws: &ModelProviderAwsAuthInfo) -> Option Result { - if BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) { + if is_supported_region(region) { Ok(format!("https://bedrock-mantle.{region}.api.aws/openai/v1")) } else { Err(CodexErr::Fatal(format!( @@ -48,14 +49,25 @@ pub(super) fn base_url(region: &str) -> Result { } } -pub(super) async fn runtime_base_url(aws: &ModelProviderAwsAuthInfo) -> Result { - let region = resolve_region(aws).await?; +pub fn is_supported_region(region: &str) -> bool { + BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) +} + +pub(super) async fn runtime_base_url( + stored_auth: Option<&AmazonBedrockAuth>, + aws: &ModelProviderAwsAuthInfo, +) -> Result { + let region = resolve_region(stored_auth, aws).await?; base_url(®ion) } -async fn resolve_region(aws: &ModelProviderAwsAuthInfo) -> Result { - match resolve_auth_method(aws).await? { - BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), +async fn resolve_region( + stored_auth: Option<&AmazonBedrockAuth>, + aws: &ModelProviderAwsAuthInfo, +) -> Result { + match resolve_auth_method(stored_auth, aws).await? { + BedrockAuthMethod::StoredBearerToken { region, .. } + | BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), BedrockAuthMethod::AwsSdkAuth { context } => Ok(context.region().to_string()), } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 8099993973..fa02c122c5 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -1,6 +1,7 @@ mod auth; mod catalog; mod mantle; +mod provider_auth; use std::path::PathBuf; use std::sync::Arc; @@ -25,17 +26,23 @@ use crate::provider::ProviderCapabilities; use auth::resolve_provider_auth; pub(crate) use catalog::static_model_catalog; use catalog::with_default_only_service_tier; +pub use mantle::is_supported_region; use mantle::runtime_base_url; +pub use provider_auth::AmazonBedrockAuth; +pub use provider_auth::delete_amazon_bedrock_auth; +pub use provider_auth::load_amazon_bedrock_auth; +pub use provider_auth::save_amazon_bedrock_auth; /// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint. #[derive(Clone, Debug)] pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, + codex_home: Option, } impl AmazonBedrockModelProvider { - pub(crate) fn new(provider_info: ModelProviderInfo) -> Self { + pub(crate) fn new(provider_info: ModelProviderInfo, codex_home: Option) -> Self { let aws = provider_info .aws .clone() @@ -46,8 +53,20 @@ impl AmazonBedrockModelProvider { Self { info: provider_info, aws, + codex_home, } } + + fn stored_auth(&self) -> Result> { + let Some(codex_home) = self.codex_home.as_deref() else { + return Ok(None); + }; + load_amazon_bedrock_auth(codex_home).map_err(|err| { + codex_protocol::error::CodexErr::Fatal(format!( + "failed to load Amazon Bedrock auth: {err}" + )) + }) + } } #[async_trait::async_trait] @@ -85,16 +104,21 @@ impl ModelProvider for AmazonBedrockModelProvider { async fn api_provider(&self) -> Result { let mut api_provider_info = self.info.clone(); - api_provider_info.base_url = Some(runtime_base_url(&self.aws).await?); + let stored_auth = self.stored_auth()?; + api_provider_info.base_url = Some(runtime_base_url(stored_auth.as_ref(), &self.aws).await?); api_provider_info.to_api_provider(/*auth_mode*/ None) } async fn runtime_base_url(&self) -> Result> { - Ok(Some(runtime_base_url(&self.aws).await?)) + let stored_auth = self.stored_auth()?; + Ok(Some( + runtime_base_url(stored_auth.as_ref(), &self.aws).await?, + )) } async fn api_auth(&self) -> Result { - resolve_provider_auth(&self.aws).await + let stored_auth = self.stored_auth()?; + resolve_provider_auth(stored_auth.as_ref(), &self.aws).await } fn models_manager( @@ -135,6 +159,7 @@ mod tests { fn capabilities_disable_unsupported_hosted_tools() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + /*codex_home*/ None, ); assert_eq!( @@ -151,6 +176,7 @@ mod tests { fn approval_review_preferred_model_uses_bedrock_gpt_5_4() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + /*codex_home*/ None, ); assert_eq!( diff --git a/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs b/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs new file mode 100644 index 0000000000..c72199e8cf --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/provider_auth.rs @@ -0,0 +1,108 @@ +use std::fs; +use std::io; +use std::path::Path; +use std::path::PathBuf; + +use serde::Deserialize; +use serde::Serialize; + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AmazonBedrockAuth { + pub bearer_token: String, + pub region: String, +} + +pub fn save_amazon_bedrock_auth( + codex_home: &Path, + bearer_token: &str, + region: &str, +) -> io::Result<()> { + let auth = AmazonBedrockAuth { + bearer_token: bearer_token.to_string(), + region: region.to_string(), + }; + let auth_path = amazon_bedrock_auth_file(codex_home); + if let Some(parent) = auth_path.parent() { + fs::create_dir_all(parent)?; + } + let contents = serde_json::to_vec_pretty(&auth)?; + fs::write(&auth_path, contents)?; + restrict_file_permissions(&auth_path) +} + +pub fn load_amazon_bedrock_auth(codex_home: &Path) -> io::Result> { + let auth_path = amazon_bedrock_auth_file(codex_home); + let contents = match fs::read_to_string(&auth_path) { + Ok(contents) => contents, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(err) => return Err(err), + }; + let auth = serde_json::from_str(&contents)?; + Ok(Some(auth)) +} + +pub fn delete_amazon_bedrock_auth(codex_home: &Path) -> io::Result { + let auth_path = amazon_bedrock_auth_file(codex_home); + match fs::remove_file(auth_path) { + Ok(()) => Ok(true), + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(false), + Err(err) => Err(err), + } +} + +fn amazon_bedrock_auth_file(codex_home: &Path) -> PathBuf { + codex_home + .join("model-providers") + .join("amazon-bedrock") + .join("auth.json") +} + +#[cfg(unix)] +fn restrict_file_permissions(auth_path: &Path) -> io::Result<()> { + use std::os::unix::fs::PermissionsExt; + + fs::set_permissions(auth_path, fs::Permissions::from_mode(0o600)) +} + +#[cfg(not(unix))] +fn restrict_file_permissions(_auth_path: &Path) -> io::Result<()> { + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::time::SystemTime; + use std::time::UNIX_EPOCH; + + use super::*; + + #[test] + fn save_load_and_delete_amazon_bedrock_auth() { + let codex_home = std::env::temp_dir().join(format!( + "codex-bedrock-auth-test-{}-{}", + std::process::id(), + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("time") + .as_nanos() + )); + + save_amazon_bedrock_auth(&codex_home, "bedrock-key", "us-east-1").expect("save auth"); + + assert_eq!( + load_amazon_bedrock_auth(&codex_home).expect("load auth"), + Some(AmazonBedrockAuth { + bearer_token: "bedrock-key".to_string(), + region: "us-east-1".to_string(), + }) + ); + assert!(delete_amazon_bedrock_auth(&codex_home).expect("delete auth")); + assert_eq!( + load_amazon_bedrock_auth(&codex_home).expect("load missing auth"), + None + ); + assert!(!delete_amazon_bedrock_auth(&codex_home).expect("delete missing auth")); + let _ = fs::remove_dir_all(&codex_home); + } +} diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 4e4660812b..1cc3540b1c 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,10 +4,15 @@ mod bearer_auth_provider; mod models_endpoint; mod provider; +pub use amazon_bedrock::delete_amazon_bedrock_auth; +pub use amazon_bedrock::is_supported_region as is_supported_amazon_bedrock_region; +pub use amazon_bedrock::load_amazon_bedrock_auth; +pub use amazon_bedrock::save_amazon_bedrock_auth; pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; +pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider; pub use provider::ProviderAccountError; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index e2ef3f3650..167990c52b 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -148,9 +148,10 @@ pub type SharedModelProvider = Arc; pub fn create_model_provider( provider_info: ModelProviderInfo, auth_manager: Option>, + codex_home: Option, ) -> SharedModelProvider { if provider_info.is_amazon_bedrock() { - Arc::new(AmazonBedrockModelProvider::new(provider_info)) + Arc::new(AmazonBedrockModelProvider::new(provider_info, codex_home)) } else { Arc::new(ConfiguredModelProvider::new(provider_info, auth_manager)) } @@ -356,6 +357,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_openai_provider(/*base_url*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!(provider.capabilities(), ProviderCapabilities::default()); @@ -366,6 +368,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_openai_provider(/*base_url*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!( @@ -379,6 +382,7 @@ mod tests { let provider = create_model_provider( provider_for("https://example.test/v1".to_string()), /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!( @@ -395,6 +399,7 @@ mod tests { let provider = create_model_provider( provider_info_with_command_auth(), /*auth_manager*/ None, + /*codex_home*/ None, ); let auth_manager = provider @@ -414,6 +419,7 @@ mod tests { Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "openai-api-key", ))), + /*codex_home*/ None, ); assert!(provider.auth_manager().is_none()); @@ -424,6 +430,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_openai_provider(/*base_url*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!( @@ -442,6 +449,7 @@ mod tests { Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "openai-api-key", ))), + /*codex_home*/ None, ); assert_eq!( @@ -464,6 +472,7 @@ mod tests { ..Default::default() }, /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!( @@ -480,6 +489,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); assert_eq!( @@ -496,6 +506,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); let manager = provider.models_manager(test_codex_home(), /*config_model_catalog*/ None); @@ -533,6 +544,7 @@ mod tests { let provider = create_model_provider( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), /*auth_manager*/ None, + /*codex_home*/ None, ); let manager = provider.models_manager( test_codex_home(), @@ -579,6 +591,7 @@ mod tests { Some(AuthManager::from_auth_for_testing( CodexAuth::create_dummy_chatgpt_auth_for_testing(), )), + /*codex_home*/ None, ); let manager = diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index b1245db3ec..c392fa00f8 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -589,8 +589,15 @@ fn active_turn_not_steerable_turn_error(error: &TypedRequestError) -> Option Option { - let provider = create_model_provider(provider.clone(), /*auth_manager*/ None); +async fn resolve_runtime_model_provider_base_url( + provider: &ModelProviderInfo, + codex_home: &Path, +) -> Option { + let provider = create_model_provider( + provider.clone(), + /*auth_manager*/ None, + Some(codex_home.to_path_buf()), + ); match provider.runtime_base_url().await { Ok(base_url) => base_url, Err(err) => { @@ -839,7 +846,8 @@ impl App { ); let runtime_model_provider_started_at = Instant::now(); let runtime_model_provider_base_url = - resolve_runtime_model_provider_base_url(&config.model_provider).await; + resolve_runtime_model_provider_base_url(&config.model_provider, &config.codex_home) + .await; let runtime_model_provider_ms = runtime_model_provider_started_at.elapsed().as_millis(); let enhanced_keys_supported = tui.enhanced_keys_supported();