diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 8a39fdc299..f8a41e274e 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1852,6 +1852,41 @@ ], "title": "AmazonBedrockLoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock AWS access key login is experimental.", + "properties": { + "accessKeyId": { + "type": "string" + }, + "region": { + "type": "string" + }, + "secretAccessKey": { + "type": "string" + }, + "sessionToken": { + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "amazonBedrockAccessKeys" + ], + "title": "AmazonBedrockAccessKeysLoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "accessKeyId", + "region", + "secretAccessKey", + "type" + ], + "title": "AmazonBedrockAccessKeysLoginAccountParams", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 4f73d298a5..6772a70680 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -573,6 +573,13 @@ "bedrockApiKey" ], "type": "string" + }, + { + "description": "Amazon Bedrock AWS access keys managed by Codex.", + "enum": [ + "bedrockAccessKeys" + ], + "type": "string" } ] }, 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 a2f100c793..706251c4e4 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 @@ -7510,6 +7510,13 @@ "bedrockApiKey" ], "type": "string" + }, + { + "description": "Amazon Bedrock AWS access keys managed by Codex.", + "enum": [ + "bedrockAccessKeys" + ], + "type": "string" } ] }, @@ -13357,6 +13364,41 @@ ], "title": "AmazonBedrockv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock AWS access key login is experimental.", + "properties": { + "accessKeyId": { + "type": "string" + }, + "region": { + "type": "string" + }, + "secretAccessKey": { + "type": "string" + }, + "sessionToken": { + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "amazonBedrockAccessKeys" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "accessKeyId", + "region", + "secretAccessKey", + "type" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" 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 9dc3448153..86b5b67c32 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 @@ -1235,6 +1235,13 @@ "bedrockApiKey" ], "type": "string" + }, + { + "description": "Amazon Bedrock AWS access keys managed by Codex.", + "enum": [ + "bedrockAccessKeys" + ], + "type": "string" } ] }, @@ -9567,6 +9574,41 @@ ], "title": "AmazonBedrockv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock AWS access key login is experimental.", + "properties": { + "accessKeyId": { + "type": "string" + }, + "region": { + "type": "string" + }, + "secretAccessKey": { + "type": "string" + }, + "sessionToken": { + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "amazonBedrockAccessKeys" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "accessKeyId", + "region", + "secretAccessKey", + "type" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" diff --git a/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json index 4e46f81cdd..c4d6bc9e09 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json @@ -52,6 +52,13 @@ "bedrockApiKey" ], "type": "string" + }, + { + "description": "Amazon Bedrock AWS access keys managed by Codex.", + "enum": [ + "bedrockAccessKeys" + ], + "type": "string" } ] }, 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 0c89965ab6..80164cc562 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json @@ -137,6 +137,41 @@ ], "title": "AmazonBedrockv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock AWS access key login is experimental.", + "properties": { + "accessKeyId": { + "type": "string" + }, + "region": { + "type": "string" + }, + "secretAccessKey": { + "type": "string" + }, + "sessionToken": { + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "amazonBedrockAccessKeys" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "accessKeyId", + "region", + "secretAccessKey", + "type" + ], + "title": "AmazonBedrockAccessKeysv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst index 0d009b2036..4f686dd0da 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 1f03c2a00c..76ac470aa0 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts b/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts index b248876173..b346c72362 100644 --- a/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts +++ b/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts @@ -5,4 +5,4 @@ /** * Authentication mode for OpenAI-backed providers. */ -export type AuthMode = "apikey" | "chatgpt" | "chatgptAuthTokens" | "headers" | "agentIdentity" | "personalAccessToken" | "bedrockApiKey"; +export type AuthMode = "apikey" | "chatgpt" | "chatgptAuthTokens" | "headers" | "agentIdentity" | "personalAccessToken" | "bedrockApiKey" | "bedrockAccessKeys"; 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 41d6075ef0..c316cac614 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts @@ -19,4 +19,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, } | { "type": "amazonBedrock", apiKey: string, region: string, }; +chatgptPlanType?: string | null, } | { "type": "amazonBedrock", apiKey: string, region: string, } | { "type": "amazonBedrockAccessKeys", accessKeyId: string, secretAccessKey: string, sessionToken?: string | null, region: string, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 0204c13b5b..ee62837e7c 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -54,6 +54,11 @@ pub enum AuthMode { #[ts(rename = "bedrockApiKey")] #[strum(serialize = "bedrockApiKey")] BedrockApiKey, + /// Amazon Bedrock AWS access keys managed by Codex. + #[serde(rename = "bedrockAccessKeys")] + #[ts(rename = "bedrockAccessKeys")] + #[strum(serialize = "bedrockAccessKeys")] + BedrockAccessKeys, } impl AuthMode { @@ -61,7 +66,11 @@ impl AuthMode { pub fn has_chatgpt_account(self) -> bool { match self { Self::Chatgpt | Self::ChatgptAuthTokens | Self::PersonalAccessToken => true, - Self::ApiKey | Self::Headers | Self::AgentIdentity | Self::BedrockApiKey => false, + Self::ApiKey + | Self::Headers + | Self::AgentIdentity + | Self::BedrockApiKey + | Self::BedrockAccessKeys => false, } } @@ -73,7 +82,7 @@ impl AuthMode { | Self::Headers | Self::AgentIdentity | Self::PersonalAccessToken => true, - Self::ApiKey | Self::BedrockApiKey => false, + Self::ApiKey | Self::BedrockApiKey | Self::BedrockAccessKeys => false, } } } @@ -3234,25 +3243,47 @@ mod tests { #[test] fn serialize_account_login_amazon_bedrock() -> Result<()> { - let request = ClientRequest::LoginAccount { - request_id: RequestId::Integer(2), - params: v2::LoginAccountParams::AmazonBedrock { - api_key: "secret".to_string(), - region: "us-west-2".to_string(), - }, - }; - assert_eq!( - json!({ - "method": "account/login/start", - "id": 2, - "params": { + for (params, expected_params) in [ + ( + v2::LoginAccountParams::AmazonBedrock { + api_key: "secret".to_string(), + region: "us-west-2".to_string(), + }, + json!({ "type": "amazonBedrock", "apiKey": "secret", "region": "us-west-2" - } - }), - serde_json::to_value(&request)?, - ); + }), + ), + ( + v2::LoginAccountParams::AmazonBedrockAccessKeys { + access_key_id: "access-key-id".to_string(), + secret_access_key: "secret-access-key".to_string(), + session_token: Some("session-token".to_string()), + region: "us-west-2".to_string(), + }, + json!({ + "type": "amazonBedrockAccessKeys", + "accessKeyId": "access-key-id", + "secretAccessKey": "secret-access-key", + "sessionToken": "session-token", + "region": "us-west-2" + }), + ), + ] { + let request = ClientRequest::LoginAccount { + request_id: RequestId::Integer(2), + params, + }; + assert_eq!( + json!({ + "method": "account/login/start", + "id": 2, + "params": expected_params, + }), + serde_json::to_value(&request)?, + ); + } assert_eq!( json!({"type": "amazonBedrock"}), serde_json::to_value(v2::LoginAccountResponse::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 bbc7959674..017f022efa 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -106,6 +106,17 @@ pub enum LoginAccountParams { #[serde(rename = "amazonBedrock", rename_all = "camelCase")] #[ts(rename = "amazonBedrock", rename_all = "camelCase")] AmazonBedrock { api_key: String, region: String }, + /// [UNSTABLE] Managed Amazon Bedrock AWS access key login is experimental. + #[experimental("account/login/start.amazonBedrockAccessKeys")] + #[serde(rename = "amazonBedrockAccessKeys", rename_all = "camelCase")] + #[ts(rename = "amazonBedrockAccessKeys", rename_all = "camelCase")] + AmazonBedrockAccessKeys { + access_key_id: String, + secret_access_key: String, + #[ts(optional = nullable)] + session_token: Option, + region: String, + }, } #[derive(Serialize, Deserialize, Debug, Default, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/bedrock.rs b/codex-rs/app-server-protocol/src/protocol/v2/bedrock.rs index 5a9ab2caa0..effd6f906a 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/bedrock.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/bedrock.rs @@ -54,15 +54,6 @@ pub enum BedrockSetupParams { #[serde(rename_all = "camelCase")] #[ts(rename_all = "camelCase")] Environment { region: String }, - #[serde(rename_all = "camelCase")] - #[ts(rename_all = "camelCase")] - AccessKeys { - access_key_id: String, - secret_access_key: String, - #[ts(optional = nullable)] - session_token: Option, - region: String, - }, } #[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 68297254fa..0114a0bad5 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -123,6 +123,7 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index e8e30e9139..6803b99b04 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -2061,6 +2061,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 76dc294f3f..188106f550 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -2271,17 +2271,19 @@ Codex supports these authentication modes. The current mode is surfaced in `acco - **API key (`apiKey`)**: Caller supplies an OpenAI API key via `account/login/start` with `type: "apiKey"`. The API key is saved and used for API requests. - **ChatGPT managed (`chatgpt`)** (recommended): Codex owns the ChatGPT OAuth flow and refresh tokens. Start via `account/login/start` with `type: "chatgpt"` for the browser flow or `type: "chatgptDeviceCode"` for device code; Codex persists tokens to disk and refreshes them automatically. -- **Codex managed Amazon Bedrock auth (`amazonBedrock`, experimental)**: Caller supplies an Amazon Bedrock API key and region via `account/login/start` with `type: "amazonBedrock"`. The client must enable the `experimentalApi` initialization capability for Codex-managed Amazon Bedrock login. Codex replaces the current primary auth with the Bedrock credential and writes `model_provider = "amazon-bedrock"` to the user config. +- **Codex managed Amazon Bedrock auth (experimental)**: Caller supplies an Amazon Bedrock API key using `type: "amazonBedrock"` or AWS access keys using `type: "amazonBedrockAccessKeys"` via `account/login/start`. The client must enable the `experimentalApi` initialization capability. Codex replaces the current primary auth with the Bedrock credential and writes `model_provider = "amazon-bedrock"` to the user config. - **Personal access token (`personalAccessToken`)**: Codex uses a ChatGPT-backed personal access token loaded outside the app-server login RPCs, such as with `codex login --with-access-token` or `CODEX_ACCESS_TOKEN`. ### API Overview - `account/read` — fetch current account info; optionally refresh tokens. -- `account/login/start` — begin login (`apiKey`, `chatgpt`, `chatgptDeviceCode`, `amazonBedrock`). +- `account/login/start` — begin login (`apiKey`, `chatgpt`, `chatgptDeviceCode`, `amazonBedrock`, `amazonBedrockAccessKeys`). +- `account/bedrock/discover` — experimental; list available AWS profiles and identify AWS access keys or Amazon Bedrock API keys visible in the app-server environment. +- `account/bedrock/setup` — experimental; validate a selected AWS profile or existing environment credentials, then persist the Amazon Bedrock provider configuration. - `account/login/completed` (notify) — emitted when a login attempt finishes (success or error). - `account/login/cancel` — cancel a pending managed ChatGPT login by `loginId`. - `account/logout` — sign out; triggers `account/updated` on success. -- `account/updated` (notify) — emitted whenever auth mode changes (`authMode`: `apikey`, `bedrockApiKey`, `chatgpt`, `personalAccessToken`, or `null`) and includes the current ChatGPT `planType` when available. +- `account/updated` (notify) — emitted whenever auth mode changes (`authMode`: `apikey`, `bedrockApiKey`, `bedrockAccessKeys`, `chatgpt`, `personalAccessToken`, or `null`) and includes the current ChatGPT `planType` when available. - `account/rateLimits/read` — fetch ChatGPT rate limits, an optional effective monthly credit limit, whether spend control has been reached, and the earned rate-limit resets currently available, including expiry details when provided by the backend. Rate-limit updates arrive via `account/rateLimits/updated` (notify); reset-credit data is snapshot-only. - `account/rateLimitResetCredit/consume` — consume one earned reset using a caller-provided idempotency key, optionally selecting a reset-credit ID returned by `account/rateLimits/read`. - `account/usage/read` — fetch ChatGPT account token-activity summary and daily buckets, or pass a valid thread UUID as `threadId` to read estimated credits, optional cost, and usage breakdowns for one thread using the app-server's active account. The optional `threadUsage` response field is absent on older servers and `null` when the billing route is unavailable. @@ -2313,7 +2315,7 @@ Field notes: - `refreshToken` (bool): set `true` to force a token refresh. - `email` is `null` when the ChatGPT account does not have an email address. - `requiresOpenaiAuth` reflects the active provider; when `false`, Codex can run without OpenAI credentials. -- Amazon Bedrock reports `usesCodexManagedCredentials: true` when it uses a Bedrock API key managed by Codex. It reports `false` for external credential paths, including the AWS credential chain and configured command auth. This identifies whether Codex-managed credentials are selected; it does not validate that the credential source can resolve credentials. +- Amazon Bedrock reports `usesCodexManagedCredentials: true` when it uses a Bedrock API key or AWS access keys managed by Codex. It reports `false` for external credential paths, including the AWS credential chain and configured command auth. This identifies whether Codex-managed credentials are selected; it does not validate that the credential source can resolve credentials. ### 2) Log in with an API key @@ -2356,7 +2358,7 @@ Field notes: `onboardingEntrypoint` is optional and is only emitted when the OAuth callback carries a recognized onboarding hint. -### 3) Log in with an Amazon Bedrock API key +### 3) Log in with Amazon Bedrock credentials This experimental flow requires the client to initialize with `experimentalApi: true`. @@ -2378,7 +2380,77 @@ This experimental flow requires the client to initialize with `experimentalApi: { "method": "account/updated", "params": { "authMode": "bedrockApiKey", "planType": null } } ``` -Codex stores the key and region as the primary Codex auth, replacing any previously stored login, and writes `model_provider = "amazon-bedrock"` to the active user config. Existing loaded sessions keep their current provider selection, so clients should restart the app-server before sending more model requests. This limitation will be addressed in a follow-up. +To log in with AWS access keys instead: + +```json +{ + "method": "account/login/start", + "id": 30, + "params": { + "type": "amazonBedrockAccessKeys", + "accessKeyId": "...", + "secretAccessKey": "...", + "sessionToken": "...", + "region": "us-west-2" + } +} +{ "id": 30, "result": { "type": "amazonBedrock" } } +{ "method": "account/login/completed", "params": { "loginId": null, "success": true, "error": null } } +{ "method": "account/updated", "params": { "authMode": "bedrockAccessKeys", "planType": null } } +``` + +The session token is optional. Both flows store credentials in the configured auth backend +(`auth.json` or keyring), replace any previously stored login, and select +`model_provider = "amazon-bedrock"`; access-key login also writes the selected AWS region to the +active user config. Neither flow changes `$CODEX_HOME/.env`. Existing loaded sessions keep their +current provider selection, so clients should restart the app-server before sending more model +requests. This limitation will be addressed in a follow-up. + +### Discover and configure AWS-managed Amazon Bedrock credentials + +These experimental methods require the client to initialize with `experimentalApi: true`. + +Discover AWS profiles and credentials already visible to the app-server process: + +```json +{ "method": "account/bedrock/discover", "id": 31, "params": {} } +{ + "id": 31, + "result": { + "profiles": [{ "name": "engineering", "region": "us-west-2" }], + "environmentCredentials": [ + { "type": "accessKeys", "region": "us-west-2" }, + { "type": "bedrockApiKey", "region": "us-west-2" } + ] + } +} +``` + +Discovery returns credential metadata only; it never includes access keys, secret access keys, +session tokens, or Bedrock API keys. A profile or environment credential's `region` is `null` +when no profile region or explicit `AWS_REGION` is available from that source. + +Set up a named AWS profile: + +```json +{ + "method": "account/bedrock/setup", + "id": 32, + "params": { "type": "profile", "profile": "engineering", "region": "us-west-2" } +} +{ "id": 32, "result": {} } +``` + +To select credentials already visible in the environment, use +`{ "type": "environment", "region": "us-west-2" }`. The provider +resolves available environment credentials through its normal authentication chain. Selecting +profile or environment credentials leaves existing keys in `$CODEX_HOME/.env` unchanged. + +Successful setup writes `model_provider = "amazon-bedrock"` and the selected AWS region to the +active user config, and additionally writes the selected profile for profile-based setup. Clients +should restart the app-server before sending more model requests. Logging out while an Amazon +Bedrock provider is selected clears the user-configured provider, profile, and region, removes +any Codex-managed credentials, and leaves AWS-managed credentials and `$CODEX_HOME/.env` unchanged. ### 4) Log in with ChatGPT (device code flow) @@ -2409,7 +2481,12 @@ Codex stores the key and region as the primary Codex auth, replacing any previou { "method": "account/updated", "params": { "authMode": null, "planType": null } } ``` -When using a Codex-managed Bedrock key, logout removes the key and clears `model_provider` if it is still set to `"amazon-bedrock"`. When using AWS-managed credentials, manage them through AWS or switch providers before logging out. +When `model_provider` is `"amazon-bedrock"` or `"amazon-bedrock-runtime"`, logout clears that +provider selection and its configured AWS profile and region, regardless of whether the +credentials are Codex-managed or AWS-managed. If the selected model is Bedrock-specific, logout +also clears `model`; `model_reasoning_effort` and other generic settings are preserved. +Codex-managed credentials are removed; AWS profiles, environment credentials, and +`$CODEX_HOME/.env` are left untouched. ### 7) Rate limits (ChatGPT) diff --git a/codex-rs/app-server/src/auth_mode.rs b/codex-rs/app-server/src/auth_mode.rs index d543470781..52831070f7 100644 --- a/codex-rs/app-server/src/auth_mode.rs +++ b/codex-rs/app-server/src/auth_mode.rs @@ -16,5 +16,6 @@ pub(crate) fn auth_mode_to_api(auth_mode: AuthMode) -> ApiAuthMode { AuthMode::AgentIdentity => ApiAuthMode::AgentIdentity, AuthMode::PersonalAccessToken => ApiAuthMode::PersonalAccessToken, AuthMode::BedrockApiKey => ApiAuthMode::BedrockApiKey, + AuthMode::BedrockAccessKeys => ApiAuthMode::BedrockAccessKeys, } } diff --git a/codex-rs/app-server/src/config_manager_service.rs b/codex-rs/app-server/src/config_manager_service.rs index 05a5f0b7fa..b7016229e8 100644 --- a/codex-rs/app-server/src/config_manager_service.rs +++ b/codex-rs/app-server/src/config_manager_service.rs @@ -199,43 +199,6 @@ impl ConfigManager { .await } - /// Clears a value from the active user config only when its current raw value matches. - pub(crate) async fn clear_user_value_if_matches( - &self, - key_path: &str, - expected_value: JsonValue, - ) -> Result<(), ConfigManagerError> { - let layers = self - .load_thread_agnostic_config() - .await - .map_err(|err| ConfigManagerError::io("failed to load configuration", err))?; - let Some(user_layer) = layers.get_active_user_layer() else { - return Ok(()); - }; - let segments = parse_key_path(key_path).map_err(|message| { - ConfigManagerError::write(ConfigWriteErrorCode::ConfigValidationError, message) - })?; - let expected_value = parse_value(expected_value).map_err(|message| { - ConfigManagerError::write(ConfigWriteErrorCode::ConfigValidationError, message) - })?; - if value_at_path(&user_layer.config, &segments) != expected_value.as_ref() { - return Ok(()); - } - let expected_version = Some(user_layer.version.clone()); - - self.apply_edits( - /*file_path*/ None, - expected_version, - vec![( - key_path.to_string(), - JsonValue::Null, - MergeStrategy::Replace, - )], - ) - .await?; - Ok(()) - } - pub(crate) async fn batch_write( &self, params: ConfigBatchWriteParams, diff --git a/codex-rs/app-server/src/config_manager_service_tests.rs b/codex-rs/app-server/src/config_manager_service_tests.rs index 580ecd417b..4020d04a8f 100644 --- a/codex-rs/app-server/src/config_manager_service_tests.rs +++ b/codex-rs/app-server/src/config_manager_service_tests.rs @@ -197,40 +197,6 @@ async fn clearing_user_setting_falls_back_to_packaged_default_without_override() Ok(()) } -#[tokio::test] -async fn clear_user_value_if_matches_clears_matching_value() -> Result<()> { - let tmp = tempdir().expect("tempdir"); - let path = tmp.path().join(CONFIG_TOML_FILE); - std::fs::write(&path, "model = \"gpt-5.2\"\napproval_policy = \"never\"\n")?; - - let service = ConfigManager::without_managed_config_for_tests(tmp.path().to_path_buf()); - service - .clear_user_value_if_matches("model", serde_json::json!("gpt-5.2")) - .await?; - - assert_eq!( - std::fs::read_to_string(&path)?, - "approval_policy = \"never\"\n" - ); - Ok(()) -} - -#[tokio::test] -async fn clear_user_value_if_matches_preserves_non_matching_value() -> Result<()> { - let tmp = tempdir().expect("tempdir"); - let path = tmp.path().join(CONFIG_TOML_FILE); - let original = "model = \"gpt-5.2\"\napproval_policy = \"never\"\n"; - std::fs::write(&path, original)?; - - let service = ConfigManager::without_managed_config_for_tests(tmp.path().to_path_buf()); - service - .clear_user_value_if_matches("model", serde_json::json!("gpt-5.3")) - .await?; - - assert_eq!(std::fs::read_to_string(&path)?, original); - Ok(()) -} - #[tokio::test] async fn write_value_rejects_legacy_profile_selector() -> Result<()> { let tmp = tempdir().expect("tempdir"); 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 084aa72652..d966fd6084 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,11 +1,14 @@ +use super::bedrock_auth::BedrockProviderConfig; use super::bedrock_auth::clear_user_model_provider_if_bedrock; -use super::bedrock_auth::configure_bedrock_for_managed_auth; +use super::bedrock_auth::configure_bedrock_provider; +use super::bedrock_auth::ensure_user_model_provider_can_be_bedrock; use super::*; use crate::auth_mode::auth_mode_to_api; use crate::external_auth::ExternalAuthBridge; use chrono::DateTime; use codex_app_server_protocol::DesktopOnboardingEntrypoint; use codex_login::LoginOnboardingEntrypoint; +use codex_login::login_with_bedrock_access_keys; use codex_model_provider::is_supported_amazon_bedrock_region; mod bedrock_setup; @@ -64,6 +67,15 @@ enum RefreshTokenRequestOutcome { FailedPermanently, } +enum BedrockLoginCredentials { + ApiKey(String), + AccessKeys { + access_key_id: String, + secret_access_key: String, + session_token: Option, + }, +} + impl Drop for ActiveLogin { fn drop(&mut self) { self.cancel(); @@ -326,8 +338,29 @@ impl AccountRequestProcessor { .await; } LoginAccountParams::AmazonBedrock { api_key, region } => { - self.login_amazon_bedrock_v2(request_id, api_key, region) - .await; + self.login_amazon_bedrock_v2( + request_id, + BedrockLoginCredentials::ApiKey(api_key), + region, + ) + .await; + } + LoginAccountParams::AmazonBedrockAccessKeys { + access_key_id, + secret_access_key, + session_token, + region, + } => { + self.login_amazon_bedrock_v2( + request_id, + BedrockLoginCredentials::AccessKeys { + access_key_id, + secret_access_key, + session_token, + }, + region, + ) + .await; } } Ok(()) @@ -420,38 +453,76 @@ impl AccountRequestProcessor { async fn login_amazon_bedrock_v2( &self, request_id: ConnectionRequestId, - api_key: String, + credentials: BedrockLoginCredentials, region: String, ) { let result = async { self.ensure_bedrock_login_allowed()?; - let api_key = api_key.trim(); - if api_key.is_empty() { - return Err(invalid_request("Amazon Bedrock API key must not be empty.")); + match &credentials { + BedrockLoginCredentials::ApiKey(api_key) => { + if api_key.trim().is_empty() { + return Err(invalid_request("Amazon Bedrock API key must not be empty.")); + } + } + BedrockLoginCredentials::AccessKeys { + access_key_id, + secret_access_key, + .. + } => { + if access_key_id.trim().is_empty() || secret_access_key.trim().is_empty() { + return Err(invalid_request( + "AWS access key ID and secret access 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}`" + "Amazon Bedrock does not support region `{region}`" ))); } - { - let mut guard = self.active_login.lock().await; - if let Some(active) = guard.take() { - drop(active); + self.cancel_active_login().await; + ensure_user_model_provider_can_be_bedrock(&self.config_manager).await?; + configure_bedrock_provider( + &self.config_manager, + BedrockProviderConfig { + region: matches!(&credentials, BedrockLoginCredentials::AccessKeys { .. }) + .then_some(region), + profile: None, + }, + ) + .await?; + + match credentials { + BedrockLoginCredentials::ApiKey(api_key) => login_with_bedrock_api_key( + &self.config.codex_home, + api_key.trim(), + region, + self.config.cli_auth_credentials_store_mode, + self.config.auth_keyring_backend_kind(), + ), + BedrockLoginCredentials::AccessKeys { + access_key_id, + secret_access_key, + session_token, + } => { + let session_token = session_token + .as_deref() + .map(str::trim) + .filter(|token| !token.is_empty()); + login_with_bedrock_access_keys( + &self.config.codex_home, + access_key_id.trim(), + secret_access_key.trim(), + session_token, + self.config.cli_auth_credentials_store_mode, + self.config.auth_keyring_backend_kind(), + ) } } - - configure_bedrock_for_managed_auth(&self.config_manager).await?; - login_with_bedrock_api_key( - &self.config.codex_home, - api_key, - region, - self.config.cli_auth_credentials_store_mode, - self.config.auth_keyring_backend_kind(), - ) .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; self.auth_manager.reload().await; self.config_manager.clear_cloud_config_bundle_loader(); @@ -883,16 +954,7 @@ impl AccountRequestProcessor { if self.auth_manager.is_workload_identity_selected() { return Err(self.configured_auth_owned_by_host_error()); } - let managed_bedrock_auth = matches!( - self.auth_manager.auth_cached(), - Some(CodexAuth::BedrockApiKey(_)) - ); let config = self.load_latest_config().await; - if config.model_provider.is_amazon_bedrock() && !managed_bedrock_auth { - return Err(invalid_request( - "cannot log out while Amazon Bedrock is using AWS-managed credentials; manage those credentials through AWS or switch model providers before logging out Codex authentication", - )); - } // Cancel any active login attempt. { @@ -909,12 +971,12 @@ impl AccountRequestProcessor { } } - self.config_manager.clear_cloud_config_bundle_loader(); - - if managed_bedrock_auth { - clear_user_model_provider_if_bedrock(&self.config_manager).await?; + if config.model_provider.is_amazon_bedrock() { + clear_user_model_provider_if_bedrock(&self.config_manager, &config).await?; } + self.config_manager.clear_cloud_config_bundle_loader(); + Self::maybe_refresh_plugin_caches_for_current_config( &self.config_manager, &self.thread_manager, diff --git a/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs b/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs index c979a41767..9e7cb31856 100644 --- a/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs +++ b/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs @@ -1,5 +1,6 @@ +use super::super::bedrock_auth::BedrockProviderConfig; +use super::super::bedrock_auth::configure_bedrock_provider; use super::super::bedrock_auth::ensure_user_model_provider_can_be_bedrock; -use super::super::config_processor::map_error as map_config_error; use super::AccountRequestProcessor; use crate::error_code::internal_error; use crate::error_code::invalid_request; @@ -11,12 +12,8 @@ use codex_app_server_protocol::BedrockEnvironmentCredential; use codex_app_server_protocol::BedrockSetupParams; use codex_app_server_protocol::BedrockSetupResponse; use codex_app_server_protocol::ClientResponsePayload; -use codex_app_server_protocol::ConfigBatchWriteParams; -use codex_app_server_protocol::ConfigEdit; use codex_app_server_protocol::JSONRPCErrorError; -use codex_app_server_protocol::MergeStrategy; use codex_login::CodexAuth; -use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; use codex_model_provider::is_supported_amazon_bedrock_region; const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; @@ -77,22 +74,21 @@ impl AccountRequestProcessor { if matches!(¶ms, BedrockSetupParams::Environment { .. }) && matches!( self.auth_manager.auth_cached(), - Some(CodexAuth::BedrockApiKey(_)) + Some(CodexAuth::BedrockApiKey(_) | CodexAuth::BedrockAccessKeys(_)) ) { return Err(invalid_request( - "A Bedrock API key is already configured and takes priority over AWS credentials. Run `codex logout` and try again.", + "Codex-managed Bedrock credentials are already configured and take priority over AWS environment credentials. Run `codex logout` and try again.", )); } let region = match ¶ms { BedrockSetupParams::Profile { region, .. } - | BedrockSetupParams::Environment { region, .. } - | BedrockSetupParams::AccessKeys { region, .. } => region.trim(), + | BedrockSetupParams::Environment { region, .. } => region.trim(), }; if !is_supported_amazon_bedrock_region(region) { return Err(invalid_request(format!( - "Amazon Bedrock Mantle does not support region `{region}`" + "Amazon Bedrock does not support region `{region}`" ))); } @@ -122,45 +118,18 @@ impl AccountRequestProcessor { } None } - BedrockSetupParams::AccessKeys { .. } => { - // TODO: Support direct AWS access key setup with Codex-managed credential storage. - return Err(invalid_request( - "Direct AWS access key setup is not supported yet.", - )); - } }; // TODO: Validate the selected credentials against Bedrock ListModels once supported. self.cancel_active_login().await; - self.config_manager - .batch_write(ConfigBatchWriteParams { - edits: [ - ( - "model_provider", - serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), - ), - ( - "model_providers.amazon-bedrock.aws.region", - serde_json::json!(region), - ), - ( - "model_providers.amazon-bedrock.aws.profile", - serde_json::json!(profile), - ), - ] - .into_iter() - .map(|(key_path, value)| ConfigEdit { - key_path: key_path.to_string(), - value, - merge_strategy: MergeStrategy::Replace, - }) - .collect(), - file_path: None, - expected_version: None, - reload_user_config: false, - }) - .await - .map_err(map_config_error)?; + configure_bedrock_provider( + &self.config_manager, + BedrockProviderConfig { + region: Some(region), + profile: profile.as_deref(), + }, + ) + .await?; Ok(Some(BedrockSetupResponse {}.into())) } diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index 429d95fa49..016b90bcf8 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -10,39 +10,59 @@ use codex_app_server_protocol::MergeStrategy; use codex_config::CONFIG_TOML_FILE; use codex_config::ConfigLayerSource; use codex_config::format_config_layer_source; +use codex_core::config::Config; +use codex_core::config::edit::ConfigEditsBuilder; use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; -pub(super) async fn configure_bedrock_for_managed_auth( - config_manager: &ConfigManager, -) -> Result<(), JSONRPCErrorError> { - ensure_user_model_provider_can_be_bedrock(config_manager).await?; +pub(super) struct BedrockProviderConfig<'a> { + pub(super) region: Option<&'a str>, + pub(super) profile: Option<&'a str>, +} - config_manager +pub(super) async fn configure_bedrock_provider( + config_manager: &ConfigManager, + config: BedrockProviderConfig<'_>, +) -> Result<(), JSONRPCErrorError> { + let mut edits = vec![ + ( + "model_provider", + serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), + ), + ( + "model_providers.amazon-bedrock.aws.profile", + serde_json::json!(config.profile), + ), + ]; + if let Some(region) = config.region { + edits.push(( + "model_providers.amazon-bedrock.aws.region", + serde_json::json!(region), + )); + } + + let response = config_manager .batch_write(ConfigBatchWriteParams { - edits: [ - ( - "model_provider", - serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), - ), - ( - "model_providers.amazon-bedrock.aws.profile", - serde_json::Value::Null, - ), - ] - .into_iter() - .map(|(key_path, value)| ConfigEdit { - key_path: key_path.to_string(), - value, - merge_strategy: MergeStrategy::Replace, - }) - .collect(), + edits: edits + .into_iter() + .map(|(key_path, value)| ConfigEdit { + key_path: key_path.to_string(), + value, + merge_strategy: MergeStrategy::Replace, + }) + .collect(), file_path: None, expected_version: None, reload_user_config: false, }) .await - .map(|_| ()) - .map_err(map_config_error) + .map_err(map_config_error)?; + if let Some(overridden) = response.overridden_metadata { + let message = overridden.message; + return Err(invalid_request(format!( + "Amazon Bedrock configuration cannot take effect: {message}" + ))); + } + Ok(()) } pub(super) async fn ensure_user_model_provider_can_be_bedrock( @@ -84,26 +104,54 @@ pub(super) async fn ensure_user_model_provider_can_be_bedrock( pub(super) async fn clear_user_model_provider_if_bedrock( config_manager: &ConfigManager, + config: &Config, ) -> Result<(), JSONRPCErrorError> { - let result = config_manager - .clear_user_value_if_matches( - "model_provider", - serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), - ) - .await; - if let Err(err) = &result - && err.write_error_code() == Some(ConfigWriteErrorCode::ConfigVersionConflict) - { - tracing::warn!( - "configuration changed while clearing the managed Amazon Bedrock model provider; retrying once" - ); - return config_manager - .clear_user_value_if_matches( - "model_provider", - serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), - ) - .await - .map_err(map_config_error); + let mut latest_config = None; + loop { + let config = latest_config.as_ref().unwrap_or(config); + let Some(key_paths) = ConfigEditsBuilder::bedrock_provider_config_paths_to_clear(config) + else { + return Ok(()); + }; + let expected_version = config + .config_layer_stack + .get_active_user_layer() + .map(|layer| layer.version.clone()); + + let result = config_manager + .batch_write(ConfigBatchWriteParams { + edits: key_paths + .into_iter() + .map(|segments| ConfigEdit { + key_path: segments.join("."), + value: serde_json::Value::Null, + merge_strategy: MergeStrategy::Replace, + }) + .collect(), + file_path: None, + expected_version, + reload_user_config: false, + }) + .await; + if let Err(err) = &result + && latest_config.is_none() + && err.write_error_code() == Some(ConfigWriteErrorCode::ConfigVersionConflict) + { + tracing::warn!( + "configuration changed while clearing Amazon Bedrock settings; retrying once" + ); + latest_config = Some( + config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|err| { + internal_error(format!( + "failed to reload Amazon Bedrock configuration: {err}" + )) + })?, + ); + continue; + } + return result.map(|_| ()).map_err(map_config_error); } - result.map_err(map_config_error) } diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index d68a49c1a1..2303f9be18 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -167,6 +167,7 @@ pub fn write_chatgpt_auth( agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index d3a31bced9..54ed7f3100 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -45,6 +45,7 @@ use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::CLIENT_ID_OVERRIDE_ENV_VAR; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; +use codex_login::auth::BedrockAccessKeysAuth; use codex_login::auth::BedrockApiKeyAuth; use codex_login::load_auth_dot_json; use codex_login::login_with_api_key; @@ -1095,8 +1096,13 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { Ok(()) } +#[test_case("amazonBedrock"; "api_key")] +#[test_case("amazonBedrockAccessKeys"; "access_keys")] #[tokio::test] -async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> Result<()> { +async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider( + credential_type: &str, +) -> Result<()> { + let managed_access_keys = credential_type == "amazonBedrockAccessKeys"; let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; let config_path = codex_home.path().join("config.toml"); @@ -1108,7 +1114,8 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R http_headers = {{ X-Existing = \"preserved\" }}\n\ [model_providers.amazon-bedrock.aws]\n\ profile = \"stale-profile\"\n\ - region = \"us-east-1\"\n" + region = \"us-east-1\"\n\ + auth_refresh = {{ command = \"aws\" }}\n" ), )?; login_with_api_key( @@ -1135,9 +1142,26 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R .as_table_mut() .expect("AWS configuration should be a table") .remove("profile"); - let request_id = mcp - .send_login_account_amazon_bedrock_request(" managed-bedrock-api-key ", " us-west-2 ") - .await?; + if managed_access_keys { + expected_config["model_providers"]["amazon-bedrock"]["aws"]["region"] = + toml::Value::String("us-west-2".to_string()); + } + let params = if managed_access_keys { + json!({ + "type": credential_type, + "accessKeyId": " test-id ", + "secretAccessKey": " test-secret ", + "sessionToken": " test-token ", + "region": " us-west-2 ", + }) + } else { + json!({ + "type": credential_type, + "apiKey": " managed-bedrock-api-key ", + "region": " us-west-2 ", + }) + }; + let request_id = mcp.send_login_account_request(params).await?; let response: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_response_message(RequestId::Integer(request_id)), @@ -1151,19 +1175,29 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R assert_eq!( load_file_auth(codex_home.path())?, Some(AuthDotJson { - auth_mode: Some(DomainAuthMode::BedrockApiKey), + auth_mode: Some(if managed_access_keys { + DomainAuthMode::BedrockAccessKeys + } else { + DomainAuthMode::BedrockApiKey + }), openai_api_key: None, tokens: None, last_refresh: None, agent_identity: None, personal_access_token: None, - bedrock_api_key: Some(BedrockApiKeyAuth { + bedrock_api_key: (!managed_access_keys).then(|| BedrockApiKeyAuth { api_key: "managed-bedrock-api-key".to_string(), region: "us-west-2".to_string(), }), + bedrock_access_keys: managed_access_keys.then(|| BedrockAccessKeysAuth { + access_key_id: "test-id".to_string(), + secret_access_key: "test-secret".to_string(), + session_token: Some("test-token".to_string()), + }), }) ); assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert!(!codex_home.path().join(".env").exists()); let notification = timeout( DEFAULT_READ_TIMEOUT, @@ -1182,7 +1216,12 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R onboarding_entrypoint: None, } ); - assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; + let auth_mode = if managed_access_keys { + AuthMode::BedrockAccessKeys + } else { + AuthMode::BedrockApiKey + }; + assert_account_updated(&mut mcp, Some(auth_mode)).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1193,6 +1232,35 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R } ); + if managed_access_keys { + let mut expected_logout_config = expected_config; + let expected_logout_config_root = expected_logout_config + .as_table_mut() + .expect("config should be a table"); + expected_logout_config_root.remove("model_provider"); + expected_logout_config_root.remove("model"); + expected_logout_config["model_providers"]["amazon-bedrock"] + .as_table_mut() + .expect("Bedrock provider config should be a table") + .remove("aws"); + + let request_id = mcp.send_logout_account_request().await?; + let response: LogoutAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; + assert_eq!(response, LogoutAccountResponse {}); + assert_eq!(load_file_auth(codex_home.path())?, None); + assert_eq!(read_config_toml(codex_home.path())?, expected_logout_config); + assert!(!codex_home.path().join(".env").exists()); + assert_account_updated(&mut mcp, /*auth_mode*/ None).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: None, + requires_openai_auth: true, + } + ); + } + Ok(()) } @@ -1255,6 +1323,54 @@ async fn login_amazon_bedrock_rejects_non_bedrock_provider_override_without_chan Ok(()) } +#[tokio::test] +async fn login_amazon_bedrock_access_keys_rejects_overridden_aws_configuration() -> Result<()> { + for config_override in [ + r#"model_providers.amazon-bedrock.aws.profile="other-account""#, + r#"model_providers.amazon-bedrock.aws.region="eu-west-1""#, + ] { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let expected_auth = load_file_auth(codex_home.path())?; + + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .with_args(&["-c", config_override]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + + let request_id = mcp + .send_login_account_request(json!({ + "type": "amazonBedrockAccessKeys", + "accessKeyId": "managed-access-key-id", + "secretAccessKey": "managed-secret-access-key", + "region": "us-west-2", + })) + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!( + error.error.message, + "Amazon Bedrock configuration cannot take effect: Overridden by session flags" + ); + assert_eq!(load_file_auth(codex_home.path())?, expected_auth); + } + + Ok(()) +} + #[tokio::test] async fn login_amazon_bedrock_allows_bedrock_provider_override() -> Result<()> { let codex_home = TempDir::new()?; @@ -1301,6 +1417,7 @@ async fn login_amazon_bedrock_allows_bedrock_provider_override() -> Result<()> { api_key: "managed-bedrock-api-key".to_string(), region: "us-west-2".to_string(), }), + bedrock_access_keys: None, }) ); assert_eq!(read_config_toml(codex_home.path())?, expected_config); @@ -1314,15 +1431,15 @@ async fn login_amazon_bedrock_allows_bedrock_provider_override() -> Result<()> { Ok(()) } +#[test_case("amazon-bedrock", "mock-model"; "mantle_clears_generic_model")] +#[test_case("amazon-bedrock-runtime", "global.openai.gpt-5.6-terra"; "runtime_clears_bedrock_model")] #[tokio::test] -async fn logout_managed_bedrock_restores_default_account() -> Result<()> { +async fn logout_managed_bedrock_restores_default_account( + model_provider_id: &str, + model: &str, +) -> Result<()> { let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; - 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::builder() .with_codex_home(codex_home.path()) @@ -1358,6 +1475,33 @@ async fn logout_managed_bedrock_restores_default_account() -> Result<()> { } ); + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)? + .replace( + "model_provider = \"amazon-bedrock\"", + &format!("model_provider = \"{model_provider_id}\""), + ) + .replace( + "model = \"mock-model\"", + &format!("model = \"{model}\"\nmodel_reasoning_effort = \"high\""), + ); + std::fs::write( + config_path, + format!( + "{config}\n[model_providers.{model_provider_id}]\nbase_url = \"https://bedrock.example.com/v1\"\n[model_providers.{model_provider_id}.aws]\nprofile = \"managed-profile\"\nregion = \"us-west-2\"\nauth_refresh = {{ command = \"aws\" }}\n" + ), + )?; + let mut expected_config = read_config_toml(codex_home.path())?; + let expected_config_root = expected_config + .as_table_mut() + .expect("config should be a table"); + expected_config_root.remove("model_provider"); + expected_config_root.remove("model"); + expected_config["model_providers"][model_provider_id] + .as_table_mut() + .expect("Bedrock provider config should be a table") + .remove("aws"); + let request_id = mcp.send_logout_account_request().await?; let response = timeout( DEFAULT_READ_TIMEOUT, @@ -1382,37 +1526,105 @@ async fn logout_managed_bedrock_restores_default_account() -> Result<()> { } #[tokio::test] -async fn logout_aws_managed_bedrock_errors_without_changing_auth_or_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())?; +async fn logout_aws_managed_bedrock_clears_provider_and_restores_default_account() -> Result<()> { + for managed_bedrock_auth in [false, true] { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)? + .replace( + "model = \"mock-model\"", + "model = \"openai.gpt-5.6-sol\"\nmodel_reasoning_effort = \"high\"", + ) + .replace( + "[model_providers.amazon-bedrock.aws]", + "[model_providers.amazon-bedrock]\nbase_url = \"https://bedrock.example.com/v1\"\n[model_providers.amazon-bedrock.aws]\nauth_refresh = { command = \"aws\" }", + ); + std::fs::write(config_path, config)?; + let dotenv_path = codex_home.path().join(".env"); + let aws_credentials_path = codex_home.path().join("aws-credentials"); + let dotenv = "AWS_ACCESS_KEY_ID=environment-id\nAWS_SECRET_ACCESS_KEY=environment-secret\n"; + let aws_credentials = "[codex-bedrock]\naws_access_key_id = profile-id\naws_secret_access_key = profile-secret\n"; + std::fs::write(&dotenv_path, dotenv)?; + std::fs::write(&aws_credentials_path, aws_credentials)?; + if managed_bedrock_auth { + login_with_bedrock_api_key( + codex_home.path(), + "managed-bedrock-api-key", + "us-east-1", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + } else { + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + } - let mut mcp = TestAppServer::builder() - .with_codex_home(codex_home.path()) - .without_auto_env() - .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) - .await?; - let expected_config = read_config_toml(codex_home.path())?; - let request_id = mcp.send_logout_account_request().await?; - let error = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_error_message(RequestId::Integer(request_id)), - ) - .await??; - assert_eq!(error.error.code, -32600); - assert_eq!( - error.error.message, - "cannot log out while Amazon Bedrock is using AWS-managed credentials; manage those credentials through AWS or switch model providers before logging out Codex authentication" - ); - assert_eq!(load_file_auth(codex_home.path())?, expected_auth); - assert_eq!(read_config_toml(codex_home.path())?, expected_config); + let aws_credentials_env_path = aws_credentials_path.to_string_lossy(); + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + ("AWS_ACCESS_KEY_ID", Some("environment-id")), + ("AWS_SECRET_ACCESS_KEY", Some("environment-secret")), + ( + "AWS_SHARED_CREDENTIALS_FILE", + Some(aws_credentials_env_path.as_ref()), + ), + ]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + uses_codex_managed_credentials: false, + }), + requires_openai_auth: false, + } + ); + let mut expected_config = read_config_toml(codex_home.path())?; + let expected_config_root = expected_config + .as_table_mut() + .expect("config should be a table"); + expected_config_root.remove("model_provider"); + expected_config_root.remove("model"); + expected_config["model_providers"]["amazon-bedrock"] + .as_table_mut() + .expect("Bedrock provider config should be a table") + .remove("aws"); + + 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::(response)?, + LogoutAccountResponse {} + ); + assert_eq!(load_file_auth(codex_home.path())?, None); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_eq!(std::fs::read_to_string(dotenv_path)?, dotenv); + assert_eq!( + std::fs::read_to_string(aws_credentials_path)?, + aws_credentials + ); + assert_account_updated(&mut mcp, /*auth_mode*/ None).await?; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: None, + requires_openai_auth: true, + } + ); + } Ok(()) } @@ -1450,6 +1662,14 @@ async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_ assert!(matches!(initialized, JSONRPCMessage::Response(_))); create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?; + std::fs::write( + config_path, + format!( + "{config}\n[model_providers.amazon-bedrock.aws]\nprofile = \"preserved\"\nregion = \"us-west-2\"\n" + ), + )?; let expected_config = read_config_toml(codex_home.path())?; let request_id = mcp.send_logout_account_request().await?; @@ -1594,7 +1814,25 @@ async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_change .await??; assert_eq!( error.error.message, - "Amazon Bedrock Mantle does not support region `us-west-1`" + "Amazon Bedrock does not support region `us-west-1`" + ); + + let request_id = mcp + .send_login_account_request(json!({ + "type": "amazonBedrockAccessKeys", + "accessKeyId": " ", + "secretAccessKey": "test-secret", + "region": "us-west-2", + })) + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "AWS access key ID and secret access key must not be empty." ); assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 06a25f3dba..84cdd36c04 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -120,6 +120,7 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), diff --git a/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs b/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs index b46cfd0bd5..87cfaf3ae4 100644 --- a/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs +++ b/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs @@ -173,7 +173,7 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { json!({"region": "us-east-1"}) ); - let config: toml::Value = toml::from_str(&std::fs::read_to_string(config_path)?)?; + let config: toml::Value = toml::from_str(&std::fs::read_to_string(&config_path)?)?; assert_eq!( config, toml::toml! { @@ -247,7 +247,7 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { .await??; assert_eq!( error.error.message, - "A Bedrock API key is already configured and takes priority over AWS credentials. Run `codex logout` and try again." + "Codex-managed Bedrock credentials are already configured and take priority over AWS environment credentials. Run `codex logout` and try again." ); assert_eq!(std::fs::read_to_string(&auth_path)?, expected_auth); assert_eq!( @@ -267,7 +267,7 @@ async fn setup_bedrock_rejects_invalid_or_conflicting_credentials() -> Result<() for (params, expected_error) in [ ( json!({"type": "profile", "profile": "engineering", "region": "us-west-1"}), - "Amazon Bedrock Mantle does not support region `us-west-1`", + "Amazon Bedrock does not support region `us-west-1`", ), ( json!({"type": "profile", "profile": " ", "region": "us-west-2"}), @@ -281,15 +281,6 @@ async fn setup_bedrock_rejects_invalid_or_conflicting_credentials() -> Result<() json!({"type": "environment", "region": "us-west-2"}), "No AWS credentials found. Please Configure AWS credentials or complete AWS sign-in, then try again.", ), - ( - json!({ - "type": "accessKeys", - "accessKeyId": "test-id", - "secretAccessKey": "test-secret", - "region": "us-west-2", - }), - "Direct AWS access key setup is not supported yet.", - ), ] { let request_id = app_server .send_raw_request("account/bedrock/setup", Some(params)) diff --git a/codex-rs/aws-auth/src/lib.rs b/codex-rs/aws-auth/src/lib.rs index e361dddb8f..8f1bc2b5d7 100644 --- a/codex-rs/aws-auth/src/lib.rs +++ b/codex-rs/aws-auth/src/lib.rs @@ -23,6 +23,27 @@ pub struct AwsAuthConfig { pub service: String, } +/// Static AWS access keys supplied by a caller instead of the default SDK chain. +#[derive(Clone, PartialEq, Eq)] +pub struct AwsAccessKeys { + pub access_key_id: String, + pub secret_access_key: String, + pub session_token: Option, +} + +impl std::fmt::Debug for AwsAccessKeys { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AwsAccessKeys") + .field("access_key_id", &"") + .field("secret_access_key", &"") + .field( + "session_token", + &self.session_token.as_ref().map(|_| ""), + ) + .finish() + } +} + /// Generic HTTP request shape consumed by SigV4 signing. #[derive(Debug, Clone, PartialEq, Eq)] pub struct AwsRequestToSign { @@ -98,6 +119,22 @@ impl AwsAuthContext { }) } + pub async fn load_with_access_keys( + config: AwsAuthConfig, + access_keys: AwsAccessKeys, + ) -> Result { + let mut context = Self::load(config).await?; + context.credentials_provider = + SharedCredentialsProvider::new(aws_credential_types::Credentials::new( + access_keys.access_key_id, + access_keys.secret_access_key, + access_keys.session_token, + /*expires_after*/ None, + "codex-managed-bedrock-access-keys", + )); + Ok(context) + } + pub async fn load_profile(config: AwsAuthConfig) -> Result { let profile = config .profile diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index 6228c77d91..3fba0e2fdf 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -1416,6 +1416,7 @@ fn stored_auth_mode(auth: &codex_login::AuthDotJson) -> &'static str { AuthMode::AgentIdentity => "agent_identity", AuthMode::PersonalAccessToken => "personal_access_token", AuthMode::BedrockApiKey => "bedrock_api_key", + AuthMode::BedrockAccessKeys => "bedrock_access_keys", } } @@ -1427,6 +1428,8 @@ fn stored_auth_mode_value(auth: &AuthDotJson) -> AuthMode { AuthMode::PersonalAccessToken } else if auth.bedrock_api_key.is_some() { AuthMode::BedrockApiKey + } else if auth.bedrock_access_keys.is_some() { + AuthMode::BedrockAccessKeys } else if auth.openai_api_key.is_some() { AuthMode::ApiKey } else { @@ -1509,6 +1512,17 @@ fn stored_auth_issues( issues.push("Bedrock API key auth is missing a Bedrock API key"); } } + AuthMode::BedrockAccessKeys => match auth.bedrock_access_keys.as_ref() { + Some(access_keys) => { + if access_keys.access_key_id.trim().is_empty() { + issues.push("Bedrock access key auth is missing an access key ID"); + } + if access_keys.secret_access_key.trim().is_empty() { + issues.push("Bedrock access key auth is missing a secret access key"); + } + } + None => issues.push("Bedrock access key auth is missing AWS access keys"), + }, } issues } @@ -2495,6 +2509,7 @@ fn auth_mode_name(auth: &CodexAuth) -> &'static str { AuthMode::AgentIdentity => "agent_identity", AuthMode::PersonalAccessToken => "personal_access_token", AuthMode::BedrockApiKey => "bedrock_api_key", + AuthMode::BedrockAccessKeys => "bedrock_access_keys", } } @@ -2644,7 +2659,9 @@ fn provider_auth_reachability_mode_from_auth( return ProviderAuthReachabilityMode::Chatgpt; } match stored_auth.map(stored_auth_mode_value) { - Some(AuthMode::ApiKey | AuthMode::BedrockApiKey) => ProviderAuthReachabilityMode::ApiKey, + Some(AuthMode::ApiKey | AuthMode::BedrockApiKey | AuthMode::BedrockAccessKeys) => { + ProviderAuthReachabilityMode::ApiKey + } Some( AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens @@ -3566,6 +3583,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; assert_eq!( @@ -3585,6 +3603,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; assert_eq!( @@ -3606,6 +3625,7 @@ mod tests { agent_identity: None, personal_access_token: Some("at-test".to_string()), bedrock_api_key: None, + bedrock_access_keys: None, }; assert_eq!(stored_auth_mode(&auth), "personal_access_token"); @@ -3629,6 +3649,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; assert_eq!( diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 5456d1a04c..22e607c6e6 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -9,6 +9,8 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::Config; +use codex_core::config::edit::ConfigEdit; +use codex_core::config::edit::ConfigEditsBuilder; use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::AuthRouteConfig; @@ -489,6 +491,10 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { eprintln!("Logged in using Amazon Bedrock API key"); std::process::exit(0); } + AuthMode::BedrockAccessKeys => { + eprintln!("Logged in using Amazon Bedrock AWS access keys"); + std::process::exit(0); + } }, Ok(None) => { eprintln!("Not logged in"); @@ -505,7 +511,7 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; let auth_route_config = config.auth_route_config(); - match logout_with_revoke( + let logged_out = match logout_with_revoke( &config.codex_home, config.cli_auth_credentials_store_mode, config.auth_keyring_backend_kind(), @@ -513,19 +519,37 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! { ) .await { - Ok(true) => { - eprintln!("Successfully logged out"); - std::process::exit(0); - } - Ok(false) => { - eprintln!("Not logged in"); - std::process::exit(0); - } - Err(e) => { - eprintln!("Error logging out: {e}"); + Ok(logged_out) => logged_out, + Err(err) => { + eprintln!("Error logging out: {err}"); std::process::exit(1); } + }; + + let cleared_bedrock_config = + if let Some(paths) = ConfigEditsBuilder::bedrock_provider_config_paths_to_clear(&config) { + let edits = paths + .into_iter() + .map(|segments| ConfigEdit::ClearPath { segments }); + if let Err(err) = ConfigEditsBuilder::for_config(&config) + .with_edits(edits) + .apply() + .await + { + eprintln!("Error clearing Amazon Bedrock configuration after logout: {err}"); + std::process::exit(1); + } + true + } else { + false + }; + + if logged_out || cleared_bedrock_config { + eprintln!("Successfully logged out"); + } else { + eprintln!("Not logged in"); } + std::process::exit(0); } async fn load_config_or_exit(cli_config_overrides: CliConfigOverrides) -> Config { diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs index 3eb81143f5..f631057e19 100644 --- a/codex-rs/cli/tests/login.rs +++ b/codex-rs/cli/tests/login.rs @@ -8,9 +8,11 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::write_chatgpt_auth; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; use codex_login::CLIENT_ID; use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR; use codex_login::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR; +use codex_login::login_with_bedrock_access_keys; use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR; use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR; use predicates::str::contains; @@ -102,6 +104,85 @@ fn login_status_validates_configured_workload_identity() -> Result<()> { Ok(()) } +#[test] +fn logout_clears_only_the_selected_bedrock_provider() -> Result<()> { + for (model_provider_id, managed_bedrock_auth, model) in [ + ("amazon-bedrock", true, "openai.gpt-5.6-sol"), + ("amazon-bedrock-runtime", true, "global.openai.gpt-5.6-sol"), + ("openai", true, "gpt-5.6-sol"), + ("amazon-bedrock", false, "gpt-5.6-sol"), + ("amazon-bedrock-runtime", false, "us.openai.gpt-5.6-sol"), + ("openai", false, "gpt-5.6-sol"), + ] { + let codex_home = TempDir::new()?; + let config_path = codex_home.path().join("config.toml"); + std::fs::write( + &config_path, + format!( + "cli_auth_credentials_store = \"file\"\n\ + model_provider = \"{model_provider_id}\"\n\ + model = \"{model}\"\n\ + model_reasoning_effort = \"high\"\n\ + [model_providers.amazon-bedrock]\n\ + base_url = \"https://mantle.example.com/v1\"\n\ + [model_providers.amazon-bedrock.aws]\n\ + profile = \"mantle-profile\"\n\ + region = \"us-west-2\"\n\ + auth_refresh = {{ command = \"aws\", args = [\"sso\", \"login\"] }}\n\ + [model_providers.amazon-bedrock-runtime]\n\ + base_url = \"https://runtime.example.com/v1\"\n\ + [model_providers.amazon-bedrock-runtime.aws]\n\ + profile = \"runtime-profile\"\n\ + region = \"us-east-1\"\n\ + auth_refresh = {{ command = \"aws\", args = [\"login\"] }}\n" + ), + )?; + if managed_bedrock_auth { + login_with_bedrock_access_keys( + codex_home.path(), + "managed-access-key-id", + "managed-secret-access-key", + Some("managed-session-token"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + } + let mut expected_config: toml::Value = + toml::from_str(&std::fs::read_to_string(&config_path)?)?; + if model_provider_id != "openai" { + let expected_root = expected_config + .as_table_mut() + .expect("config should be a table"); + expected_root.remove("model_provider"); + expected_root.remove("model"); + expected_root["model_providers"][model_provider_id] + .as_table_mut() + .expect("selected Bedrock provider should be a table") + .remove("aws"); + } + let expected_message = if managed_bedrock_auth || model_provider_id != "openai" { + "Successfully logged out" + } else { + "Not logged in" + }; + + codex_command(codex_home.path())? + .env_remove(CODEX_ACCESS_TOKEN_ENV_VAR) + .env("AWS_ACCESS_KEY_ID", "environment-access-key-id") + .env("AWS_SECRET_ACCESS_KEY", "environment-secret-access-key") + .args(["logout"]) + .assert() + .success() + .stderr(contains(expected_message)); + + assert!(!codex_home.path().join("auth.json").exists()); + let actual_config: toml::Value = toml::from_str(&std::fs::read_to_string(&config_path)?)?; + assert_eq!(actual_config, expected_config); + } + + Ok(()) +} + #[test] fn login_with_access_token_rejects_invalid_jwt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core-plugins/src/test_support.rs b/codex-rs/core-plugins/src/test_support.rs index 0cae12bcca..dc16dad05f 100644 --- a/codex-rs/core-plugins/src/test_support.rs +++ b/codex-rs/core-plugins/src/test_support.rs @@ -27,6 +27,7 @@ use codex_login::CodexAuth; use codex_login::ExternalAuth; use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; +use codex_login::auth::BedrockAccessKeysAuth; use codex_login::auth::BedrockApiKeyAuth; use codex_login::test_support::auth_manager_from_optional_auth; use codex_protocol::auth::AuthMode; @@ -118,6 +119,11 @@ fn test_codex_auth(auth_mode: Option) -> Option { api_key: "test-api-key".to_string(), region: "us-east-1".to_string(), }), + AuthMode::BedrockAccessKeys => CodexAuth::BedrockAccessKeys(BedrockAccessKeysAuth { + access_key_id: "test-access-key-id".to_string(), + secret_access_key: "test-secret-access-key".to_string(), + session_token: None, + }), AuthMode::AgentIdentity | AuthMode::PersonalAccessToken => { panic!("test auth mode requires a purpose-built CodexAuth") } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index b3e85d18d6..ed0f10d3d4 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2225,7 +2225,9 @@ impl AuthRequestTelemetryContext { let auth_telemetry = auth_header_telemetry(api_auth); Self { auth_mode: auth_mode.map(|mode| match mode { - AuthMode::ApiKey | AuthMode::BedrockApiKey => "ApiKey", + AuthMode::ApiKey | AuthMode::BedrockApiKey | AuthMode::BedrockAccessKeys => { + "ApiKey" + } AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::Headers diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index be4237429b..6d13d9a755 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -26,6 +26,7 @@ use toml_edit::value; const NOTICE_TABLE_KEY: &str = "notice"; +mod bedrock; mod document_helpers; /// Discrete config mutations supported by the persistence engine. diff --git a/codex-rs/core/src/config/edit/bedrock.rs b/codex-rs/core/src/config/edit/bedrock.rs new file mode 100644 index 0000000000..74e2eba24b --- /dev/null +++ b/codex-rs/core/src/config/edit/bedrock.rs @@ -0,0 +1,37 @@ +use super::ConfigEditsBuilder; +use crate::config::Config; +use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; +use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_PROVIDER_ID; + +impl ConfigEditsBuilder { + /// Return the user config paths to clear for the selected Bedrock provider. + pub fn bedrock_provider_config_paths_to_clear(config: &Config) -> Option>> { + let provider_id = config.model_provider_id.as_str(); + if !matches!( + provider_id, + AMAZON_BEDROCK_PROVIDER_ID | AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ) { + return None; + } + + let user_layer = config.config_layer_stack.get_active_user_layer()?; + if user_layer + .config + .get("model_provider") + .and_then(toml::Value::as_str) + != Some(provider_id) + { + return None; + } + + Some(vec![ + vec!["model_provider".to_string()], + vec![ + "model_providers".to_string(), + provider_id.to_string(), + "aws".to_string(), + ], + vec!["model".to_string()], + ]) + } +} diff --git a/codex-rs/core/tests/suite/external_auth.rs b/codex-rs/core/tests/suite/external_auth.rs index 2b5593c23d..5d2e01434f 100644 --- a/codex-rs/core/tests/suite/external_auth.rs +++ b/codex-rs/core/tests/suite/external_auth.rs @@ -3,6 +3,7 @@ use codex_login::CodexAuth; use codex_login::ExternalAuth; use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; +use codex_login::auth::BedrockAccessKeysAuth; use codex_model_provider_info::AwsAuthRefreshConfig; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; @@ -181,6 +182,51 @@ async fn custom_provider_uses_explicit_bearer_without_ambient_account() -> anyho Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn amazon_bedrock_managed_access_keys_sign_requests() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response_mock = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let mut provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-east-1".to_string()), + auth_refresh: None, + })); + provider.base_url = Some(format!("{}/v1", server.uri())); + + let mut builder = test_codex() + .with_model("openai.gpt-5.5") + .with_auth(CodexAuth::BedrockAccessKeys(BedrockAccessKeysAuth { + access_key_id: "managed-access-key-id".to_string(), + secret_access_key: "managed-secret-access-key".to_string(), + session_token: Some("managed-session-token".to_string()), + })) + .with_config(move |config| { + config.model_provider_id = provider.name.clone(); + config.model_provider = provider; + }); + let test = builder.build_with_auto_env(&server).await?; + test.submit_turn("hello").await?; + + let request = response_mock.single_request(); + let authorization = request + .header("authorization") + .expect("managed AWS credentials should produce SigV4 authorization"); + assert!(authorization.starts_with("AWS4-HMAC-SHA256 ")); + assert!(authorization.contains("Credential=managed-access-key-id/")); + assert_eq!( + request.header("x-amz-security-token").as_deref(), + Some("managed-session-token") + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/exec/tests/suite/apply_patch.rs b/codex-rs/exec/tests/suite/apply_patch.rs index 9683efbd95..8376b2e765 100644 --- a/codex-rs/exec/tests/suite/apply_patch.rs +++ b/codex-rs/exec/tests/suite/apply_patch.rs @@ -208,6 +208,7 @@ async fn shutdown_flushes_completed_turn_and_file_diff() -> anyhow::Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, codex_login::AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 8a642b0a54..128273381e 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -370,6 +370,7 @@ async fn stored_agent_identity_jwt_keeps_auth_json_unchanged() -> anyhow::Result agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity.clone())), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::Direct, @@ -449,6 +450,7 @@ async fn login_with_access_token_writes_only_personal_access_token() { agent_identity: None, personal_access_token: Some("at-login-test".to_string()), bedrock_api_key: None, + bedrock_access_keys: None, } ); assert_eq!(auth.resolved_mode(), AuthMode::PersonalAccessToken); @@ -1068,6 +1070,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, auth_dot_json ); @@ -1116,6 +1119,7 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; super::save_auth( dir.path(), @@ -2331,6 +2335,7 @@ async fn workspace_policy_rejects_agent_identity_before_hydration() { agent_identity: Some(stored_agent_identity), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::Direct, @@ -2572,6 +2577,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), diff --git a/codex-rs/login/src/auth/bedrock_access_keys.rs b/codex-rs/login/src/auth/bedrock_access_keys.rs new file mode 100644 index 0000000000..304400c353 --- /dev/null +++ b/codex-rs/login/src/auth/bedrock_access_keys.rs @@ -0,0 +1,63 @@ +use std::path::Path; + +use codex_config::types::AuthCredentialsStoreMode; +use codex_protocol::auth::AuthMode; +use serde::Deserialize; +use serde::Serialize; + +use super::manager::save_auth; +use super::storage::AuthDotJson; +use super::storage::AuthKeyringBackendKind; + +/// Managed Amazon Bedrock AWS access keys persisted in auth storage. +#[derive(Deserialize, Serialize, Clone, PartialEq, Eq)] +pub struct BedrockAccessKeysAuth { + pub access_key_id: String, + pub secret_access_key: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_token: Option, +} + +impl std::fmt::Debug for BedrockAccessKeysAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BedrockAccessKeysAuth") + .field("access_key_id", &"") + .field("secret_access_key", &"") + .field( + "session_token", + &self.session_token.as_ref().map(|_| ""), + ) + .finish() + } +} + +/// Writes auth storage containing only Amazon Bedrock AWS access keys. +pub fn login_with_bedrock_access_keys( + codex_home: &Path, + access_key_id: &str, + secret_access_key: &str, + session_token: Option<&str>, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::BedrockAccessKeys), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + bedrock_access_keys: Some(BedrockAccessKeysAuth { + access_key_id: access_key_id.to_string(), + secret_access_key: secret_access_key.to_string(), + session_token: session_token.map(str::to_string), + }), + }; + save_auth( + codex_home, + &auth_dot_json, + auth_credentials_store_mode, + keyring_backend_kind, + ) +} diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs index 3445878365..cae6ddf548 100644 --- a/codex-rs/login/src/auth/bedrock_api_key.rs +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -35,6 +35,7 @@ pub fn login_with_bedrock_api_key( api_key: api_key.to_string(), region: region.to_string(), }), + bedrock_access_keys: None, }; save_auth( codex_home, diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index 282fff98f9..70764abff8 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -20,6 +20,7 @@ fn api_key_auth() -> AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } @@ -32,6 +33,7 @@ fn bedrock_only_auth() -> AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: Some(bedrock_auth()), + bedrock_access_keys: None, } } @@ -76,6 +78,7 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> agent_identity: None, personal_access_token: None, bedrock_api_key: Some(bedrock_auth()), + bedrock_access_keys: None, }; assert_eq!(loaded, expected); assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); @@ -87,7 +90,8 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::Headers(_) | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockAccessKeys(_) => None, }), Some(bedrock_auth()) ); @@ -124,6 +128,60 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +#[serial(codex_auth_env)] +async fn access_keys_auth_round_trips_and_logs_out() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + crate::auth::login_with_bedrock_access_keys( + codex_home.path(), + "access-key-id", + "secret-access-key", + Some("session-token"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + Some(vec!["allowed-workspace".to_string()]), + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + crate::test_support::transport_default_auth_route_config(), + ) + .await; + + assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockAccessKeys)); + assert_eq!( + auth_manager.auth_cached().and_then(|auth| match auth { + CodexAuth::BedrockAccessKeys(auth) => Some(auth), + CodexAuth::ApiKey(_) + | CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::Headers(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockApiKey(_) => None, + }), + Some(crate::auth::BedrockAccessKeysAuth { + access_key_id: "access-key-id".to_string(), + secret_access_key: "secret-access-key".to_string(), + session_token: Some("session-token".to_string()), + }) + ); + let auth_change_rx = auth_manager.auth_change_receiver(); + let auth_revision = *auth_change_rx.borrow(); + let mut auth_without_mode = storage.load()?.expect("access keys should be stored"); + auth_without_mode.auth_mode = None; + storage.save(&auth_without_mode)?; + assert!(!auth_manager.reload().await); + assert_eq!(*auth_change_rx.borrow(), auth_revision); + assert!(auth_manager.logout().await?); + assert_eq!(storage.load()?, None); + Ok(()) +} + #[tokio::test] #[serial(codex_auth_env)] async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> { @@ -151,7 +209,8 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::Headers(_) | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockAccessKeys(_) => None, }), Some(bedrock_auth()) ); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 01c039c083..f4f415e253 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -43,6 +43,7 @@ use super::workload_identity::WorkloadIdentitySessionError; use crate::auth::AuthHeaders; pub use crate::auth::agent_identity::AgentIdentityAuth; pub use crate::auth::agent_identity::AgentIdentityAuthError; +pub use crate::auth::bedrock_access_keys::BedrockAccessKeysAuth; pub use crate::auth::bedrock_api_key::BedrockApiKeyAuth; pub use crate::auth::personal_access_token::PersonalAccessTokenAuth; pub use crate::auth::storage::AgentIdentityAuthRecord; @@ -81,6 +82,7 @@ pub enum CodexAuth { AgentIdentity(AgentIdentityAuth), PersonalAccessToken(PersonalAccessTokenAuth), BedrockApiKey(BedrockApiKeyAuth), + BedrockAccessKeys(BedrockAccessKeysAuth), } /// Policy for resolving Agent Identity auth from a broader Codex auth snapshot. @@ -155,6 +157,7 @@ impl PartialEq for CodexAuth { (Self::Headers(a), Self::Headers(b)) => a == b, (Self::PersonalAccessToken(a), Self::PersonalAccessToken(b)) => a == b, (Self::BedrockApiKey(a), Self::BedrockApiKey(b)) => a == b, + (Self::BedrockAccessKeys(a), Self::BedrockAccessKeys(b)) => a == b, _ => self.api_auth_mode() == other.api_auth_mode(), } } @@ -359,6 +362,14 @@ impl CodexAuth { }; return Ok(Self::BedrockApiKey(auth)); } + if auth_mode == AuthMode::BedrockAccessKeys { + let Some(auth) = auth_dot_json.bedrock_access_keys else { + return Err(std::io::Error::other( + "Bedrock access keys auth is missing AWS access keys.", + )); + }; + return Ok(Self::BedrockAccessKeys(auth)); + } if auth_mode == AuthMode::Headers { return Err(std::io::Error::other( "externally provided auth cannot be loaded from auth storage.", @@ -391,6 +402,9 @@ impl CodexAuth { unreachable!("personal access token mode is handled above") } AuthMode::BedrockApiKey => unreachable!("bedrock api key mode is handled above"), + AuthMode::BedrockAccessKeys => { + unreachable!("bedrock access keys mode is handled above") + } } } @@ -473,6 +487,7 @@ impl CodexAuth { Self::AgentIdentity(_) => AuthMode::AgentIdentity, Self::PersonalAccessToken(_) => AuthMode::PersonalAccessToken, Self::BedrockApiKey(_) => AuthMode::BedrockApiKey, + Self::BedrockAccessKeys(_) => AuthMode::BedrockAccessKeys, } } @@ -486,6 +501,7 @@ impl CodexAuth { Self::AgentIdentity(_) => AuthMode::AgentIdentity, Self::PersonalAccessToken(_) => AuthMode::PersonalAccessToken, Self::BedrockApiKey(_) => AuthMode::BedrockApiKey, + Self::BedrockAccessKeys(_) => AuthMode::BedrockAccessKeys, } } @@ -525,7 +541,8 @@ impl CodexAuth { | Self::Headers(_) | Self::AgentIdentity(_) | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => None, + | Self::BedrockApiKey(_) + | Self::BedrockAccessKeys(_) => None, } } @@ -557,7 +574,7 @@ impl CodexAuth { "header auth does not expose a bearer token", )), Self::PersonalAccessToken(auth) => Ok(auth.access_token().to_string()), - Self::BedrockApiKey(_) => Err(std::io::Error::other( + Self::BedrockApiKey(_) | Self::BedrockAccessKeys(_) => Err(std::io::Error::other( "Bedrock API key auth does not expose a Codex bearer token", )), } @@ -648,7 +665,8 @@ impl CodexAuth { | Self::Headers(_) | Self::AgentIdentity(_) | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => return None, + | Self::BedrockApiKey(_) + | Self::BedrockAccessKeys(_) => return None, }; #[expect(clippy::unwrap_used)] state.auth_dot_json.lock().unwrap().clone() @@ -693,7 +711,8 @@ impl CodexAuth { | Self::ChatgptAuthTokens(_) | Self::Headers(_) | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => Ok(None), + | Self::BedrockApiKey(_) + | Self::BedrockAccessKeys(_) => Ok(None), Self::Chatgpt(_) => { if policy == AgentIdentityAuthPolicy::JwtOnly { return Ok(None); @@ -766,6 +785,7 @@ impl CodexAuth { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; let state = ChatgptAuthState { @@ -970,6 +990,7 @@ pub fn login_with_api_key( agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( codex_home, @@ -1003,6 +1024,7 @@ pub async fn login_with_access_token( agent_identity: None, personal_access_token: Some(access_token.to_string()), bedrock_api_key: None, + bedrock_access_keys: None, } } CodexAccessToken::AgentIdentityJwt(jwt) => { @@ -1021,6 +1043,7 @@ pub async fn login_with_access_token( agent_identity: Some(AgentIdentityStorage::Jwt(jwt.to_string())), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } }; @@ -1219,7 +1242,10 @@ fn validate_auth_restrictions( let Some(expected_workspaces) = expected_workspaces else { return Ok(()); }; - if matches!(auth, CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_)) { + if matches!( + auth, + CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) | CodexAuth::BedrockAccessKeys(_) + ) { return Ok(()); } @@ -1281,7 +1307,8 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( if let Some(required_method) = config.forced_login_method { let method_violation = match (required_method, auth.auth_mode()) { (ForcedLoginMethod::Api, AuthMode::ApiKey) - | (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) => None, + | (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) + | (ForcedLoginMethod::Api, AuthMode::BedrockAccessKeys) => None, (ForcedLoginMethod::Chatgpt, AuthMode::Chatgpt) | (ForcedLoginMethod::Chatgpt, AuthMode::ChatgptAuthTokens) | (ForcedLoginMethod::Chatgpt, AuthMode::Headers) @@ -1296,7 +1323,8 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( .to_string(), ), (ForcedLoginMethod::Chatgpt, AuthMode::ApiKey) - | (ForcedLoginMethod::Chatgpt, AuthMode::BedrockApiKey) => Some( + | (ForcedLoginMethod::Chatgpt, AuthMode::BedrockApiKey) + | (ForcedLoginMethod::Chatgpt, AuthMode::BedrockAccessKeys) => Some( "ChatGPT login is required, but an API key is currently being used. Logging out." .to_string(), ), @@ -1314,7 +1342,9 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { let chatgpt_account_id = match &auth { - CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => { + CodexAuth::ApiKey(_) + | CodexAuth::BedrockApiKey(_) + | CodexAuth::BedrockAccessKeys(_) => { return Ok(()); } CodexAuth::Headers(_) @@ -1717,6 +1747,7 @@ impl AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }) } @@ -1730,6 +1761,9 @@ impl AuthDotJson { if self.bedrock_api_key.is_some() { return AuthMode::BedrockApiKey; } + if self.bedrock_access_keys.is_some() { + return AuthMode::BedrockAccessKeys; + } if self.openai_api_key.is_some() { return AuthMode::ApiKey; } @@ -2449,6 +2483,7 @@ impl AuthManager { }, (AuthMode::PersonalAccessToken, AuthMode::PersonalAccessToken) => a == b, (AuthMode::BedrockApiKey, AuthMode::BedrockApiKey) => a == b, + (AuthMode::BedrockAccessKeys, AuthMode::BedrockAccessKeys) => a == b, _ => false, }, _ => false, @@ -2810,7 +2845,8 @@ impl AuthManager { | CodexAuth::Headers(_) | CodexAuth::AgentIdentity(_) | CodexAuth::PersonalAccessToken(_) - | CodexAuth::BedrockApiKey(_) => Ok(()), + | CodexAuth::BedrockApiKey(_) + | CodexAuth::BedrockAccessKeys(_) => Ok(()), } }; if let Err(RefreshTokenError::Permanent(error)) = &result { diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 8c76713b8f..e4898f0976 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -1,6 +1,7 @@ mod access_token; mod agent_identity; mod auth_headers; +mod bedrock_access_keys; mod bedrock_api_key; pub mod default_client; pub mod error; @@ -14,6 +15,8 @@ mod manager; mod revoke; pub use auth_headers::AuthHeaders; +pub use bedrock_access_keys::BedrockAccessKeysAuth; +pub use bedrock_access_keys::login_with_bedrock_access_keys; pub use bedrock_api_key::BedrockApiKeyAuth; pub use bedrock_api_key::login_with_bedrock_api_key; pub use error::RefreshTokenFailedError; diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index ae195f3bff..8a9d75f448 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use std::sync::Mutex; use tracing::warn; +use super::BedrockAccessKeysAuth; use super::BedrockApiKeyAuth; use crate::token_data::TokenData; use codex_agent_identity::AgentIdentityJwtClaims; @@ -58,6 +59,9 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub bedrock_api_key: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub bedrock_access_keys: Option, } #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 647af79b65..0c69343808 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -26,6 +26,7 @@ async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage @@ -49,6 +50,7 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; let file = get_auth_file(codex_home.path()); @@ -84,6 +86,7 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth_dot_json)?; @@ -115,6 +118,7 @@ async fn file_storage_round_trips_registered_agent_identity_auth() -> anyhow::Re agent_identity: Some(AgentIdentityStorage::Record(record)), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth_dot_json)?; @@ -166,6 +170,7 @@ async fn file_storage_loads_empty_agent_identity_email_as_none() -> anyhow::Resu })), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }) ); Ok(()) @@ -192,6 +197,7 @@ async fn file_storage_writes_missing_agent_identity_email_as_empty_string() -> a })), personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth_dot_json)?; @@ -215,6 +221,7 @@ async fn file_storage_round_trips_personal_access_token_auth() -> anyhow::Result agent_identity: None, personal_access_token: Some("at-example".to_string()), bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth_dot_json)?; @@ -266,6 +273,7 @@ fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; let storage = create_auth_storage( dir.path().to_path_buf(), @@ -297,6 +305,7 @@ fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth_dot_json)?; @@ -430,6 +439,7 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } @@ -457,6 +467,7 @@ fn secrets_keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Resu agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; seed_secrets_backend_with_auth(&mock_keyring, codex_home.path(), &expected)?; @@ -595,6 +606,7 @@ fn secrets_keyring_auth_storage_save_persists_and_removes_fallback_file() -> any agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; storage.save(&auth)?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 717bd0f4b0..33659722fd 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -54,6 +54,7 @@ pub use auth::is_workload_identity_selected; pub use auth::load_auth_dot_json; pub use auth::login_with_access_token; pub use auth::login_with_api_key; +pub use auth::login_with_bedrock_access_keys; pub use auth::login_with_bedrock_api_key; pub use auth::logout; pub use auth::logout_with_revoke; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index dbfddc3d3b..677c740da2 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -915,6 +915,7 @@ pub(crate) async fn persist_tokens_async( agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( &codex_home, diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index 68a82feb2d..3a8579282d 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -154,6 +154,7 @@ async fn refresh_token_honors_respect_system_proxy() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( codex_home.path(), @@ -207,6 +208,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -283,6 +285,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -350,6 +353,7 @@ async fn auth_refreshes_when_access_token_is_near_expiry() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -402,6 +406,7 @@ async fn auth_skips_access_token_outside_refresh_window() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -440,6 +445,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -452,6 +458,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -510,6 +517,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -523,6 +531,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -585,6 +594,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -635,6 +645,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -687,6 +698,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -700,6 +712,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -755,6 +768,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -768,6 +782,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -821,6 +836,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -876,6 +892,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -945,6 +962,7 @@ async fn refresh_token_does_not_retry_after_bad_request_reused_failure() -> Resu agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1012,6 +1030,7 @@ async fn refresh_token_does_not_retry_after_standard_invalid_grant_failure() -> agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1079,6 +1098,7 @@ async fn refresh_token_does_not_cache_other_bad_request_failure() -> Result<()> agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1144,6 +1164,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1168,6 +1189,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -1230,6 +1252,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1285,6 +1308,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1297,6 +1321,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -1384,6 +1409,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&initial_auth).await?; @@ -1397,6 +1423,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( ctx.codex_home.path(), @@ -1458,6 +1485,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; ctx.write_auth(&auth).await?; diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index 2d06522a71..3ecdf0150f 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -263,6 +263,7 @@ fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 316e38e348..706fb6e283 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use codex_api::AuthError; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; +use codex_aws_auth::AwsAccessKeys; use codex_aws_auth::AwsAuthContext; use codex_aws_auth::AwsAuthError; use codex_aws_auth::AwsRequestToSign; @@ -11,7 +12,6 @@ use codex_http_client::RequestBody; use codex_http_client::RequestCompression; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; @@ -36,6 +36,7 @@ pub(super) enum BedrockAuthSource { CommandBearerToken, ConfiguredAwsProfile, ManagedBearerToken, + ManagedAccessKeys, EnvBearerToken, EnvAwsCredentials, AwsSdk, @@ -65,6 +66,11 @@ pub(super) fn auth_source( Some(CodexAuth::BedrockApiKey(_)) ) { BedrockAuthSource::ManagedBearerToken + } else if matches!( + auth_manager.and_then(AuthManager::auth_cached), + Some(CodexAuth::BedrockAccessKeys(_)) + ) { + BedrockAuthSource::ManagedAccessKeys } else if non_empty_env_var_from(AWS_BEARER_TOKEN_BEDROCK_ENV_VAR, env_var).is_some() { BedrockAuthSource::EnvBearerToken } else if non_empty_env_var_from(AWS_ACCESS_KEY_ID_ENV_VAR, env_var).is_some() @@ -78,7 +84,7 @@ pub(super) fn auth_source( pub(super) async fn resolve_auth_method( source: BedrockAuthSource, - managed_auth: Option<&BedrockApiKeyAuth>, + managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, ) -> Result { @@ -88,15 +94,15 @@ pub(super) async fn resolve_auth_method( .to_string(), )), BedrockAuthSource::ManagedBearerToken => { - let managed_auth = managed_auth.ok_or_else(|| { - CodexErr::Fatal( + let Some(CodexAuth::BedrockApiKey(auth)) = managed_auth else { + return Err(CodexErr::Fatal( "selected Codex-managed Amazon Bedrock API key is no longer available" .to_string(), - ) - })?; + )); + }; Ok(BedrockAuthMethod::ManagedBearerToken { - token: managed_auth.api_key.clone(), - region: managed_auth.region.clone(), + token: auth.api_key.clone(), + region: auth.region.clone(), }) } BedrockAuthSource::EnvBearerToken => { @@ -120,6 +126,27 @@ pub(super) async fn resolve_auth_method( .map_err(aws_auth_error_to_codex_error)?; Ok(BedrockAuthMethod::AwsSdkAuth { context }) } + BedrockAuthSource::ManagedAccessKeys => { + let Some(CodexAuth::BedrockAccessKeys(auth)) = managed_auth else { + return Err(CodexErr::Fatal( + "selected Codex-managed Amazon Bedrock access keys are no longer available" + .to_string(), + )); + }; + let access_keys = AwsAccessKeys { + access_key_id: auth.access_key_id.clone(), + secret_access_key: auth.secret_access_key.clone(), + session_token: auth.session_token.clone(), + }; + let config = match endpoint { + BedrockEndpoint::Mantle => aws_auth_config(aws), + BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), + }; + let context = AwsAuthContext::load_with_access_keys(config, access_keys) + .await + .map_err(aws_auth_error_to_codex_error)?; + Ok(BedrockAuthMethod::AwsSdkAuth { context }) + } BedrockAuthSource::EnvAwsCredentials | BedrockAuthSource::AwsSdk => { let config = match endpoint { BedrockEndpoint::Mantle => aws_auth_config(aws), @@ -135,7 +162,7 @@ pub(super) async fn resolve_auth_method( pub(super) async fn resolve_provider_auth( source: BedrockAuthSource, - managed_auth: Option<&BedrockApiKeyAuth>, + managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, ) -> Result { @@ -154,7 +181,7 @@ pub(super) async fn resolve_provider_auth( pub(super) async fn resolve_region( source: BedrockAuthSource, - managed_auth: Option<&BedrockApiKeyAuth>, + managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, ) -> Result { @@ -277,6 +304,8 @@ impl AuthProvider for BedrockSigV4AuthProvider { #[cfg(test)] mod tests { use codex_api::AuthProvider; + use codex_login::auth::BedrockAccessKeysAuth; + use codex_login::auth::BedrockApiKeyAuth; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -300,6 +329,13 @@ mod tests { api_key: "managed-bedrock-api-key".to_string(), region: "us-east-1".to_string(), })); + let managed_access_keys = AuthManager::from_auth_for_testing(CodexAuth::BedrockAccessKeys( + BedrockAccessKeysAuth { + access_key_id: "managed-access-key-id".to_string(), + secret_access_key: "managed-secret-access-key".to_string(), + session_token: None, + }, + )); let cases: &[( &ModelProviderInfo, Option<&AuthManager>, @@ -322,6 +358,18 @@ mod tests { &[AWS_BEARER_TOKEN_BEDROCK_ENV_VAR], BedrockAuthSource::ManagedBearerToken, ), + ( + &configured_profile_provider, + Some(managed_access_keys.as_ref()), + &[AWS_BEARER_TOKEN_BEDROCK_ENV_VAR], + BedrockAuthSource::ConfiguredAwsProfile, + ), + ( + &provider, + Some(managed_access_keys.as_ref()), + &[AWS_BEARER_TOKEN_BEDROCK_ENV_VAR], + BedrockAuthSource::ManagedAccessKeys, + ), ( &provider, None, diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 4cf9939fab..60cc69d4c9 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -1,5 +1,5 @@ use codex_aws_auth::AwsAuthConfig; -use codex_login::auth::BedrockApiKeyAuth; +use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; @@ -50,14 +50,14 @@ pub(super) fn base_url(region: &str) -> Result { Ok(format!("https://bedrock-mantle.{region}.api.aws/openai/v1")) } else { Err(CodexErr::Fatal(format!( - "Amazon Bedrock Mantle does not support region `{region}`" + "Amazon Bedrock does not support region `{region}`" ))) } } pub(super) async fn bedrock_mantle_runtime_base_url( source: BedrockAuthSource, - managed_auth: Option<&BedrockApiKeyAuth>, + managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { let region = resolve_region(source, managed_auth, aws, BedrockEndpoint::Mantle).await?; @@ -84,7 +84,7 @@ mod tests { assert_eq!( err.to_string(), - "Fatal error: Amazon Bedrock Mantle does not support region `us-west-1`" + "Fatal error: Amazon Bedrock does not support region `us-west-1`" ); } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index a56aaab1fa..94b2421564 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -15,7 +15,6 @@ use codex_api::SharedAuthProvider; use codex_api::TransportError; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID; use codex_model_provider_info::AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID; @@ -105,21 +104,22 @@ impl AmazonBedrockModelProvider { auth::auth_source(&self.info, self.auth_manager.as_deref(), std::env::var) } - fn managed_auth(&self) -> Option { - if self.auth_source() != auth::BedrockAuthSource::ManagedBearerToken { - return None; - } + fn managed_auth(&self) -> Option { + let source = self.auth_source(); self.auth_manager - .as_ref() - .and_then(|auth_manager| auth_manager.auth_cached()) - .and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, + .as_deref() + .and_then(AuthManager::auth_cached) + .filter(|auth| { + matches!( + (source, auth), + ( + auth::BedrockAuthSource::ManagedBearerToken, + CodexAuth::BedrockApiKey(_) + ) | ( + auth::BedrockAuthSource::ManagedAccessKeys, + CodexAuth::BedrockAccessKeys(_) + ) + ) }) } @@ -137,9 +137,8 @@ impl AmazonBedrockModelProvider { Some(auth_manager) => auth_manager.auth().await, None => None, }, - auth::BedrockAuthSource::ManagedBearerToken => { - self.managed_auth().map(CodexAuth::BedrockApiKey) - } + auth::BedrockAuthSource::ManagedBearerToken + | auth::BedrockAuthSource::ManagedAccessKeys => self.managed_auth(), auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::EnvBearerToken | auth::BedrockAuthSource::EnvAwsCredentials @@ -229,7 +228,8 @@ impl ModelProvider for AmazonBedrockModelProvider { fn auth_manager(&self) -> Option> { match self.auth_source() { auth::BedrockAuthSource::CommandBearerToken - | auth::BedrockAuthSource::ManagedBearerToken => self.auth_manager.clone(), + | auth::BedrockAuthSource::ManagedBearerToken + | auth::BedrockAuthSource::ManagedAccessKeys => self.auth_manager.clone(), auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::EnvBearerToken | auth::BedrockAuthSource::EnvAwsCredentials @@ -274,8 +274,11 @@ impl ModelProvider for AmazonBedrockModelProvider { fn account_state(&self) -> ProviderAccountResult { Ok(ProviderAccountState { account: Some(ProviderAccount::AmazonBedrock { - uses_codex_managed_credentials: self.auth_source() - == auth::BedrockAuthSource::ManagedBearerToken, + uses_codex_managed_credentials: matches!( + self.auth_source(), + auth::BedrockAuthSource::ManagedBearerToken + | auth::BedrockAuthSource::ManagedAccessKeys + ), }), requires_openai_auth: false, }) @@ -329,6 +332,8 @@ mod error_tests; mod tests { use std::num::NonZeroU64; + use codex_login::auth::BedrockAccessKeysAuth; + use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::AwsAuthRefreshConfig; use codex_protocol::config_types::ModelProviderAuthInfo; use http::HeaderValue; @@ -482,42 +487,52 @@ mod tests { ); assert!(!provider.uses_aws_auth_recovery()); - let configured_profile_provider = AmazonBedrockModelProvider::new( - ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { - profile: Some("configured-aws-profile".to_string()), - ..aws - })), - Some(auth_manager), + let access_keys_auth_manager = AuthManager::from_auth_for_testing( + CodexAuth::BedrockAccessKeys(BedrockAccessKeysAuth { + access_key_id: "managed-access-key-id".to_string(), + secret_access_key: "managed-secret-access-key".to_string(), + session_token: None, + }), ); - assert!(configured_profile_provider.auth_manager().is_none()); - assert_eq!(configured_profile_provider.auth().await, None); - assert_eq!( - configured_profile_provider.account_state(), - Ok(ProviderAccountState { - account: Some(ProviderAccount::AmazonBedrock { - uses_codex_managed_credentials: false, - }), - requires_openai_auth: false, - }) - ); - assert_eq!( - configured_profile_provider - .runtime_base_url() - .await - .expect("configured AWS profile region should resolve"), - Some("https://bedrock-mantle.us-west-2.api.aws/openai/v1".to_string()) - ); - assert!( - configured_profile_provider - .api_auth() - .await - .expect("configured AWS profile auth should resolve") - .to_auth_headers() - .get(http::header::AUTHORIZATION) - .is_none() - ); - assert!(configured_profile_provider.uses_aws_auth_recovery()); + for auth_manager in [auth_manager, access_keys_auth_manager] { + let configured_profile_provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: Some("configured-aws-profile".to_string()), + ..aws.clone() + })), + Some(auth_manager), + ); + + assert!(configured_profile_provider.auth_manager().is_none()); + assert_eq!(configured_profile_provider.auth().await, None); + assert_eq!( + configured_profile_provider.account_state(), + Ok(ProviderAccountState { + account: Some(ProviderAccount::AmazonBedrock { + uses_codex_managed_credentials: false, + }), + requires_openai_auth: false, + }) + ); + assert_eq!( + configured_profile_provider + .runtime_base_url() + .await + .expect("configured AWS profile region should resolve"), + Some("https://bedrock-mantle.us-west-2.api.aws/openai/v1".to_string()) + ); + assert!( + configured_profile_provider + .api_auth() + .await + .expect("configured AWS profile auth should resolve") + .to_auth_headers() + .get(http::header::AUTHORIZATION) + .is_none() + ); + assert!(configured_profile_provider.uses_aws_auth_recovery()); + } } #[tokio::test] diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs index 3ab6e6c9c3..7a4d7f410e 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/runtime.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs @@ -1,5 +1,5 @@ use codex_aws_auth::AwsAuthConfig; -use codex_login::auth::BedrockApiKeyAuth; +use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_protocol::error::Result; @@ -24,7 +24,7 @@ pub(super) fn base_url(region: &str) -> String { pub(super) async fn bedrock_runtime_base_url( source: BedrockAuthSource, - managed_auth: Option<&BedrockApiKeyAuth>, + managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, ) -> Result { let region = resolve_region(source, managed_auth, aws, BedrockEndpoint::Runtime).await?; diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 8907ce1578..f4d276465d 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -206,7 +206,10 @@ pub(crate) fn resolve_provider_auth( return Ok(unauthenticated_auth_provider()); } - if matches!(auth, Some(CodexAuth::BedrockApiKey(_))) { + if matches!( + auth, + Some(CodexAuth::BedrockApiKey(_) | CodexAuth::BedrockAccessKeys(_)) + ) { return Err(CodexErr::UnsupportedOperation( BEDROCK_API_KEY_UNSUPPORTED_MESSAGE.to_string(), )); @@ -307,7 +310,9 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { Arc::new(AgentIdentityAuthProvider { auth: auth.clone() }) } CodexAuth::Headers(auth) => Arc::new(HeaderAuthProvider { auth: auth.clone() }), - CodexAuth::BedrockApiKey(_) => unreachable!("{BEDROCK_API_KEY_UNSUPPORTED_MESSAGE}"), + CodexAuth::BedrockApiKey(_) | CodexAuth::BedrockAccessKeys(_) => { + unreachable!("{BEDROCK_API_KEY_UNSUPPORTED_MESSAGE}") + } CodexAuth::ApiKey(_) | CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 44310865a7..6aa6941a3c 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -402,7 +402,7 @@ impl ModelProvider for ConfiguredModelProvider { }) .map(|auth| match &auth { CodexAuth::ApiKey(_) => Ok(ProviderAccount::ApiKey), - CodexAuth::BedrockApiKey(_) => { + CodexAuth::BedrockApiKey(_) | CodexAuth::BedrockAccessKeys(_) => { Err(ProviderAccountError::UnsupportedBedrockApiKeyAuth) } CodexAuth::Chatgpt(_) diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 74ea9389ec..324ca42ceb 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -576,6 +576,7 @@ c2ln", agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; std::fs::create_dir_all(codex_home).expect("codex home should be created"); std::fs::write( diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index e75167d021..5a08881331 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -57,7 +57,9 @@ pub enum TelemetryAuthMode { impl From for TelemetryAuthMode { fn from(mode: AuthMode) -> Self { match mode { - AuthMode::ApiKey | AuthMode::BedrockApiKey => Self::ApiKey, + AuthMode::ApiKey | AuthMode::BedrockApiKey | AuthMode::BedrockAccessKeys => { + Self::ApiKey + } AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::Headers diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs index 67b7e64b88..5f7debc617 100644 --- a/codex-rs/protocol/src/auth.rs +++ b/codex-rs/protocol/src/auth.rs @@ -31,6 +31,10 @@ pub enum AuthMode { #[serde(rename = "bedrockApiKey")] #[strum(serialize = "bedrockApiKey")] BedrockApiKey, + /// Amazon Bedrock AWS access keys managed by Codex. + #[serde(rename = "bedrockAccessKeys")] + #[strum(serialize = "bedrockAccessKeys")] + BedrockAccessKeys, } impl AuthMode { @@ -38,7 +42,11 @@ impl AuthMode { pub fn has_chatgpt_account(self) -> bool { match self { Self::Chatgpt | Self::ChatgptAuthTokens | Self::PersonalAccessToken => true, - Self::ApiKey | Self::Headers | Self::AgentIdentity | Self::BedrockApiKey => false, + Self::ApiKey + | Self::Headers + | Self::AgentIdentity + | Self::BedrockApiKey + | Self::BedrockAccessKeys => false, } } @@ -50,7 +58,7 @@ impl AuthMode { | Self::Headers | Self::AgentIdentity | Self::PersonalAccessToken => true, - Self::ApiKey | Self::BedrockApiKey => false, + Self::ApiKey | Self::BedrockApiKey | Self::BedrockAccessKeys => false, } } } diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index f2de313d03..1b89e841f3 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -1602,7 +1602,9 @@ pub(crate) fn status_account_display_from_auth_mode( email: None, plan: plan_type.map(plan_type_display_name), }), - Some(AuthMode::Headers) | Some(AuthMode::BedrockApiKey) => None, + Some(AuthMode::Headers) + | Some(AuthMode::BedrockApiKey) + | Some(AuthMode::BedrockAccessKeys) => None, None => None, } } diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index a09d8bc9b9..f09ac2263f 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -116,6 +116,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }; save_auth( codex_home, @@ -170,6 +171,7 @@ mod tests { agent_identity: None, personal_access_token: None, bedrock_api_key: None, + bedrock_access_keys: None, }, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index b438c2addb..0bbb899f4a 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -963,6 +963,7 @@ impl AuthModeWidget { ApiAuthMode::AgentIdentity => AuthMode::AgentIdentity, ApiAuthMode::PersonalAccessToken => AuthMode::PersonalAccessToken, ApiAuthMode::BedrockApiKey => AuthMode::BedrockApiKey, + ApiAuthMode::BedrockAccessKeys => AuthMode::BedrockAccessKeys, }) }) .unwrap_or(LoginStatus::NotAuthenticated);