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:
committed by
GitHub
Unverified
parent
e18fe7a07f
commit
1cd3ad1f49
@@ -25,6 +25,7 @@ pub use crate::default_client::CodexHttpClient;
|
||||
pub use crate::default_client::CodexRequestBuilder;
|
||||
pub use crate::error::StreamError;
|
||||
pub use crate::error::TransportError;
|
||||
pub use crate::request::PreparedRequestBody;
|
||||
pub use crate::request::Request;
|
||||
pub use crate::request::RequestBody;
|
||||
pub use crate::request::RequestCompression;
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use bytes::Bytes;
|
||||
use http::Method;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::HeaderValue;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
@@ -27,6 +28,18 @@ impl RequestBody {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PreparedRequestBody {
|
||||
pub headers: HeaderMap,
|
||||
pub body: Option<Bytes>,
|
||||
}
|
||||
|
||||
impl PreparedRequestBody {
|
||||
pub fn body_bytes(&self) -> Bytes {
|
||||
self.body.clone().unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Request {
|
||||
pub method: Method,
|
||||
@@ -63,6 +76,135 @@ impl Request {
|
||||
self.compression = compression;
|
||||
self
|
||||
}
|
||||
|
||||
/// Convert the request body into the exact bytes that will be sent.
|
||||
///
|
||||
/// Auth schemes such as AWS SigV4 need to sign the final body bytes, including
|
||||
/// compression and content headers. Calling this method does not mutate the
|
||||
/// request.
|
||||
pub fn prepare_body_for_send(&self) -> Result<PreparedRequestBody, String> {
|
||||
let mut headers = self.headers.clone();
|
||||
match self.body.as_ref() {
|
||||
Some(RequestBody::Raw(raw_body)) => {
|
||||
if self.compression != RequestCompression::None {
|
||||
return Err("request compression cannot be used with raw bodies".to_string());
|
||||
}
|
||||
Ok(PreparedRequestBody {
|
||||
headers,
|
||||
body: Some(raw_body.clone()),
|
||||
})
|
||||
}
|
||||
Some(RequestBody::Json(body)) => {
|
||||
let json = serde_json::to_vec(&body).map_err(|err| err.to_string())?;
|
||||
let bytes = if self.compression != RequestCompression::None {
|
||||
if headers.contains_key(http::header::CONTENT_ENCODING) {
|
||||
return Err(
|
||||
"request compression was requested but content-encoding is already set"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
let pre_compression_bytes = json.len();
|
||||
let compression_start = std::time::Instant::now();
|
||||
let (compressed, content_encoding) = match self.compression {
|
||||
RequestCompression::None => unreachable!("guarded by compression != None"),
|
||||
RequestCompression::Zstd => (
|
||||
zstd::stream::encode_all(std::io::Cursor::new(json), 3)
|
||||
.map_err(|err| err.to_string())?,
|
||||
HeaderValue::from_static("zstd"),
|
||||
),
|
||||
};
|
||||
let post_compression_bytes = compressed.len();
|
||||
let compression_duration = compression_start.elapsed();
|
||||
|
||||
headers.insert(http::header::CONTENT_ENCODING, content_encoding);
|
||||
|
||||
tracing::debug!(
|
||||
pre_compression_bytes,
|
||||
post_compression_bytes,
|
||||
compression_duration_ms = compression_duration.as_millis(),
|
||||
"Compressed request body with zstd"
|
||||
);
|
||||
|
||||
compressed
|
||||
} else {
|
||||
json
|
||||
};
|
||||
|
||||
if !headers.contains_key(http::header::CONTENT_TYPE) {
|
||||
headers.insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/json"),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(PreparedRequestBody {
|
||||
headers,
|
||||
body: Some(Bytes::from(bytes)),
|
||||
})
|
||||
}
|
||||
None => Ok(PreparedRequestBody {
|
||||
headers,
|
||||
body: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use http::HeaderValue;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn prepare_body_for_send_serializes_json_and_sets_content_type() {
|
||||
let request = Request::new(Method::POST, "https://example.com/v1/responses".to_string())
|
||||
.with_json(&json!({"model": "test-model"}));
|
||||
|
||||
let prepared = request
|
||||
.prepare_body_for_send()
|
||||
.expect("body should prepare");
|
||||
|
||||
assert_eq!(
|
||||
prepared.body,
|
||||
Some(Bytes::from_static(br#"{"model":"test-model"}"#))
|
||||
);
|
||||
assert_eq!(
|
||||
prepared
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
assert_eq!(
|
||||
request.body,
|
||||
Some(RequestBody::Json(json!({"model": "test-model"})))
|
||||
);
|
||||
assert_eq!(request.compression, RequestCompression::None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_body_for_send_rejects_existing_content_encoding_when_compressing() {
|
||||
let mut request =
|
||||
Request::new(Method::POST, "https://example.com/v1/responses".to_string())
|
||||
.with_json(&json!({"model": "test-model"}))
|
||||
.with_compression(RequestCompression::Zstd);
|
||||
request.headers.insert(
|
||||
http::header::CONTENT_ENCODING,
|
||||
HeaderValue::from_static("gzip"),
|
||||
);
|
||||
|
||||
let err = request
|
||||
.prepare_body_for_send()
|
||||
.expect_err("conflicting content-encoding should fail");
|
||||
|
||||
assert_eq!(
|
||||
err,
|
||||
"request compression was requested but content-encoding is already set"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -3,7 +3,6 @@ use crate::default_client::CodexRequestBuilder;
|
||||
use crate::error::TransportError;
|
||||
use crate::request::Request;
|
||||
use crate::request::RequestBody;
|
||||
use crate::request::RequestCompression;
|
||||
use crate::request::Response;
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
@@ -43,12 +42,14 @@ impl ReqwestTransport {
|
||||
}
|
||||
|
||||
fn build(&self, req: Request) -> Result<CodexRequestBuilder, TransportError> {
|
||||
let prepared = req.prepare_body_for_send().map_err(TransportError::Build)?;
|
||||
|
||||
let Request {
|
||||
method,
|
||||
url,
|
||||
mut headers,
|
||||
body,
|
||||
compression,
|
||||
headers: _,
|
||||
body: _,
|
||||
compression: _,
|
||||
timeout,
|
||||
} = req;
|
||||
|
||||
@@ -61,63 +62,9 @@ impl ReqwestTransport {
|
||||
builder = builder.timeout(timeout);
|
||||
}
|
||||
|
||||
match body {
|
||||
Some(RequestBody::Raw(raw_body)) => {
|
||||
if compression != RequestCompression::None {
|
||||
return Err(TransportError::Build(
|
||||
"request compression cannot be used with raw bodies".to_string(),
|
||||
));
|
||||
}
|
||||
builder = builder.headers(headers).body(raw_body);
|
||||
}
|
||||
Some(RequestBody::Json(body)) => {
|
||||
if compression != RequestCompression::None {
|
||||
if headers.contains_key(http::header::CONTENT_ENCODING) {
|
||||
return Err(TransportError::Build(
|
||||
"request compression was requested but content-encoding is already set"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let json = serde_json::to_vec(&body)
|
||||
.map_err(|err| TransportError::Build(err.to_string()))?;
|
||||
let pre_compression_bytes = json.len();
|
||||
let compression_start = std::time::Instant::now();
|
||||
let (compressed, content_encoding) = match compression {
|
||||
RequestCompression::None => unreachable!("guarded by compression != None"),
|
||||
RequestCompression::Zstd => (
|
||||
zstd::stream::encode_all(std::io::Cursor::new(json), 3)
|
||||
.map_err(|err| TransportError::Build(err.to_string()))?,
|
||||
http::HeaderValue::from_static("zstd"),
|
||||
),
|
||||
};
|
||||
let post_compression_bytes = compressed.len();
|
||||
let compression_duration = compression_start.elapsed();
|
||||
|
||||
// Ensure the server knows to unpack the request body.
|
||||
headers.insert(http::header::CONTENT_ENCODING, content_encoding);
|
||||
if !headers.contains_key(http::header::CONTENT_TYPE) {
|
||||
headers.insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/json"),
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
pre_compression_bytes,
|
||||
post_compression_bytes,
|
||||
compression_duration_ms = compression_duration.as_millis(),
|
||||
"Compressed request body with zstd"
|
||||
);
|
||||
|
||||
builder = builder.headers(headers).body(compressed);
|
||||
} else {
|
||||
builder = builder.headers(headers).json(&body);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
builder = builder.headers(headers);
|
||||
}
|
||||
builder = builder.headers(prepared.headers);
|
||||
if let Some(body) = prepared.body {
|
||||
builder = builder.body(body);
|
||||
}
|
||||
Ok(builder)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user