Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions codex-rs/app-server/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2060,6 +2060,8 @@ Codex stores the key and region as the primary Codex auth, replacing any previou
{ "method": "account/updated", "params": { "authMode": null, "planType": null } }
```

When Codex-managed Amazon Bedrock auth is stored, logout removes that credential and clears the user-level `model_provider` only when it is still `"amazon-bedrock"`. A concurrent user config change is preserved. For AWS-managed Bedrock auth, logout is a no-op because the AWS SDK credential chain is managed outside Codex.

### 7) Rate limits (ChatGPT)

```json
Expand Down
1 change: 0 additions & 1 deletion codex-rs/app-server/src/request_processors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ use codex_app_server_protocol::AppTemplateUnavailableReason;
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::CancelLoginAccountParams;
use codex_app_server_protocol::CancelLoginAccountResponse;
use codex_app_server_protocol::CancelLoginAccountStatus;
Expand Down
43 changes: 26 additions & 17 deletions codex-rs/app-server/src/request_processors/account_processor.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use super::bedrock_auth::clear_user_model_provider_if_bedrock;
use super::bedrock_auth::set_user_model_provider_to_bedrock;
use super::bedrock_auth::user_model_provider_state;
use super::*;
use crate::auth_mode::auth_mode_to_api;
use chrono::DateTime;
Expand Down Expand Up @@ -811,7 +813,7 @@ impl AccountRequestProcessor {
}
}

async fn logout_common(&self) -> std::result::Result<Option<AuthMode>, JSONRPCErrorError> {
async fn logout_common(&self) -> Result<(), JSONRPCErrorError> {
// Cancel any active login attempt.
{
let mut guard = self.active_login.lock().await;
Expand All @@ -820,6 +822,19 @@ impl AccountRequestProcessor {
}
}

let managed_bedrock_auth = matches!(
self.auth_manager.auth_cached(),
Some(CodexAuth::BedrockApiKey(_))
);
let user_model_provider = if managed_bedrock_auth {
Some(user_model_provider_state(&self.config_manager).await?)
} else {
None
};
if self.config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth {
return Ok(());
}

match self.auth_manager.logout_with_revoke().await {
Ok(_) => {}
Err(err) => {
Expand All @@ -834,26 +849,20 @@ impl AccountRequestProcessor {
)
.await;

// Reflect the current auth method after logout (likely None).
Ok(self
.auth_manager
.auth_cached()
.as_ref()
.map(CodexAuth::api_auth_mode)
.map(auth_mode_to_api))
if let Some(user_model_provider) = user_model_provider {
clear_user_model_provider_if_bedrock(&self.config_manager, user_model_provider).await?;
}

Ok(())
}

async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> {
let result = self.logout_common().await;
let account_updated =
result
.as_ref()
.ok()
.cloned()
.map(|auth_mode| AccountUpdatedNotification {
auth_mode,
plan_type: None,
});
let account_updated = if result.is_ok() {
Some(self.current_account_updated_notification())
} else {
None
};
self.outgoing
.send_result(request_id, result.map(|_| LogoutAccountResponse {}))
.await;
Expand Down
64 changes: 64 additions & 0 deletions codex-rs/app-server/src/request_processors/bedrock_auth.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,18 @@
use super::config_processor::map_error as map_config_error;
use crate::config_manager::ConfigManager;
use crate::error_code::internal_error;
use codex_app_server_protocol::ConfigLayerSource;
use codex_app_server_protocol::ConfigReadParams;
use codex_app_server_protocol::ConfigValueWriteParams;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::MergeStrategy;
use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID;

pub(super) struct UserModelProviderState {
model_provider: Option<String>,
version: Option<String>,
}

pub(super) async fn set_user_model_provider_to_bedrock(
config_manager: &ConfigManager,
) -> Result<(), JSONRPCErrorError> {
Expand All @@ -16,6 +24,62 @@ pub(super) async fn set_user_model_provider_to_bedrock(
.await
}

pub(super) async fn clear_user_model_provider_if_bedrock(
config_manager: &ConfigManager,
user_model_provider: UserModelProviderState,
) -> Result<(), JSONRPCErrorError> {
if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) {
write_user_model_provider(
config_manager,
serde_json::Value::Null,
user_model_provider.version,
)
.await?;
}

Ok(())
}

pub(super) async fn user_model_provider_state(
config_manager: &ConfigManager,
) -> Result<UserModelProviderState, JSONRPCErrorError> {
let user_config_path = config_manager
.user_config_path()
.map_err(|err| internal_error(format!("failed to resolve user config path: {err}")))?;
let response = config_manager
.read(ConfigReadParams {
include_layers: true,
cwd: None,
})
.await
.map_err(|err| internal_error(format!("failed to read user config: {err}")))?;
let layer = response
.layers
.unwrap_or_default()
.into_iter()
.find(|layer| {
matches!(
&layer.name,
ConfigLayerSource::User { file, .. } if file == &user_config_path
)
});
let Some(layer) = layer else {
return Ok(UserModelProviderState {
model_provider: None,
version: None,
});
};
let model_provider = layer
.config
.get("model_provider")
.and_then(serde_json::Value::as_str)
.map(str::to_string);
Ok(UserModelProviderState {
model_provider,
version: Some(layer.version),
})
}

async fn write_user_model_provider(
config_manager: &ConfigManager,
value: serde_json::Value,
Expand Down
135 changes: 135 additions & 0 deletions codex-rs/app-server/tests/suite/v2/account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1108,6 +1108,141 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R
Ok(())
}

#[tokio::test]
async fn logout_managed_bedrock_restores_aws_managed_account() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), aws_managed_bedrock_config())?;
let mut expected_config = read_config_toml(codex_home.path())?;
expected_config
.as_table_mut()
.expect("config should be a table")
.remove("model_provider");

let mut mcp =
TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).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_eq!(load_file_auth(codex_home.path())?, None);
assert_eq!(read_config_toml(codex_home.path())?, expected_config);
assert_account_updated(&mut mcp, /*auth_mode*/ None).await?;
assert_eq!(
read_account(&mut mcp).await?,
GetAccountResponse {
account: Some(Account::AmazonBedrock {
credential_source: AmazonBedrockCredentialSource::AwsManaged,
}),
requires_openai_auth: false,
}
);
Ok(())
}

