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 6904ef7165ec..ccce74bc0ca5 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 @@ -5220,6 +5220,22 @@ ], "title": "ChatgptAccount", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockAccountType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockAccount", + "type": "object" } ] }, 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 bc00828e68d4..8560fb6c843c 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 @@ -46,6 +46,22 @@ ], "title": "ChatgptAccount", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockAccountType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockAccount", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json index 8534927157a1..ec333708b76b 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json @@ -42,6 +42,22 @@ ], "title": "ChatgptAccount", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockAccountType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockAccount", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts b/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts index f91677499e74..4c3a58e8d6a3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/Account.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { PlanType } from "../PlanType"; -export type Account = { "type": "apiKey", } | { "type": "chatgpt", email: string, planType: PlanType, }; +export type Account = { "type": "apiKey", } | { "type": "chatgpt", email: string, planType: PlanType, } | { "type": "amazonBedrock", }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index b6726679fb1a..97165a509261 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -7,6 +7,7 @@ use crate::RequestId; use crate::protocol::common::AuthMode; use codex_experimental_api_macros::ExperimentalApi; use codex_protocol::account::PlanType; +use codex_protocol::account::ProviderAccount; use codex_protocol::approvals::ElicitationRequest as CoreElicitationRequest; use codex_protocol::approvals::ExecPolicyAmendment as CoreExecPolicyAmendment; use codex_protocol::approvals::GuardianAssessmentAction as CoreGuardianAssessmentAction; @@ -2015,6 +2016,20 @@ pub enum Account { #[serde(rename = "chatgpt", rename_all = "camelCase")] #[ts(rename = "chatgpt", rename_all = "camelCase")] Chatgpt { email: String, plan_type: PlanType }, + + #[serde(rename = "amazonBedrock", rename_all = "camelCase")] + #[ts(rename = "amazonBedrock", rename_all = "camelCase")] + AmazonBedrock {}, +} + +impl From for Account { + fn from(account: ProviderAccount) -> Self { + match account { + ProviderAccount::ApiKey => Self::ApiKey {}, + ProviderAccount::Chatgpt { email, plan_type } => Self::Chatgpt { email, plan_type }, + ProviderAccount::AmazonBedrock => Self::AmazonBedrock {}, + } + } } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, ExperimentalApi)] diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c94568947653..9d671996e883 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -37,7 +37,6 @@ use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::AuthMode; -use codex_app_server_protocol::AuthMode as CoreAuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; @@ -302,6 +301,8 @@ use codex_mcp::discover_supported_scopes; use codex_mcp::effective_mcp_servers; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; +use codex_model_provider::ProviderAccountError; +use codex_model_provider::create_model_provider; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; @@ -1844,51 +1845,28 @@ impl CodexMessageProcessor { self.refresh_token_if_requested(do_refresh).await; - // Whether auth is required for the active model provider. - let requires_openai_auth = self.config.model_provider.requires_openai_auth; - - if !requires_openai_auth { - let response = GetAccountResponse { - account: None, - requires_openai_auth, - }; - self.outgoing.send_response(request_id, response).await; - return; - } - - let account = match self.auth_manager.auth_cached() { - Some(auth) => match auth.auth_mode() { - CoreAuthMode::ApiKey => Some(Account::ApiKey {}), - CoreAuthMode::Chatgpt - | CoreAuthMode::ChatgptAuthTokens - | CoreAuthMode::AgentIdentity => { - let email = auth.get_account_email(); - let plan_type = auth.account_plan_type(); - - match (email, plan_type) { - (Some(email), Some(plan_type)) => { - Some(Account::Chatgpt { email, plan_type }) - } - _ => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: - "email and plan type are required for chatgpt authentication" - .to_string(), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - } - } - }, - None => None, + let provider = create_model_provider( + self.config.model_provider.clone(), + Some(self.auth_manager.clone()), + ); + let account_state = match provider.account_state() { + Ok(account_state) => account_state, + Err(ProviderAccountError::MissingChatgptAccountDetails) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: "email and plan type are required for chatgpt authentication" + .to_string(), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } }; + let account = account_state.account.map(Account::from); let response = GetAccountResponse { account, - requires_openai_auth, + requires_openai_auth: account_state.requires_openai_auth, }; self.outgoing.send_response(request_id, response).await; } diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 3c88bcb7a430..2d75fd10a271 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -55,6 +55,8 @@ struct CreateConfigTomlParams { forced_workspace_id: Option, requires_openai_auth: Option, base_url: Option, + model_provider_id: Option, + extra_provider_config: Option, } fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std::io::Result<()> { @@ -77,6 +79,23 @@ fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std: Some(false) => String::new(), None => String::new(), }; + let model_provider_id = params + .model_provider_id + .unwrap_or_else(|| "mock_provider".to_string()); + let provider_section = if model_provider_id == "mock_provider" { + format!( + r#"[model_providers.mock_provider] +name = "Mock provider for test" +base_url = "{base_url}" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +{requires_line} +"# + ) + } else { + params.extra_provider_config.unwrap_or_default() + }; let contents = format!( r#" model = "mock-model" @@ -85,18 +104,12 @@ sandbox_mode = "danger-full-access" {forced_line} {forced_workspace_line} -model_provider = "mock_provider" +model_provider = "{model_provider_id}" [features] shell_snapshot = false -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{base_url}" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -{requires_line} +{provider_section} "# ); std::fs::write(config_toml, contents) @@ -1545,6 +1558,47 @@ async fn get_account_when_auth_not_required() -> Result<()> { Ok(()) } +#[tokio::test] +async fn get_account_with_aws_provider() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + model_provider_id: Some("amazon-bedrock".to_string()), + extra_provider_config: Some( + r#"[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +region = "us-west-2" +"# + .to_string(), + ), + ..Default::default() + }, + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let params = GetAccountParams { + refresh_token: false, + }; + let request_id = mcp.send_get_account_request(params).await?; + + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let received: GetAccountResponse = to_response(resp)?; + + let expected = GetAccountResponse { + account: Some(Account::AmazonBedrock {}), + requires_openai_auth: false, + }; + assert_eq!(received, expected); + Ok(()) +} + #[tokio::test] async fn get_account_with_chatgpt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index a28262fb7d17..af7ac8714ce1 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -9,9 +9,12 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::account::ProviderAccount; use codex_protocol::error::Result; use crate::provider::ModelProvider; +use crate::provider::ProviderAccountResult; +use crate::provider::ProviderAccountState; use auth::resolve_provider_auth; use auth::resolve_region; use mantle::base_url; @@ -37,6 +40,13 @@ impl ModelProvider for AmazonBedrockModelProvider { None } + fn account_state(&self) -> ProviderAccountResult { + Ok(ProviderAccountState { + account: Some(ProviderAccount::AmazonBedrock), + requires_openai_auth: false, + }) + } + async fn api_provider(&self) -> Result { let region = resolve_region(&self.aws).await?; let mut api_provider_info = self.info.clone(); diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index f5454edd3f4b..11c180db1141 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -7,6 +7,10 @@ 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_protocol::account::ProviderAccount; pub use provider::ModelProvider; +pub use provider::ProviderAccountError; +pub use provider::ProviderAccountResult; +pub use provider::ProviderAccountState; pub use provider::SharedModelProvider; pub use provider::create_model_provider; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 3075c2a318a8..7cd14bbc49b4 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -7,11 +7,42 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::account::ProviderAccount; use crate::amazon_bedrock::AmazonBedrockModelProvider; use crate::auth::auth_manager_for_provider; use crate::auth::resolve_provider_auth; +/// Current app-visible account state for a model provider. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderAccountState { + pub account: Option, + pub requires_openai_auth: bool, +} + +/// Error returned when a provider cannot construct its app-visible account state. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProviderAccountError { + MissingChatgptAccountDetails, +} + +impl fmt::Display for ProviderAccountError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::MissingChatgptAccountDetails => { + write!( + f, + "email and plan type are required for chatgpt authentication" + ) + } + } + } +} + +impl std::error::Error for ProviderAccountError {} + +pub type ProviderAccountResult = std::result::Result; + /// Runtime provider abstraction used by model execution. /// /// Implementations own provider-specific behavior for a model backend. The @@ -33,6 +64,9 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { /// Returns the current provider-scoped auth value, if one is configured. async fn auth(&self) -> Option; + /// Returns the current app-visible account state for this provider. + fn account_state(&self) -> ProviderAccountResult; + /// Returns provider configuration adapted for the API client. async fn api_provider(&self) -> codex_protocol::error::Result { let auth = self.auth().await; @@ -99,6 +133,38 @@ impl ModelProvider for ConfiguredModelProvider { None => None, } } + + fn account_state(&self) -> ProviderAccountResult { + let account = if self.info.requires_openai_auth { + self.auth_manager + .as_ref() + .and_then(|auth_manager| auth_manager.auth_cached()) + .map(|auth| match &auth { + CodexAuth::ApiKey(_) => Ok(ProviderAccount::ApiKey), + CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::AgentIdentity(_) => { + let email = auth.get_account_email(); + let plan_type = auth.account_plan_type(); + + match (email, plan_type) { + (Some(email), Some(plan_type)) => { + Ok(ProviderAccount::Chatgpt { email, plan_type }) + } + _ => Err(ProviderAccountError::MissingChatgptAccountDetails), + } + } + }) + .transpose()? + } else { + None + }; + + Ok(ProviderAccountState { + account, + requires_openai_auth: self.info.requires_openai_auth, + }) + } } #[cfg(test)] @@ -106,7 +172,9 @@ mod tests { use std::num::NonZeroU64; use codex_model_provider_info::ModelProviderAwsAuthInfo; + use codex_model_provider_info::WireApi; use codex_protocol::config_types::ModelProviderAuthInfo; + use pretty_assertions::assert_eq; use super::*; @@ -155,4 +223,76 @@ mod tests { assert!(provider.auth_manager().is_none()); } + + #[test] + fn openai_provider_returns_unauthenticated_openai_account_state() { + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + /*auth_manager*/ None, + ); + + assert_eq!( + provider.account_state(), + Ok(ProviderAccountState { + account: None, + requires_openai_auth: true, + }) + ); + } + + #[test] + fn openai_provider_returns_api_key_account_state() { + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "openai-api-key", + ))), + ); + + assert_eq!( + provider.account_state(), + Ok(ProviderAccountState { + account: Some(ProviderAccount::ApiKey), + requires_openai_auth: true, + }) + ); + } + + #[test] + fn custom_non_openai_provider_returns_no_account_state() { + let provider = create_model_provider( + ModelProviderInfo { + name: "Custom".to_string(), + base_url: Some("http://localhost:1234/v1".to_string()), + wire_api: WireApi::Responses, + requires_openai_auth: false, + ..Default::default() + }, + /*auth_manager*/ None, + ); + + assert_eq!( + provider.account_state(), + Ok(ProviderAccountState { + account: None, + requires_openai_auth: false, + }) + ); + } + + #[test] + fn amazon_bedrock_provider_returns_bedrock_account_state() { + let provider = create_model_provider( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + /*auth_manager*/ None, + ); + + assert_eq!( + provider.account_state(), + Ok(ProviderAccountState { + account: Some(ProviderAccount::AmazonBedrock), + requires_openai_auth: false, + }) + ); + } } diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index 5832381cbc8c..aea9ad843ac7 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -27,6 +27,14 @@ pub enum PlanType { Unknown, } +/// Account state returned by a model provider before it is adapted to an app-facing wire type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProviderAccount { + ApiKey, + Chatgpt { email: String, plan_type: PlanType }, + AmazonBedrock, +} + impl PlanType { pub fn is_team_like(self) -> bool { matches!(self, Self::Team | Self::SelfServeBusinessUsageBased) diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 8992d97d6798..262b540c6215 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -273,6 +273,9 @@ impl AppServerSession { true, ) } + Some(Account::AmazonBedrock {}) => { + (None, None, None, None, FeedbackAudience::External, false) + } None => (None, None, None, None, FeedbackAudience::External, false), }; Ok(AppServerBootstrap { diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index ee50d0a178f8..095a2f347726 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -1669,6 +1669,7 @@ async fn get_login_status( Ok(match account.account { Some(AppServerAccount::ApiKey {}) => LoginStatus::AuthMode(AppServerAuthMode::ApiKey), Some(AppServerAccount::Chatgpt { .. }) => LoginStatus::AuthMode(AppServerAuthMode::Chatgpt), + Some(AppServerAccount::AmazonBedrock {}) => LoginStatus::NotAuthenticated, None => LoginStatus::NotAuthenticated, }) }