feat: Enabling Sigv4 Auth Support for Q (#207)

Manojtuguru · Manoj Tuguru · omursalo · web-flow · commit 42b16763d058 · 2025-06-20T16:06:39.000-07:00
* Initial Commit for Q CLI changes * Add SigV4 authentication support to login command * Adding with the Auth Strategy * Remove references for AuthStrategy * Added back the Q Dev client and Q_USE_SENDMESSAGE env variable * fix: Add origin to SendMessage request and change env variable for Sigv4 * code-cleanup:removed sigv4 auth references * chore: cleanup AuthStrategy definitions * fix/chore: skip login when using sigv4 and cleanup unused code * Fix/mcp disable (#257) * Update crates/chat-cli/src/api_client/clients/streaming_client.rs Co-authored-by: Chay Nabors <github@chaynabors.com> * Update crates/chat-cli/src/api_client/clients/streaming_client.rs Co-authored-by: Chay Nabors <github@chaynabors.com> * Update crates/chat-cli/src/auth/builder_id.rs Co-authored-by: Chay Nabors <github@chaynabors.com> * Chore:Addressing comments and Clean up * Refactored the builder ID condition check to AMAZON_Q_SIGV4 * removing unused imports * fixes --------- Co-authored-by: Manoj Tuguru <manojtug@amazon.com> Co-authored-by: Olena Mursalova <omursalo@amazon.com> Co-authored-by: Josh Rutkowski <54906829+joshrutkowski@users.noreply.github.com> Co-authored-by: Chay Nabors <github@chaynabors.com> Co-authored-by: Chay Nabors <nabochay@amazon.com>
diff --git a/crates/chat-cli/src/api_client/clients/shared.rs b/crates/chat-cli/src/api_client/clients/shared.rs
@@ -8,7 +8,11 @@ use aws_credential_types::provider::ProvideCredentials;
 use aws_types::SdkConfig;
 use aws_types::sdk_config::StalledStreamProtectionConfig;
 
-use crate::api_client::Endpoint;
+use crate::api_client::credentials::CredentialsChain;
+use crate::api_client::{
+    ApiClientError,
+    Endpoint,
+};
 use crate::aws_common::behavior_version;
 use crate::database::Database;
 use crate::database::settings::Setting;
@@ -55,3 +59,13 @@ pub async fn bearer_sdk_config(database: &Database, endpoint: &Endpoint) -> SdkC
     let credentials = Credentials::new("xxx", "xxx", None, None, "xxx");
     base_sdk_config(database, endpoint.region().clone(), credentials).await
 }
+
+pub async fn sigv4_sdk_config(database: &Database, endpoint: &Endpoint) -> Result<SdkConfig, ApiClientError> {
+    let credentials_chain = CredentialsChain::new().await;
+
+    if let Err(err) = credentials_chain.provide_credentials().await {
+        return Err(ApiClientError::Credentials(err));
+    };
+
+    Ok(base_sdk_config(database, endpoint.region().clone(), credentials_chain).await)
+}
diff --git a/crates/chat-cli/src/api_client/clients/streaming_client.rs b/crates/chat-cli/src/api_client/clients/streaming_client.rs
@@ -4,6 +4,8 @@ use std::sync::{
 };
 
 use amzn_codewhisperer_streaming_client::Client as CodewhispererStreamingClient;
+use amzn_qdeveloper_streaming_client::Client as QDeveloperStreamingClient;
+use amzn_qdeveloper_streaming_client::types::Origin;
 use aws_types::request_id::RequestId;
 use tracing::{
     debug,
@@ -12,6 +14,7 @@ use tracing::{
 
 use super::shared::{
     bearer_sdk_config,
+    sigv4_sdk_config,
     stalled_stream_protection_config,
 };
 use crate::api_client::interceptor::opt_out::OptOutInterceptor;
@@ -40,12 +43,14 @@ mod inner {
     };
 
     use amzn_codewhisperer_streaming_client::Client as CodewhispererStreamingClient;
+    use amzn_qdeveloper_streaming_client::Client as QDeveloperStreamingClient;
 
     use crate::api_client::model::ChatResponseStream;
 
     #[derive(Clone, Debug)]
     pub enum Inner {
         Codewhisperer(CodewhispererStreamingClient),
+        QDeveloper(QDeveloperStreamingClient),
         Mock(Arc<Mutex<std::vec::IntoIter<Vec<ChatResponseStream>>>>),
     }
 }
@@ -58,7 +63,13 @@ pub struct StreamingClient {
 
 impl StreamingClient {
     pub async fn new(database: &mut Database) -> Result<Self, ApiClientError> {
-        Self::new_codewhisperer_client(database, &Endpoint::load_codewhisperer(database)).await
+        // If SIGV4_AUTH_ENABLED is true, use Q developer client
+        if std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) {
+            Self::new_qdeveloper_client(database, &Endpoint::load_q(database)).await
+        } else {
+            // Default to CodeWhisperer client
+            Self::new_codewhisperer_client(database, &Endpoint::load_codewhisperer(database)).await
+        }
     }
 
     pub fn mock(events: Vec<Vec<ChatResponseStream>>) -> Self {
@@ -96,6 +107,25 @@ impl StreamingClient {
         Ok(Self { inner, profile })
     }
 
+    // Add SigV4 client creation method
+    pub async fn new_qdeveloper_client(database: &Database, endpoint: &Endpoint) -> Result<Self, ApiClientError> {
+        let conf_builder: amzn_qdeveloper_streaming_client::config::Builder =
+            (&sigv4_sdk_config(database, endpoint).await?).into();
+        let conf = conf_builder
+            .http_client(crate::aws_common::http_client::client())
+            .interceptor(OptOutInterceptor::new(database))
+            .interceptor(UserAgentOverrideInterceptor::new())
+            .app_name(app_name())
+            .endpoint_url(endpoint.url())
+            .stalled_stream_protection(stalled_stream_protection_config())
+            .build();
+        let client = QDeveloperStreamingClient::from_conf(conf);
+        Ok(Self {
+            inner: inner::Inner::QDeveloper(client),
+            profile: None,
+        })
+    }
+
     pub async fn send_message(&self, conversation: ConversationState) -> Result<SendMessageOutput, ApiClientError> {
         debug!("Sending conversation: {:#?}", conversation);
         let ConversationState {
@@ -180,6 +210,51 @@ impl StreamingClient {
                     },
                 }
             },
+            inner::Inner::QDeveloper(client) => {
+                let conversation_state = amzn_qdeveloper_streaming_client::types::ConversationState::builder()
+                    .set_conversation_id(conversation_id)
+                    .current_message(amzn_qdeveloper_streaming_client::types::ChatMessage::UserInputMessage(
+                        user_input_message.into(),
+                    ))
+                    .chat_trigger_type(amzn_qdeveloper_streaming_client::types::ChatTriggerType::Manual)
+                    .set_history(
+                        history
+                            .map(|v| v.into_iter().map(|i| i.try_into()).collect::<Result<Vec<_>, _>>())
+                            .transpose()?,
+                    )
+                    .build()
+                    .expect("building conversation_state should not fail");
+
+                let response = client
+                    .send_message()
+                    .conversation_state(conversation_state)
+                    .set_source(Some(Origin::from("CLI")))
+                    .send()
+                    .await;
+
+                match response {
+                    Ok(resp) => Ok(SendMessageOutput::QDeveloper(resp)),
+                    Err(e) => {
+                        let status_code = e.raw_response().map(|res| res.status().as_u16());
+                        let is_quota_breach = e.raw_response().is_some_and(|resp| resp.status().as_u16() == 429);
+                        let is_context_window_overflow = e.as_service_error().is_some_and(|err| {
+                            matches!(err, err if err.meta().code() == Some("ValidationException")
+                                && err.meta().message() == Some("Input is too long."))
+                        });
+
+                        if is_quota_breach {
+                            Err(ApiClientError::QuotaBreach {
+                                message: "quota has reached its limit",
+                                status_code,
+                            })
+                        } else if is_context_window_overflow {
+                            Err(ApiClientError::ContextWindowOverflow { status_code })
+                        } else {
+                            Err(e.into())
+                        }
+                    },
+                }
+            },
             inner::Inner::Mock(events) => {
                 let mut new_events = events.lock().unwrap().next().unwrap_or_default().clone();
                 new_events.reverse();
@@ -194,13 +269,15 @@ pub enum SendMessageOutput {
     Codewhisperer(
         amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseOutput,
     ),
+    QDeveloper(amzn_qdeveloper_streaming_client::operation::send_message::SendMessageOutput),
     Mock(Vec<ChatResponseStream>),
 }
 
 impl SendMessageOutput {
     pub fn request_id(&self) -> Option<&str> {
         match self {
             SendMessageOutput::Codewhisperer(output) => output.request_id(),
+            SendMessageOutput::QDeveloper(output) => output.request_id(),
             SendMessageOutput::Mock(_) => None,
         }
     }
@@ -212,6 +289,7 @@ impl SendMessageOutput {
                 .recv()
                 .await?
                 .map(|s| s.into())),
+            SendMessageOutput::QDeveloper(output) => Ok(output.send_message_response.recv().await?.map(|s| s.into())),
             SendMessageOutput::Mock(vec) => Ok(vec.pop()),
         }
     }
@@ -221,6 +299,7 @@ impl RequestId for SendMessageOutput {
     fn request_id(&self) -> Option<&str> {
         match self {
             SendMessageOutput::Codewhisperer(output) => output.request_id(),
+            SendMessageOutput::QDeveloper(output) => output.request_id(),
             SendMessageOutput::Mock(_) => Some("<mock-request-id>"),
         }
     }
@@ -242,6 +321,7 @@ mod tests {
 
         let _ = StreamingClient::new(&mut database).await;
         let _ = StreamingClient::new_codewhisperer_client(&mut database, &endpoint).await;
+        let _ = StreamingClient::new_qdeveloper_client(&database, &endpoint).await;
     }
 
     #[tokio::test]
diff --git a/crates/chat-cli/src/api_client/consts.rs b/crates/chat-cli/src/api_client/consts.rs
@@ -3,6 +3,8 @@ use aws_config::Region;
 // Endpoint constants
 pub const PROD_CODEWHISPERER_ENDPOINT_URL: &str = "https://codewhisperer.us-east-1.amazonaws.com";
 pub const PROD_CODEWHISPERER_ENDPOINT_REGION: Region = Region::from_static("us-east-1");
+pub const PROD_Q_ENDPOINT_URL: &str = "https://q.us-east-1.amazonaws.com";
+pub const PROD_Q_ENDPOINT_REGION: Region = Region::from_static("us-east-1");
 
 // FRA endpoint constants
 pub const PROD_CODEWHISPERER_FRA_ENDPOINT_URL: &str = "https://q.eu-central-1.amazonaws.com/";
diff --git a/crates/chat-cli/src/api_client/credentials/mod.rs b/crates/chat-cli/src/api_client/credentials/mod.rs
@@ -0,0 +1,80 @@
+use aws_config::default_provider::region::DefaultRegionChain;
+use aws_config::ecs::EcsCredentialsProvider;
+use aws_config::environment::credentials::EnvironmentVariableCredentialsProvider;
+use aws_config::imds::credentials::ImdsCredentialsProvider;
+use aws_config::meta::credentials::CredentialsProviderChain;
+use aws_config::profile::ProfileFileCredentialsProvider;
+use aws_config::provider_config::ProviderConfig;
+use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider;
+use aws_credential_types::Credentials;
+use aws_credential_types::provider::{
+    self,
+    ProvideCredentials,
+    future,
+};
+use tracing::Instrument;
+
+#[derive(Debug)]
+pub struct CredentialsChain {
+    provider_chain: CredentialsProviderChain,
+}
+
+impl CredentialsChain {
+    /// Based on code the code for
+    /// [aws_config::default_provider::credentials::DefaultCredentialsChain]
+    pub async fn new() -> Self {
+        let region = DefaultRegionChain::builder().build().region().await;
+        let config = ProviderConfig::default().with_region(region.clone());
+
+        let env_provider = EnvironmentVariableCredentialsProvider::new();
+        let profile_provider = ProfileFileCredentialsProvider::builder().configure(&config).build();
+        let web_identity_token_provider = WebIdentityTokenCredentialsProvider::builder()
+            .configure(&config)
+            .build();
+        let imds_provider = ImdsCredentialsProvider::builder().configure(&config).build();
+        let ecs_provider = EcsCredentialsProvider::builder().configure(&config).build();
+
+        let mut provider_chain = CredentialsProviderChain::first_try("Environment", env_provider);
+
+        provider_chain = provider_chain
+            .or_else("Profile", profile_provider)
+            .or_else("WebIdentityToken", web_identity_token_provider)
+            .or_else("EcsContainer", ecs_provider)
+            .or_else("Ec2InstanceMetadata", imds_provider);
+
+        CredentialsChain { provider_chain }
+    }
+
+    async fn credentials(&self) -> provider::Result {
+        self.provider_chain
+            .provide_credentials()
+            .instrument(tracing::debug_span!("provide_credentials", provider = %"default_chain"))
+            .await
+    }
+}
+
+impl ProvideCredentials for CredentialsChain {
+    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(self.credentials())
+    }
+
+    fn fallback_on_interrupt(&self) -> Option<Credentials> {
+        self.provider_chain.fallback_on_interrupt()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_credentials_chain() {
+        let credentials_chain = CredentialsChain::new().await;
+        let credentials_res = credentials_chain.provide_credentials().await;
+        let fallback_on_interrupt_res = credentials_chain.fallback_on_interrupt();
+        println!("credentials_res: {credentials_res:?}, fallback_on_interrupt_res: {fallback_on_interrupt_res:?}");
+    }
+}
diff --git a/crates/chat-cli/src/api_client/endpoints.rs b/crates/chat-cli/src/api_client/endpoints.rs
@@ -9,6 +9,8 @@ use crate::api_client::consts::{
     PROD_CODEWHISPERER_ENDPOINT_URL,
     PROD_CODEWHISPERER_FRA_ENDPOINT_REGION,
     PROD_CODEWHISPERER_FRA_ENDPOINT_URL,
+    PROD_Q_ENDPOINT_REGION,
+    PROD_Q_ENDPOINT_URL,
 };
 use crate::database::Database;
 use crate::database::settings::Setting;
@@ -28,6 +30,10 @@ impl Endpoint {
         url: Cow::Borrowed(PROD_CODEWHISPERER_ENDPOINT_URL),
         region: PROD_CODEWHISPERER_ENDPOINT_REGION,
     };
+    pub const Q_ENDPOINT: Self = Self {
+        url: Cow::Borrowed(PROD_Q_ENDPOINT_URL),
+        region: PROD_Q_ENDPOINT_REGION,
+    };
 
     pub fn load_codewhisperer(database: &Database) -> Self {
         let (endpoint, region) = if let Some(Value::Object(o)) = database.settings.get(Setting::ApiCodeWhispererService)
@@ -63,6 +69,24 @@ impl Endpoint {
         }
     }
 
+    pub fn load_q(database: &Database) -> Self {
+        match database.settings.get(Setting::ApiQService) {
+            Some(Value::Object(o)) => {
+                let endpoint = o.get("endpoint").and_then(|v| v.as_str());
+                let region = o.get("region").and_then(|v| v.as_str());
+
+                match (endpoint, region) {
+                    (Some(endpoint), Some(region)) => Self {
+                        url: endpoint.to_owned().into(),
+                        region: Region::new(region.to_owned()),
+                    },
+                    _ => Endpoint::Q_ENDPOINT,
+                }
+            },
+            _ => Endpoint::Q_ENDPOINT,
+        }
+    }
+
     pub(crate) fn url(&self) -> &str {
         &self.url
     }
@@ -82,6 +106,7 @@ mod tests {
     async fn test_endpoints() {
         let database = Database::new().await.unwrap();
         let _ = Endpoint::load_codewhisperer(&database);
+        let _ = Endpoint::load_q(&database);
 
         let prod = &Endpoint::DEFAULT_ENDPOINT;
         Url::parse(prod.url()).unwrap();
diff --git a/crates/chat-cli/src/api_client/error.rs b/crates/chat-cli/src/api_client/error.rs
@@ -9,6 +9,7 @@ use amzn_consolas_client::operation::generate_recommendations::GenerateRecommend
 use amzn_consolas_client::operation::list_customizations::ListCustomizationsError;
 use amzn_qdeveloper_streaming_client::operation::send_message::SendMessageError as QDeveloperSendMessageError;
 use amzn_qdeveloper_streaming_client::types::error::ChatResponseStreamError as QDeveloperChatResponseStreamError;
+use aws_credential_types::provider::error::CredentialsError;
 use aws_sdk_ssooidc::error::ProvideErrorMetadata;
 use aws_smithy_runtime_api::client::orchestrator::HttpResponse;
 pub use aws_smithy_runtime_api::client::result::SdkError;
@@ -88,6 +89,10 @@ pub enum ApiClientError {
         request_id: Option<String>,
         status_code: Option<u16>,
     },
+
+    // Credential errors
+    #[error("failed to load credentials: {}", .0)]
+    Credentials(CredentialsError),
 }
 
 impl ApiClientError {
@@ -110,6 +115,7 @@ impl ApiClientError {
             ApiClientError::AuthError(_) => None,
             ApiClientError::ModelOverloadedError { status_code, .. } => *status_code,
             ApiClientError::MonthlyLimitReached { status_code } => *status_code,
+            ApiClientError::Credentials(_e) => None,
         }
     }
 }
@@ -134,6 +140,7 @@ impl ReasonCode for ApiClientError {
             ApiClientError::AuthError(_) => "AuthError".to_string(),
             ApiClientError::ModelOverloadedError { .. } => "ModelOverloadedError".to_string(),
             ApiClientError::MonthlyLimitReached { .. } => "MonthlyLimitReached".to_string(),
+            ApiClientError::Credentials(_) => "CredentialsError".to_string(),
         }
     }
 }
@@ -171,6 +178,7 @@ mod tests {
 
     fn all_errors() -> Vec<ApiClientError> {
         vec![
+            ApiClientError::Credentials(CredentialsError::unhandled("<unhandled>")),
             ApiClientError::GenerateCompletions(SdkError::service_error(
                 GenerateCompletionsError::unhandled("<unhandled>"),
                 response(),
diff --git a/crates/chat-cli/src/api_client/mod.rs b/crates/chat-cli/src/api_client/mod.rs
@@ -1,16 +1,14 @@
 pub mod clients;
 pub(crate) mod consts;
+pub(crate) mod credentials;
 pub mod customization;
 mod endpoints;
 mod error;
 pub(crate) mod interceptor;
 pub mod model;
 pub mod profile;
 
-pub use clients::{
-    Client,
-    StreamingClient,
-};
+pub use clients::Client;
 pub use endpoints::Endpoint;
 pub use error::ApiClientError;
 pub use profile::list_available_profiles;
diff --git a/crates/chat-cli/src/auth/builder_id.rs b/crates/chat-cli/src/auth/builder_id.rs
diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs
diff --git a/crates/chat-cli/src/database/settings.rs b/crates/chat-cli/src/database/settings.rs

Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ use amzn_consolas_client::operation::generate_recommendations::GenerateRecommend`
`9`	`9`	`use amzn_consolas_client::operation::list_customizations::ListCustomizationsError;`
`10`	`10`	`use amzn_qdeveloper_streaming_client::operation::send_message::SendMessageError as QDeveloperSendMessageError;`
`11`	`11`	`use amzn_qdeveloper_streaming_client::types::error::ChatResponseStreamError as QDeveloperChatResponseStreamError;`
	`12`	`+use aws_credential_types::provider::error::CredentialsError;`
`12`	`13`	`use aws_sdk_ssooidc::error::ProvideErrorMetadata;`
`13`	`14`	`use aws_smithy_runtime_api::client::orchestrator::HttpResponse;`
`14`	`15`	`pub use aws_smithy_runtime_api::client::result::SdkError;`
`@@ -88,6 +89,10 @@ pub enum ApiClientError {`
`88`	`89`	`request_id: Option<String>,`
`89`	`90`	`status_code: Option<u16>,`
`90`	`91`	`},`
	`92`	`+`
	`93`	`+ // Credential errors`
	`94`	`+ #[error("failed to load credentials: {}", .0)]`
	`95`	`+ Credentials(CredentialsError),`
`91`	`96`	`}`
`92`	`97`
`93`	`98`	`impl ApiClientError {`
`@@ -110,6 +115,7 @@ impl ApiClientError {`
`110`	`115`	`ApiClientError::AuthError(_) => None,`
`111`	`116`	`ApiClientError::ModelOverloadedError { status_code, .. } => *status_code,`
`112`	`117`	`ApiClientError::MonthlyLimitReached { status_code } => *status_code,`
	`118`	`+ ApiClientError::Credentials(_e) => None,`
`113`	`119`	`}`
`114`	`120`	`}`
`115`	`121`	`}`
`@@ -134,6 +140,7 @@ impl ReasonCode for ApiClientError {`
`134`	`140`	`ApiClientError::AuthError(_) => "AuthError".to_string(),`
`135`	`141`	`ApiClientError::ModelOverloadedError { .. } => "ModelOverloadedError".to_string(),`
`136`	`142`	`ApiClientError::MonthlyLimitReached { .. } => "MonthlyLimitReached".to_string(),`
	`143`	`+ ApiClientError::Credentials(_) => "CredentialsError".to_string(),`
`137`	`144`	`}`
`138`	`145`	`}`
`139`	`146`	`}`
`@@ -171,6 +178,7 @@ mod tests {`
`171`	`178`
`172`	`179`	`fn all_errors() -> Vec<ApiClientError> {`
`173`	`180`	`vec![`
	`181`	`+ ApiClientError::Credentials(CredentialsError::unhandled("<unhandled>")),`
`174`	`182`	`ApiClientError::GenerateCompletions(SdkError::service_error(`
`175`	`183`	`GenerateCompletionsError::unhandled("<unhandled>"),`
`176`	`184`	`response(),`