From 9df4d86a68ed1bdf3b9e21c91a049e0905e5aaa8 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 22 Apr 2026 15:26:34 -0700 Subject: [PATCH] Expose AWS account state in app-server --- codex-rs/Cargo.lock | 1 + .../codex_app_server_protocol.schemas.json | 16 ++ .../codex_app_server_protocol.v2.schemas.json | 16 ++ .../schema/json/v2/GetAccountResponse.json | 16 ++ .../schema/typescript/v2/Account.ts | 2 +- .../app-server-protocol/src/protocol/v2.rs | 15 ++ codex-rs/app-server/Cargo.toml | 1 + .../app-server/src/codex_message_processor.rs | 62 +++----- codex-rs/app-server/tests/suite/v2/account.rs | 70 ++++++++- .../model-provider/src/amazon_bedrock/mod.rs | 10 ++ codex-rs/model-provider/src/lib.rs | 4 + codex-rs/model-provider/src/provider.rs | 140 ++++++++++++++++++ codex-rs/protocol/src/account.rs | 8 + codex-rs/tui/src/app_server_session.rs | 3 + codex-rs/tui/src/lib.rs | 1 + 15 files changed, 314 insertions(+), 51 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c3b607eb66b3..fcdca99252ea 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1856,6 +1856,7 @@ dependencies = [ "codex-git-utils", "codex-login", "codex-mcp", + "codex-model-provider", "codex-model-provider-info", "codex-models-manager", "codex-otel", 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 cf54aca2829b..d47cc10c0ea4 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 @@ -5149,6 +5149,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 506d92e604c9..3be027b874fb 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 f8bc91c01eec..88940424c681 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; @@ -1889,6 +1890,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/Cargo.toml b/codex-rs/app-server/Cargo.toml index 339bc20f10f0..b38b1e28120c 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -49,6 +49,7 @@ codex-file-search = { workspace = true } codex-chatgpt = { workspace = true } codex-login = { workspace = true } codex-mcp = { workspace = true } +codex-model-provider = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } codex-app-server-protocol = { workspace = true } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0eb465e54af4..f4ac19849206 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; @@ -296,6 +295,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; @@ -1832,51 +1833,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 f12c6a914a92..1b1a615ccd40 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -5,6 +5,10 @@ mod 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 bb46329a51d9..48ff0115cdd1 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 26746b8b19cc..4e2c4736739d 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -266,6 +266,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 8814cddda30f..ed78cfd20d17 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, }) }