feat: Support Qv2 CLI social login

Cargo.toml

@@ -130,6 +130,7 @@ winreg = "0.55.0"
 schemars = "1.0.4"
 jsonschema = "0.30.0"
 zip = "2.2.0"
+serde_yaml = "0.9"
 rmcp = { version = "0.7.0", features = ["client", "transport-sse-client-reqwest", "reqwest", "transport-streamable-http-client-reqwest", "transport-child-process", "tower", "auth"] }
 
 [workspace.lints.rust]

crates/chat-cli/Cargo.toml

@@ -119,6 +119,8 @@ schemars.workspace = true
 jsonschema.workspace = true
 zip.workspace = true
 rmcp.workspace = true
+serde_yaml.workspace = true
+urlencoding = "2.1.3"
 
 [target.'cfg(unix)'.dependencies]
 nix.workspace = true

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

@@ -50,7 +50,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,
@@ -122,7 +122,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(),
@@ -183,7 +183,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_url(endpoint.url())
                         .retry_classifier(retry_classifier::QCliRetryClassifier::new())

crates/chat-cli/src/auth/consts.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/index.html

@@ -156,23 +156,51 @@ <h4>Request denied</h4>
                 const params = new URLSearchParams(window.location.search)
 
                 const error = params.get('error')
+                const success = params.get('success');
                 if (error) {
                     showErrorMessage(error)
                     return
                 }
 
-                const productName = 'Amazon Q for CLI'
-                document.getElementById(
-                    'approvalMessage'
-                ).innerText = `${productName} have been given requested permissions`
-                document.getElementById(
-                    'footerText'
-                ).innerText = `You can close this window and start using ${productName}`
-
-                function showErrorMessage(errorText) {
-                    document.getElementById('approved-auth').classList.add('hidden')
-                    document.getElementById('denied-auth').classList.remove('hidden')
-                    document.getElementById('errorMessage').innerText = errorText
+
+                // Only show success message if explicitly marked as successful
+                // This prevents showing success when token exchange might still fail
+                if (success === 'true') {
+                    const productName = 'KIRO CLI';
+                    document.getElementById('approvalMessage').innerText = 
+                        `${productName} has been granted given requested permissions.`;
+                    document.getElementById('footerText').innerText = 
+                        `You can close this window and start using ${productName}`;
+                } else {
+                    // Show neutral "processing" message
+                    document.getElementById('approvalMessage').innerText = 
+                        'Processing authentication...';
+                    document.getElementById('footerText').innerText = 
+                        'You can close this window and check the terminal for results.';
+                }
+
+                function showErrorMessage(errorParam) {
+                    document.getElementById('approved-auth').classList.add('hidden');
+                    document.getElementById('denied-auth').classList.remove('hidden');
+
+                    // Map error codes to user-friendly messages
+                    let errorMessage = '';
+                    let retryMessage = 'You can close this window and try again.';
+
+                    const decodedError = decodeURIComponent(errorParam);
+
+                    if (decodedError === 'access_denied' || decodedError.includes('denied')) {
+                        errorMessage = 'You cancelled the authentication or denied access. Kiro CLI requires the requested permissions to function properly.';
+                    } else if (decodedError.includes('invalid') || decodedError.includes('invitation')) {
+                        errorMessage = 'Invalid access code. Please check your invitation code and try again.';
+                    } else if (decodedError.includes('timeout')) {
+                        errorMessage = 'Authentication timed out. Please try again.';
+                    } else {
+                        errorMessage = `Authentication error: ${decodedError}`;
+                    }
+
+                    document.getElementById('errorMessage').innerText = errorMessage;
+                    document.getElementById('retryText').innerText = retryMessage;
                 }
             }
         </script>

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

@@ -3,10 +3,21 @@ mod consts;
 pub mod pkce;
 mod scope;
 
+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,
     logout,