#[tokio::test]
async fn logout_aws_managed_bedrock_preserves_openai_auth_and_config() -> Result<()> {
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), aws_managed_bedrock_config())?;
login_with_api_key(
codex_home.path(),
"sk-test-key",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
let expected_auth = load_file_auth(codex_home.path())?;
let expected_config = read_config_toml(codex_home.path())?;

let mut mcp =
TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).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_eq!(load_file_auth(codex_home.path())?, expected_auth);
assert_eq!(read_config_toml(codex_home.path())?, expected_config);
assert_account_updated(&mut mcp, Some(AuthMode::ApiKey)).await?;
Ok(())
}

#[tokio::test]
async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_api() -> Result<()>
{
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
login_with_bedrock_api_key(
codex_home.path(),
"managed-bedrock-api-key",
"us-west-2",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
let expected_config = read_config_toml(codex_home.path())?;

let mut mcp = TestAppServer::new(codex_home.path()).await?;
let initialized = mcp
.initialize_with_capabilities(
ClientInfo {
name: DEFAULT_CLIENT_NAME.to_string(),
title: None,
version: "0.1.0".to_string(),
},
Some(InitializeCapabilities {
experimental_api: false,
..Default::default()
}),
)
.await?;
assert!(matches!(initialized, JSONRPCMessage::Response(_)));

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_eq!(load_file_auth(codex_home.path())?, None);
assert_eq!(read_config_toml(codex_home.path())?, expected_config);
assert_account_updated(&mut mcp, /*auth_mode*/ None).await?;
Ok(())
}

#[tokio::test]
async fn managed_bedrock_login_requires_experimental_api() -> Result<()> {
let codex_home = TempDir::new()?;
Expand Down
Loading