claude-api 0.5.3

//! AWS Bedrock support: a [`RequestSigner`] that signs HTTP requests
//! with sigv4.
//!
//! `claude-api` is the Anthropic API client; this module exposes the
//! signing primitive only. Users who want to talk to Bedrock Anthropic
//! models still need to:
//!
//! 1. Set `Client::builder().base_url("https://bedrock-runtime.{region}.amazonaws.com")`,
//! 2. Use Bedrock's URL shape (`/model/{model_id}/invoke`),
//! 3. Inject `anthropic_version: "bedrock-2023-05-31"` and remove the
//!    `model` field from the body (Bedrock takes the model in the URL).
//!
//! The typed namespace handles still use Anthropic's URL shape; Bedrock
//! model IDs and the `anthropic_version` header must be supplied by the
//! caller. See `examples/bedrock.rs` for a complete working sketch.
//!
//! Gated on the `bedrock` feature.
//!
//! # Set up the client
//!
//! ```no_run
//! use std::sync::Arc;
//! use claude_api::{Client, bedrock::{BedrockCredentials, BedrockSigner}};
//! # fn run() -> Result<(), claude_api::Error> {
//! let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".into());
//! let creds = BedrockCredentials::from_env()
//!     .expect("AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY must be set");
//! let client = Client::builder()
//!     .signer(Arc::new(BedrockSigner::new(creds, &region)))
//!     .base_url(format!("https://bedrock-runtime.{region}.amazonaws.com"))
//!     .build()?;
//! # Ok(())
//! # }
//! ```

#![cfg(feature = "bedrock")]
#![cfg_attr(docsrs, doc(cfg(feature = "bedrock")))]

use std::str::FromStr;
use std::time::SystemTime;

use aws_credential_types::Credentials;
use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign};
use aws_sigv4::sign::v4::SigningParams;

use crate::auth::{RequestSigner, SignerResult};

/// AWS access credentials. Carries the same fields as
/// [`aws_credential_types::Credentials`] but is owned, `Clone`, and
/// kept opaque -- the underlying secret is moved into a fresh
/// `Credentials` per call so this type is safe to share via `Arc`.
#[derive(Clone)]
pub struct BedrockCredentials {
    access_key_id: String,
    secret_access_key: String,
    session_token: Option<String>,
}

impl BedrockCredentials {
    /// Construct from access key ID + secret access key. Use
    /// [`Self::with_session_token`] for STS-issued temporary credentials.
    #[must_use]
    pub fn new(access_key_id: impl Into<String>, secret_access_key: impl Into<String>) -> Self {
        Self {
            access_key_id: access_key_id.into(),
            secret_access_key: secret_access_key.into(),
            session_token: None,
        }
    }

    /// Attach an STS session token. Required when credentials come from
    /// `AssumeRole`, IMDS, or any other temporary-credential source.
    #[must_use]
    pub fn with_session_token(mut self, token: impl Into<String>) -> Self {
        self.session_token = Some(token.into());
        self
    }

    /// Read credentials from the standard AWS environment variables
    /// (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, optional
    /// `AWS_SESSION_TOKEN`). Returns `None` if either of the required
    /// pair is missing.
    #[must_use]
    pub fn from_env() -> Option<Self> {
        let access = std::env::var("AWS_ACCESS_KEY_ID").ok()?;
        let secret = std::env::var("AWS_SECRET_ACCESS_KEY").ok()?;
        let mut creds = Self::new(access, secret);
        if let Ok(token) = std::env::var("AWS_SESSION_TOKEN") {
            creds = creds.with_session_token(token);
        }
        Some(creds)
    }

    fn to_aws(&self) -> Credentials {
        Credentials::new(
            self.access_key_id.clone(),
            self.secret_access_key.clone(),
            self.session_token.clone(),
            None,
            "claude-api-bedrock-signer",
        )
    }
}

impl std::fmt::Debug for BedrockCredentials {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BedrockCredentials")
            .field("access_key_id", &"<redacted>")
            .field("secret_access_key", &"<redacted>")
            .field(
                "session_token",
                &self.session_token.as_ref().map(|_| "<redacted>"),
            )
            .finish()
    }
}

/// AWS sigv4 signer for the `bedrock` service.
///
/// Install on a [`Client`](crate::Client) via
/// [`ClientBuilder::signer`](crate::ClientBuilder::signer).
#[derive(Debug, Clone)]
pub struct BedrockSigner {
    credentials: BedrockCredentials,
    region: String,
    /// AWS service name used in the canonical request. Defaults to
    /// `"bedrock"`. Override only if you're targeting a sister service
    /// that reuses this signer.
    service: String,
}

impl BedrockSigner {
    /// Build a signer for `service = "bedrock"` in the given region.
    #[must_use]
    pub fn new(credentials: BedrockCredentials, region: impl Into<String>) -> Self {
        Self {
            credentials,
            region: region.into(),
            service: "bedrock".into(),
        }
    }

    /// Override the service name used in the canonical request.
    #[must_use]
    pub fn with_service(mut self, service: impl Into<String>) -> Self {
        self.service = service.into();
        self
    }
}

