feat: Support Qv2 CLI social login (#3060)

* first pass

* reuse 2 pkce func

* add support for always inclusion steering files (#3047)

* redirct index when user cancel auth

* add a tmp page when token exchange does not finished

* add socail profile

* add unified token resolver

* only prompt inviation code when need

* reuse pkce handle call back logic

* telemtry

* fmt

* ci

* add some ut

* rename amazonq to kito

* add more attempt to connect

* comment on ports

* add support for always inclusion steering files (#3047)

* feat: Adds continuation id logic (#3114)

* feat: New continuation id logic

* chore: Centralize metadata

* chore: Moves struct to right place

---------

Co-authored-by: Kenneth S. <[email protected]>

* first pass for login with social option using auth portal

* recover log change, pkce, telemetry

* add which mwinit

* delete telemetry and add portal res enum

* add auth status to portal

* refactor code struc

* user friendly msg

* allow 3 connectoins

* update endpoint for testing

* update endpoint for testing

* point to prod

* clippy

---------

Co-authored-by: Zoe Lin <[email protected]>
Co-authored-by: Kenneth Sanchez V <[email protected]>
Co-authored-by: Kenneth S. <[email protected]>

Author: 4 people

Cargo.lock

@@ -122,6 +122,7 @@ zip.workspace = true
 rmcp.workspace = true
 chat-cli-ui.workspace = true 
 serde_yaml.workspace = true
+urlencoding = "2.1.3"
 
 [target.'cfg(unix)'.dependencies]
 nix.workspace = true

crates/chat-cli/Cargo.toml

@@ -56,7 +56,7 @@ use crate::api_client::model::{
 };
 use crate::api_client::opt_out::OptOutInterceptor;
 use crate::api_client::send_message_output::SendMessageOutput;
-use crate::auth::builder_id::BearerResolver;
+use crate::auth::UnifiedBearerResolver;
 use crate::aws_common::{
     UserAgentOverrideInterceptor,
     app_name,
@@ -147,7 +147,7 @@ impl ApiClient {
                 .http_client(crate::aws_common::http_client::client())
                 .interceptor(OptOutInterceptor::new(database))
                 .interceptor(UserAgentOverrideInterceptor::new())
-                .bearer_token_resolver(BearerResolver)
+                .bearer_token_resolver(UnifiedBearerResolver)
                 .app_name(app_name())
                 .endpoint_url(endpoint.url())
                 .build(),
@@ -208,7 +208,7 @@ impl ApiClient {
                         .interceptor(OptOutInterceptor::new(database))
                         .interceptor(UserAgentOverrideInterceptor::new())
                         .interceptor(DelayTrackingInterceptor::new())
-                        .bearer_token_resolver(BearerResolver)
+                        .bearer_token_resolver(UnifiedBearerResolver)
                         .app_name(app_name())
                         .endpoint_resolver(StaticEndpointResolver::new(endpoint.url().to_string()))
                         .retry_classifier(retry_classifier::QCliRetryClassifier::new())

crates/chat-cli/src/api_client/mod.rs

@@ -543,7 +543,7 @@ pub async fn poll_create_token(
     }
 }
 
-pub async fn is_logged_in(database: &mut Database) -> bool {
+pub async fn is_builder_id_logged_in(database: &mut Database) -> bool {
     // Check for BuilderId if not using Sigv4
     if std::env::var("AMAZON_Q_SIGV4").is_ok_and(|v| !v.is_empty()) {
         debug!("logged in using sigv4 credentials");

crates/chat-cli/src/auth/builder_id.rs

@@ -18,6 +18,8 @@ pub(crate) const SCOPES: &[&str] = &[
 
 pub(crate) const CLIENT_TYPE: &str = "public";
 
+pub const SOCIAL_AUTH_SERVICE_ENDPOINT: &str = "https://prod.us-east-1.auth.desktop.kiro.dev";
+
 // The start URL for public builder ID users
 pub const START_URL: &str = "https://view.awsapps.com/start";
 

crates/chat-cli/src/auth/consts.rs

@@ -3,18 +3,31 @@ mod consts;
 pub mod pkce;
 mod scope;
 
+pub mod portal;
+pub mod social;
+use aws_sdk_ssooidc::config::{
+    ConfigBag,
+    RuntimeComponents,
+};
 use aws_sdk_ssooidc::error::SdkError;
 use aws_sdk_ssooidc::operation::create_token::CreateTokenError;
 use aws_sdk_ssooidc::operation::register_client::RegisterClientError;
 use aws_sdk_ssooidc::operation::start_device_authorization::StartDeviceAuthorizationError;
+use aws_smithy_runtime_api::client::identity::http::Token;
+use aws_smithy_runtime_api::client::identity::{
+    Identity,
+    IdentityFuture,
+    ResolveIdentity,
+};
 pub use builder_id::{
-    is_logged_in,
+    is_builder_id_logged_in,
     logout,
 };
 pub use consts::START_URL;
 use thiserror::Error;
 
 use crate::aws_common::SdkErrorDisplay;
+use crate::database::Database;
 
 #[derive(Debug, Error)]
 pub enum AuthError {
@@ -48,6 +61,12 @@ pub enum AuthError {
     OAuthCustomError(String),
     #[error(transparent)]
     DatabaseError(#[from] crate::database::DatabaseError),
+    #[error(transparent)]
+    Reqwest(#[from] reqwest::Error),
+    #[error("HTTP error: {0}")]
+    HttpStatus(reqwest::StatusCode),
+    #[error("Authentication failed: {0}")]
+    SocialAuthProviderFailure(String),
 }
 
 impl From<aws_sdk_ssooidc::Error> for AuthError {
@@ -73,3 +92,33 @@ impl From<SdkError<StartDeviceAuthorizationError>> for AuthError {
         Self::SdkStartDeviceAuthorization(Box::new(value))
     }
 }
+/// Unified bearer token resolver that tries both social and builder ID tokens
+#[derive(Debug, Clone)]
+pub struct UnifiedBearerResolver;
+
+impl ResolveIdentity for UnifiedBearerResolver {
+    fn resolve_identity<'a>(
+        &'a self,
+        _runtime_components: &'a RuntimeComponents,
+        _config_bag: &'a ConfigBag,
+    ) -> IdentityFuture<'a> {
+        IdentityFuture::new_boxed(Box::pin(async {
+            let database = Database::new().await?;
+
+            if let Ok(Some(token)) = builder_id::BuilderIdToken::load(&database).await {
+                return Ok(Identity::new(
+                    Token::new(token.access_token.0.clone(), Some(token.expires_at.into())),
+                    Some(token.expires_at.into()),
+                ));
+            }
+
+            if let Ok(Some(token)) = social::SocialToken::load(&database).await {
+                return Ok(Identity::new(
+                    Token::new(token.access_token.0.clone(), Some(token.expires_at.into())),
+                    Some(token.expires_at.into()),
+                ));
+            }
+            Err(AuthError::NoToken.into())
+        }))
+    }
+}

crates/chat-cli/src/auth/mod.rs

@@ -458,14 +458,14 @@ impl PkceQueryParams {
 /// Generates a random 43-octet URL safe string according to the RFC recommendation.
 ///
 /// Reference: https://datatracker.ietf.org/doc/html/rfc7636#section-4.1
-fn generate_code_verifier() -> String {
+pub fn generate_code_verifier() -> String {
     URL_SAFE.encode(rand::random::<[u8; 32]>()).replace('=', "")
 }
 
 /// Base64 URL encoded sha256 hash of the code verifier.
 ///
 /// Reference: https://datatracker.ietf.org/doc/html/rfc7636#section-4.2
-fn generate_code_challenge(code_verifier: &str) -> String {
+pub fn generate_code_challenge(code_verifier: &str) -> String {
     use sha2::{
         Digest,
         Sha256,