@@ -15,6 +26,7 @@ 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 +60,19 @@ 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),
+    // Social auth specific errors
+    #[error(
+        "Authentication failed: The identity provider denied access. Please ensure you grant all required permissions."
+    )]
+    SocialAuthProviderDeniedAccess,
+    #[error("Authentication failed: The identity provider reported an error: {0}")]
+    SocialAuthProviderFailure(String),
+    #[error("Invalid access code. Please check your invitation code and try again.")]
+    SocialInvalidInvitationCode,
 }
 
 impl From<aws_sdk_ssooidc::Error> for AuthError {
@@ -73,3 +98,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/pkce.rs

@@ -21,6 +21,7 @@
 
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::Arc;
 use std::time::Duration;
 
 pub use aws_sdk_ssooidc::client::Client;
@@ -282,42 +283,83 @@ impl PkceRegistration {
         Ok(())
     }
 
-    async fn recv_code(listener: TcpListener, expected_state: String) -> Result<String, AuthError> {
+    async fn recv_code(listener: tokio::net::TcpListener, expected_state: String) -> Result<String, AuthError> {
+        Self::recv_code_with_extra_accepts(listener, expected_state, 0).await
+    }
+
+    // NEW: extended version that accepts N extra connections to serve /index.html again
+    pub async fn recv_code_with_extra_accepts(
+        listener: tokio::net::TcpListener,
+        expected_state: String,
+        extra_accepts: usize,
+    ) -> Result<String, AuthError> {
+        // Channel for delivering (code, state) from any served connection
         let (code_tx, mut code_rx) = tokio::sync::mpsc::channel::<Result<(String, String), AuthError>>(1);
-        let (stream, _) = listener.accept().await?;
-        let stream = TokioIo::new(stream); // Wrapper to implement Hyper IO traits for Tokio types.
+        let code_tx_arc = Arc::new(code_tx);
+
+        // Host used by the service to 302 -> /index.html
         let host = listener.local_addr()?.to_string();
-        tokio::spawn(async move {
-            if let Err(err) = http1::Builder::new()
-                .serve_connection(stream, PkceHttpService {
-                    code_tx: std::sync::Arc::new(code_tx),
-                    host,
-                })
-                .await
-            {
-                error!(?err, "Error occurred serving the connection");
+
+        // Serve multiple connections to handle both /oauth/callback and /index.html
+        // Keep the behavior as close as possible to the previous version:
+        // first connection is typically the callback, but we allow extra accepts
+        // (callback + potential error redirect + index.html).
+        let max_accepts = 1 + extra_accepts;
+        let server_handle = tokio::spawn({
+            let code_tx_arc = code_tx_arc.clone();
+            let host = host.clone();
+            async move {
+                for _ in 0..max_accepts {
+                    if let Ok((stream, _)) = listener.accept().await {
+                        let stream = TokioIo::new(stream);
+                        let service = PkceHttpService {
+                            code_tx: code_tx_arc.clone(),
+                            host: host.clone(),
+                        };
+                        tokio::spawn(async move {
+                            if let Err(err) = http1::Builder::new().serve_connection(stream, service).await {
+                                debug!(?err, "Error serving connection");
+                            }
+                        });
+                    } else {
+                        // Accept error; break to avoid a tight loop
+                        break;
+                    }
+                }
             }
         });
-        match code_rx.recv().await {
-            Some(Ok((code, state))) => {
-                debug!(code = "<redacted>", state, "Received code and state");
-                if state != expected_state {
-                    return Err(AuthError::OAuthStateMismatch {
-                        actual: state,
-                        expected: expected_state,
-                    });
+
+        // Wait for authorization code (or an error) with the same timeout semantics
+        let result = tokio::select! {
+            msg = code_rx.recv() => {
+                match msg {
+                    Some(Ok((code, state))) => {
+                        debug!(code = "<redacted>", state, "Received code and state");
+                        if state != expected_state {
+                            Err(AuthError::OAuthStateMismatch {
+                                actual: state,
+                                expected: expected_state,
+                            })
+                        } else {
+                            Ok(code)
+                        }
+                    }
+                    Some(Err(e)) => Err(e),
+                    None => Err(AuthError::OAuthMissingCode),
                 }
-                // Give time for the user to be redirected to index.html.
-                tokio::time::sleep(Duration::from_millis(200)).await;
-                Ok(code)
-            },
-            Some(Err(err)) => {
-                // Give time for the user to be redirected to index.html.
-                tokio::time::sleep(Duration::from_millis(200)).await;
-                Err(err)
-            },
-            None => Err(AuthError::OAuthMissingCode),
-        }
+            }
+            _ = tokio::time::sleep(DEFAULT_AUTHORIZATION_TIMEOUT) => {
+                Err(AuthError::OAuthTimeout)
+            }
+        };
+
+        // Small grace period so the browser can load /index.html after the 302
+        tokio::time::sleep(Duration::from_millis(200)).await;
+
+        // Let the server task finish (best-effort)
+        let _ = tokio::time::timeout(Duration::from_secs(2), server_handle).await;
+
+        result
     }
 }
 
@@ -459,14 +501,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,