impl RequestSigner for BedrockSigner {
    fn sign(&self, request: &mut reqwest::Request) -> SignerResult {
        let identity = self.credentials.to_aws().into();

        let settings = SigningSettings::default();
        let params: aws_sigv4::http_request::SigningParams = SigningParams::builder()
            .identity(&identity)
            .region(&self.region)
            .name(&self.service)
            .time(SystemTime::now())
            .settings(settings)
            .build()
            .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?
            .into();

        // SignableRequest wants borrowed (&str, &str) header pairs.
        // Collect into a Vec<(String, String)> first to satisfy the
        // lifetime, then borrow.
        let header_strings: Vec<(String, String)> = request
            .headers()
            .iter()
            .filter_map(|(name, value)| {
                value
                    .to_str()
                    .ok()
                    .map(|v| (name.as_str().to_owned(), v.to_owned()))
            })
            .collect();
        let headers_iter = header_strings.iter().map(|(k, v)| (k.as_str(), v.as_str()));

        let body_bytes = request.body().and_then(|b| b.as_bytes()).unwrap_or(&[]);
        let signable_body = SignableBody::Bytes(body_bytes);

        let url = request.url().as_str().to_owned();
        let signable =
            SignableRequest::new(request.method().as_str(), &url, headers_iter, signable_body)
                .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;

        let signing_output = sign(signable, &params)
            .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
        let (instructions, _signature) = signing_output.into_parts();

        for (name, value) in instructions.headers() {
            let header_name = http::HeaderName::from_str(name)?;
            let header_value = http::HeaderValue::from_str(value)?;
            request.headers_mut().insert(header_name, header_value);
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_request() -> reqwest::Request {
        let client = reqwest::Client::new();
        client
            .post("https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-sonnet-20240620-v1:0/invoke")
            .body(r#"{"messages":[{"role":"user","content":"hi"}]}"#)
            .build()
            .unwrap()
    }

    fn fixed_signer() -> BedrockSigner {
        BedrockSigner::new(
            BedrockCredentials::new("AKIDEXAMPLE", "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"),
            "us-east-1",
        )
    }

    #[test]
    fn bedrock_signer_adds_authorization_header() {
        let signer = fixed_signer();
        let mut req = make_request();
        signer.sign(&mut req).expect("sign succeeds");

        let auth = req
            .headers()
            .get("authorization")
            .expect("Authorization header set by signer");
        let auth_str = auth.to_str().expect("Authorization is ASCII");
        assert!(
            auth_str.starts_with("AWS4-HMAC-SHA256 "),
            "expected sigv4 algorithm prefix: {auth_str}"
        );
        assert!(
            auth_str.contains("Credential=AKIDEXAMPLE/"),
            "expected access key in credential scope: {auth_str}"
        );
        assert!(
            auth_str.contains("/us-east-1/bedrock/aws4_request"),
            "expected region+service in credential scope: {auth_str}"
        );
        assert!(
            auth_str.contains("SignedHeaders="),
            "expected SignedHeaders component: {auth_str}"
        );
        assert!(
            auth_str.contains("Signature="),
            "expected Signature component: {auth_str}"
        );
    }

    #[test]
    fn bedrock_signer_adds_x_amz_date_header() {
        let signer = fixed_signer();
        let mut req = make_request();
        signer.sign(&mut req).unwrap();
        let date = req
            .headers()
            .get("x-amz-date")
            .expect("X-Amz-Date header set by signer");
        let s = date.to_str().unwrap();
        // ISO 8601 basic format: YYYYMMDDTHHMMSSZ -> 16 chars.
        assert_eq!(s.len(), 16, "date should be 16-char ISO 8601 basic: {s}");
        assert!(s.ends_with('Z'), "date should be UTC: {s}");
    }

    #[test]
    fn bedrock_signer_includes_session_token_when_present() {
        let creds =
            BedrockCredentials::new("AKID", "SECRET").with_session_token("session-token-value");
        let signer = BedrockSigner::new(creds, "us-west-2");
        let mut req = make_request();
        signer.sign(&mut req).unwrap();
        let token = req
            .headers()
            .get("x-amz-security-token")
            .expect("X-Amz-Security-Token forwarded by signer");
        assert_eq!(token.to_str().unwrap(), "session-token-value");
    }

    #[test]
    fn bedrock_credentials_redact_secret_in_debug() {
        let creds =
            BedrockCredentials::new("AKID", "VERY-SECRET").with_session_token("ALSO-SECRET");
        let dbg = format!("{creds:?}");
        assert!(!dbg.contains("VERY-SECRET"), "{dbg}");
        assert!(!dbg.contains("ALSO-SECRET"), "{dbg}");
        assert!(dbg.contains("redacted"), "{dbg}");
    }

    #[test]
    fn from_env_returns_none_when_missing() {
        // We can't reliably scrub the env in this test (other tests may
        // rely on AWS_*). Instead just verify the function compiles
        // and returns *some* outcome. Coverage of the path is via the
        // signer integration tests below.
        let _: Option<BedrockCredentials> = BedrockCredentials::from_env();
    }

    #[test]
    fn signer_default_service_name_is_bedrock() {
        let signer = fixed_signer();
        assert_eq!(signer.service, "bedrock");
    }

    #[test]
    fn signer_with_service_override() {
        let signer = fixed_signer().with_service("bedrock-runtime");
        assert_eq!(signer.service, "bedrock-runtime");
    }
}