mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: add AWS SigV4 auth for OpenAI-compatible model providers (#17820)
## Summary Add first-class Amazon Bedrock Mantle provider support so Codex can keep using its existing Responses API transport with OpenAI-compatible AWS-hosted endpoints such as AOA/Mantle. This is needed for the AWS launch path, where provider traffic should authenticate with AWS credentials instead of OpenAI bearer credentials. Requests are authenticated immediately before transport send, so SigV4 signs the final method, URL, headers, and body bytes that `reqwest` will send. ## What Changed - Added a new `codex-aws-auth` crate for loading AWS SDK config, resolving credentials, and signing finalized HTTP requests with AWS SigV4. - Added a built-in `amazon-bedrock` provider that targets Bedrock Mantle Responses endpoints, defaults to `us-east-1`, supports region/profile overrides, disables WebSockets, and does not require OpenAI auth. - Added Amazon Bedrock auth resolution in `codex-model-provider`: prefer `AWS_BEARER_TOKEN_BEDROCK` when set, otherwise use AWS SDK credentials and SigV4 signing. - Added `AuthProvider::apply_auth` and `Request::prepare_body_for_send` so request-signing providers can sign the exact outbound request after JSON serialization/compression. - Determine the region by taking the `aws.region` config first (required for bearer token codepath), and fallback to SDK default region. ## Testing Amazon Bedrock Mantle Responses paths: - Built the local Codex binary with `cargo build`. - Verified the custom proxy-backed `aws` provider using `env_key = "AWS_BEARER_TOKEN_BEDROCK"` streamed raw `responses` output with `response.output_text.delta`, `response.completed`, and `mantle-env-ok`. - Verified a full `codex exec --profile aws` turn returned `mantle-env-ok`. - Confirmed the custom provider used the bearer env var, not AWS profile auth: bogus `AWS_PROFILE` still passed, empty env var failed locally, and malformed env var reached Mantle and failed with `401 invalid_api_key`. - Verified built-in `amazon-bedrock` with `AWS_BEARER_TOKEN_BEDROCK` set passed despite bogus AWS profiles, returning `amazon-bedrock-env-ok`. - Verified built-in `amazon-bedrock` SDK/SigV4 auth passed with `AWS_BEARER_TOKEN_BEDROCK` unset and temporary AWS session env credentials, returning `amazon-bedrock-sdk-env-ok`.
This commit is contained in:
@@ -0,0 +1,232 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_api::AuthError;
|
||||
use codex_api::AuthProvider;
|
||||
use codex_api::SharedAuthProvider;
|
||||
use codex_aws_auth::AwsAuthConfig;
|
||||
use codex_aws_auth::AwsAuthContext;
|
||||
use codex_aws_auth::AwsAuthError;
|
||||
use codex_aws_auth::AwsRequestToSign;
|
||||
use codex_client::Request;
|
||||
use codex_client::RequestBody;
|
||||
use codex_client::RequestCompression;
|
||||
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::error::Result;
|
||||
use http::HeaderMap;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
use crate::BearerAuthProvider;
|
||||
|
||||
use super::mantle::aws_auth_config;
|
||||
use super::mantle::region_from_config;
|
||||
|
||||
const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
const LEGACY_SESSION_ID_HEADER: &str = "session_id";
|
||||
|
||||
enum BedrockAuthMethod {
|
||||
EnvBearerToken {
|
||||
token: String,
|
||||
region: String,
|
||||
},
|
||||
AwsSdkAuth {
|
||||
config: AwsAuthConfig,
|
||||
context: AwsAuthContext,
|
||||
},
|
||||
}
|
||||
|
||||
async fn resolve_auth_method(aws: &ModelProviderAwsAuthInfo) -> Result<BedrockAuthMethod> {
|
||||
if let Some(token) = bearer_token_from_env() {
|
||||
let region = bearer_token_region_from_config(aws)?;
|
||||
return Ok(BedrockAuthMethod::EnvBearerToken { token, region });
|
||||
}
|
||||
|
||||
let config = aws_auth_config(aws);
|
||||
let context = AwsAuthContext::load(config.clone())
|
||||
.await
|
||||
.map_err(aws_auth_error_to_codex_error)?;
|
||||
Ok(BedrockAuthMethod::AwsSdkAuth { config, context })
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_provider_auth(
|
||||
aws: &ModelProviderAwsAuthInfo,
|
||||
) -> Result<SharedAuthProvider> {
|
||||
match resolve_auth_method(aws).await? {
|
||||
BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider {
|
||||
token: Some(token),
|
||||
account_id: None,
|
||||
is_fedramp_account: false,
|
||||
})),
|
||||
BedrockAuthMethod::AwsSdkAuth { config, context } => Ok(Arc::new(
|
||||
BedrockMantleSigV4AuthProvider::with_context(config, context),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn resolve_region(aws: &ModelProviderAwsAuthInfo) -> Result<String> {
|
||||
match resolve_auth_method(aws).await? {
|
||||
BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region),
|
||||
BedrockAuthMethod::AwsSdkAuth { context, .. } => Ok(context.region().to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn bearer_token_from_env() -> Option<String> {
|
||||
std::env::var(AWS_BEARER_TOKEN_BEDROCK_ENV_VAR)
|
||||
.ok()
|
||||
.map(|token| token.trim().to_string())
|
||||
.filter(|token| !token.is_empty())
|
||||
}
|
||||
|
||||
fn bearer_token_region_from_config(aws: &ModelProviderAwsAuthInfo) -> Result<String> {
|
||||
region_from_config(aws).ok_or_else(|| {
|
||||
CodexErr::Fatal(
|
||||
"Amazon Bedrock bearer token auth requires \
|
||||
`model_providers.amazon-bedrock.aws.region`"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn aws_auth_error_to_codex_error(error: AwsAuthError) -> CodexErr {
|
||||
CodexErr::Fatal(format!("failed to resolve Amazon Bedrock auth: {error}"))
|
||||
}
|
||||
|
||||
fn aws_auth_error_to_auth_error(error: AwsAuthError) -> AuthError {
|
||||
if error.is_retryable() {
|
||||
AuthError::Transient(error.to_string())
|
||||
} else {
|
||||
AuthError::Build(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_headers_not_preserved_by_bedrock_mantle(headers: &mut HeaderMap) {
|
||||
// The Bedrock Mantle front door does not preserve this legacy OpenAI header
|
||||
// for SigV4 verification. Signing it makes the richer Codex agent request
|
||||
// fail even though raw Responses requests work.
|
||||
headers.remove(LEGACY_SESSION_ID_HEADER);
|
||||
}
|
||||
|
||||
/// AWS SigV4 auth provider for Bedrock Mantle OpenAI-compatible requests.
|
||||
#[derive(Debug)]
|
||||
struct BedrockMantleSigV4AuthProvider {
|
||||
config: AwsAuthConfig,
|
||||
context: OnceCell<AwsAuthContext>,
|
||||
}
|
||||
|
||||
impl BedrockMantleSigV4AuthProvider {
|
||||
fn with_context(config: AwsAuthConfig, context: AwsAuthContext) -> Self {
|
||||
let cell = OnceCell::new();
|
||||
let _ = cell.set(context);
|
||||
Self {
|
||||
config,
|
||||
context: cell,
|
||||
}
|
||||
}
|
||||
|
||||
async fn context(&self) -> std::result::Result<&AwsAuthContext, AuthError> {
|
||||
self.context
|
||||
.get_or_try_init(|| AwsAuthContext::load(self.config.clone()))
|
||||
.await
|
||||
.map_err(aws_auth_error_to_auth_error)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl AuthProvider for BedrockMantleSigV4AuthProvider {
|
||||
fn add_auth_headers(&self, _headers: &mut HeaderMap) {}
|
||||
|
||||
async fn apply_auth(&self, request: Request) -> std::result::Result<Request, AuthError> {
|
||||
let mut request = request;
|
||||
remove_headers_not_preserved_by_bedrock_mantle(&mut request.headers);
|
||||
let prepared = request.prepare_body_for_send().map_err(AuthError::Build)?;
|
||||
let context = self.context().await?;
|
||||
let signed = context
|
||||
.sign(AwsRequestToSign {
|
||||
method: request.method.clone(),
|
||||
url: request.url.clone(),
|
||||
headers: prepared.headers.clone(),
|
||||
body: prepared.body_bytes(),
|
||||
})
|
||||
.await
|
||||
.map_err(aws_auth_error_to_auth_error)?;
|
||||
|
||||
request.url = signed.url;
|
||||
request.headers = signed.headers;
|
||||
request.body = prepared.body.map(RequestBody::Raw);
|
||||
request.compression = RequestCompression::None;
|
||||
Ok(request)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use codex_api::AuthProvider;
|
||||
use http::HeaderValue;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn bedrock_bearer_auth_uses_configured_region_and_header() {
|
||||
let token = "bedrock-api-key-test".to_string();
|
||||
let region = bearer_token_region_from_config(&ModelProviderAwsAuthInfo {
|
||||
profile: None,
|
||||
region: Some(" us-west-2 ".to_string()),
|
||||
})
|
||||
.expect("configured region should resolve");
|
||||
let provider = BearerAuthProvider {
|
||||
token: Some(token),
|
||||
account_id: None,
|
||||
is_fedramp_account: false,
|
||||
};
|
||||
let mut headers = http::HeaderMap::new();
|
||||
|
||||
provider.add_auth_headers(&mut headers);
|
||||
|
||||
assert_eq!(region, "us-west-2");
|
||||
assert!(
|
||||
headers
|
||||
.get(http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.is_some_and(|value| value.starts_with("Bearer bedrock-api-key-"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_bearer_auth_rejects_missing_configured_region() {
|
||||
let err = bearer_token_region_from_config(&ModelProviderAwsAuthInfo {
|
||||
profile: None,
|
||||
region: None,
|
||||
})
|
||||
.expect_err("missing region should fail");
|
||||
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
"Fatal error: Amazon Bedrock bearer token auth requires \
|
||||
`model_providers.amazon-bedrock.aws.region`"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_mantle_sigv4_strips_legacy_session_id_header() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
LEGACY_SESSION_ID_HEADER,
|
||||
HeaderValue::from_static("019dae79-15c3-70c3-8736-3219b8602b37"),
|
||||
);
|
||||
headers.insert(
|
||||
"x-client-request-id",
|
||||
HeaderValue::from_static("request-id"),
|
||||
);
|
||||
|
||||
remove_headers_not_preserved_by_bedrock_mantle(&mut headers);
|
||||
|
||||
assert!(!headers.contains_key(LEGACY_SESSION_ID_HEADER));
|
||||
assert_eq!(
|
||||
headers
|
||||
.get("x-client-request-id")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("request-id")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
use codex_aws_auth::AwsAuthConfig;
|
||||
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::error::Result;
|
||||
|
||||
const BEDROCK_MANTLE_SERVICE_NAME: &str = "bedrock-mantle";
|
||||
const BEDROCK_MANTLE_SUPPORTED_REGIONS: [&str; 12] = [
|
||||
"us-east-2",
|
||||
"us-east-1",
|
||||
"us-west-2",
|
||||
"ap-southeast-3",
|
||||
"ap-south-1",
|
||||
"ap-northeast-1",
|
||||
"eu-central-1",
|
||||
"eu-west-1",
|
||||
"eu-west-2",
|
||||
"eu-south-1",
|
||||
"eu-north-1",
|
||||
"sa-east-1",
|
||||
];
|
||||
|
||||
pub(super) fn aws_auth_config(aws: &ModelProviderAwsAuthInfo) -> AwsAuthConfig {
|
||||
AwsAuthConfig {
|
||||
profile: aws.profile.clone(),
|
||||
region: region_from_config(aws),
|
||||
service: BEDROCK_MANTLE_SERVICE_NAME.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn region_from_config(aws: &ModelProviderAwsAuthInfo) -> Option<String> {
|
||||
aws.region
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|region| !region.is_empty())
|
||||
.map(str::to_string)
|
||||
}
|
||||
|
||||
pub(super) fn base_url(region: &str) -> Result<String> {
|
||||
if BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) {
|
||||
Ok(format!("https://bedrock-mantle.{region}.api.aws/v1"))
|
||||
} else {
|
||||
Err(CodexErr::Fatal(format!(
|
||||
"Amazon Bedrock Mantle does not support region `{region}`"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn base_url_uses_region_endpoint() {
|
||||
assert_eq!(
|
||||
base_url("ap-northeast-1").expect("supported region"),
|
||||
"https://bedrock-mantle.ap-northeast-1.api.aws/v1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn base_url_rejects_unsupported_region() {
|
||||
let err = base_url("us-west-1").expect_err("unsupported region");
|
||||
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
"Fatal error: Amazon Bedrock Mantle does not support region `us-west-1`"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aws_auth_config_uses_profile_and_mantle_service() {
|
||||
assert_eq!(
|
||||
aws_auth_config(&ModelProviderAwsAuthInfo {
|
||||
profile: Some("codex-bedrock".to_string()),
|
||||
region: None,
|
||||
}),
|
||||
AwsAuthConfig {
|
||||
profile: Some("codex-bedrock".to_string()),
|
||||
region: None,
|
||||
service: "bedrock-mantle".to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aws_auth_config_uses_configured_region() {
|
||||
assert_eq!(
|
||||
aws_auth_config(&ModelProviderAwsAuthInfo {
|
||||
profile: None,
|
||||
region: Some(" us-west-2 ".to_string()),
|
||||
}),
|
||||
AwsAuthConfig {
|
||||
profile: None,
|
||||
region: Some("us-west-2".to_string()),
|
||||
service: "bedrock-mantle".to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
mod auth;
|
||||
mod mantle;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_api::Provider;
|
||||
use codex_api::SharedAuthProvider;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
||||
use codex_model_provider_info::ModelProviderInfo;
|
||||
use codex_protocol::error::Result;
|
||||
|
||||
use crate::provider::ModelProvider;
|
||||
use auth::resolve_provider_auth;
|
||||
use auth::resolve_region;
|
||||
use mantle::base_url;
|
||||
|
||||
/// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint.
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct AmazonBedrockModelProvider {
|
||||
pub(crate) info: ModelProviderInfo,
|
||||
pub(crate) aws: ModelProviderAwsAuthInfo,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ModelProvider for AmazonBedrockModelProvider {
|
||||
fn info(&self) -> &ModelProviderInfo {
|
||||
&self.info
|
||||
}
|
||||
|
||||
fn auth_manager(&self) -> Option<Arc<AuthManager>> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn auth(&self) -> Option<CodexAuth> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn api_provider(&self) -> Result<Provider> {
|
||||
let region = resolve_region(&self.aws).await?;
|
||||
let mut api_provider_info = self.info.clone();
|
||||
api_provider_info.base_url = Some(base_url(®ion)?);
|
||||
api_provider_info.to_api_provider(/*auth_mode*/ None)
|
||||
}
|
||||
|
||||
async fn api_auth(&self) -> Result<SharedAuthProvider> {
|
||||
resolve_provider_auth(&self.aws).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn api_provider_for_bedrock_bearer_token_uses_configured_region_endpoint() {
|
||||
let region = "eu-central-1";
|
||||
let mut api_provider_info =
|
||||
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None);
|
||||
api_provider_info.base_url = Some(base_url(region).expect("supported region"));
|
||||
let api_provider = api_provider_info
|
||||
.to_api_provider(/*auth_mode*/ None)
|
||||
.expect("api provider should build");
|
||||
|
||||
assert_eq!(
|
||||
api_provider.base_url,
|
||||
"https://bedrock-mantle.eu-central-1.api.aws/v1"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod amazon_bedrock;
|
||||
mod auth;
|
||||
mod bearer_auth_provider;
|
||||
mod provider;
|
||||
|
||||
@@ -5,8 +5,10 @@ use codex_api::Provider;
|
||||
use codex_api::SharedAuthProvider;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
||||
use codex_model_provider_info::ModelProviderInfo;
|
||||
|
||||
use crate::amazon_bedrock::AmazonBedrockModelProvider;
|
||||
use crate::auth::auth_manager_for_provider;
|
||||
use crate::auth::resolve_provider_auth;
|
||||
|
||||
@@ -53,6 +55,20 @@ pub fn create_model_provider(
|
||||
provider_info: ModelProviderInfo,
|
||||
auth_manager: Option<Arc<AuthManager>>,
|
||||
) -> SharedModelProvider {
|
||||
if provider_info.is_amazon_bedrock() {
|
||||
let aws = provider_info
|
||||
.aws
|
||||
.clone()
|
||||
.unwrap_or(ModelProviderAwsAuthInfo {
|
||||
profile: None,
|
||||
region: None,
|
||||
});
|
||||
return Arc::new(AmazonBedrockModelProvider {
|
||||
info: provider_info,
|
||||
aws,
|
||||
});
|
||||
}
|
||||
|
||||
let auth_manager = auth_manager_for_provider(auth_manager, &provider_info);
|
||||
Arc::new(ConfiguredModelProvider {
|
||||
info: provider_info,
|
||||
@@ -89,6 +105,7 @@ impl ModelProvider for ConfiguredModelProvider {
|
||||
mod tests {
|
||||
use std::num::NonZeroU64;
|
||||
|
||||
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
||||
use codex_protocol::config_types::ModelProviderAuthInfo;
|
||||
|
||||
use super::*;
|
||||
@@ -123,4 +140,19 @@ mod tests {
|
||||
|
||||
assert!(auth_manager.has_external_auth());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_model_provider_does_not_use_openai_auth_manager_for_amazon_bedrock_provider() {
|
||||
let provider = create_model_provider(
|
||||
ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo {
|
||||
profile: Some("codex-bedrock".to_string()),
|
||||
region: None,
|
||||
})),
|
||||
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
|
||||
"openai-api-key",
|
||||
))),
|
||||
);
|
||||
|
||||
assert!(provider.auth_manager().is_none());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user