feat: add delay tracking interceptor for retry notifications (#2607)

Author: kkashilk

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

@@ -0,0 +1,92 @@
+use std::time::{
+    Duration,
+    Instant,
+};
+
+use aws_smithy_runtime_api::box_error::BoxError;
+use aws_smithy_runtime_api::client::interceptors::Intercept;
+use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextRef;
+use aws_smithy_runtime_api::client::retries::RequestAttempts;
+use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
+use aws_smithy_types::config_bag::{
+    ConfigBag,
+    Storable,
+    StoreReplace,
+};
+use crossterm::style::Color;
+use crossterm::{
+    execute,
+    style,
+};
+
+#[derive(Debug, Clone)]
+pub struct DelayTrackingInterceptor {
+    minor_delay_threshold: Duration,
+    major_delay_threshold: Duration,
+}
+
+impl DelayTrackingInterceptor {
+    pub fn new() -> Self {
+        Self {
+            minor_delay_threshold: Duration::from_secs(2),
+            major_delay_threshold: Duration::from_secs(5),
+        }
+    }
+
+    fn print_warning(message: String) {
+        let mut stderr = std::io::stderr();
+        let _ = execute!(
+            stderr,
+            style::SetForegroundColor(Color::Yellow),
+            style::Print("\nWARNING: "),
+            style::SetForegroundColor(Color::Reset),
+            style::Print(message),
+            style::Print("\n")
+        );
+    }
+}
+
+impl Intercept for DelayTrackingInterceptor {
+    fn name(&self) -> &'static str {
+        "DelayTrackingInterceptor"
+    }
+
+    fn read_before_transmit(
+        &self,
+        _: &BeforeTransmitInterceptorContextRef<'_>,
+        _: &RuntimeComponents,
+        cfg: &mut ConfigBag,
+    ) -> Result<(), BoxError> {
+        let attempt_number = cfg.load::<RequestAttempts>().map_or(1, |attempts| attempts.attempts());
+
+        let now = Instant::now();
+
+        if let Some(last_attempt_time) = cfg.load::<LastAttemptTime>() {
+            let delay = now.duration_since(last_attempt_time.0);
+
+            if delay >= self.major_delay_threshold {
+                Self::print_warning(format!(
+                    "Auto Retry #{} delayed by {:.1}s. Service is under heavy load - consider switching models.",
+                    attempt_number,
+                    delay.as_secs_f64()
+                ));
+            } else if delay >= self.minor_delay_threshold {
+                Self::print_warning(format!(
+                    "Auto Retry #{} delayed by {:.1}s due to transient issues.",
+                    attempt_number,
+                    delay.as_secs_f64()
+                ));
+            }
+        }
+
+        cfg.interceptor_state().store_put(LastAttemptTime(Instant::now()));
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone)]
+struct LastAttemptTime(Instant);
+
+impl Storable for LastAttemptTime {
+    type Storer = StoreReplace<Self>;
+}

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

@@ -1,5 +1,6 @@
 mod credentials;
 pub mod customization;
+mod delay_interceptor;
 mod endpoints;
 mod error;
 pub mod model;
@@ -41,6 +42,7 @@ use tracing::{
 };
 
 use crate::api_client::credentials::CredentialsChain;
+use crate::api_client::delay_interceptor::DelayTrackingInterceptor;
 use crate::api_client::model::{
     ChatResponseStream,
     ConversationState,
@@ -163,6 +165,7 @@ impl ApiClient {
                     .http_client(crate::aws_common::http_client::client())
                     .interceptor(OptOutInterceptor::new(database))
                     .interceptor(UserAgentOverrideInterceptor::new())
+                    .interceptor(DelayTrackingInterceptor::new())
                     .app_name(app_name())
                     .endpoint_url(endpoint.url())
                     .retry_classifier(retry_classifier::QCliRetryClassifier::new())
@@ -176,6 +179,7 @@ impl ApiClient {
                         .http_client(crate::aws_common::http_client::client())
                         .interceptor(OptOutInterceptor::new(database))
                         .interceptor(UserAgentOverrideInterceptor::new())
+                        .interceptor(DelayTrackingInterceptor::new())
                         .bearer_token_resolver(BearerResolver)
                         .app_name(app_name())
                         .endpoint_url(endpoint.url())