Show credit usage upon user turn completion (#3382)

* Show credit usage upon user turn completion

Author: akhileshb

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

@@ -365,7 +365,7 @@ impl ApiClient {
         let request = self
             .client
             .list_available_models()
-            .set_origin(Some(Origin::KiroCli))
+            .set_origin(Some(Origin::AiEditor))
             .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()));
         let mut paginator = request.into_paginator().send();
 
@@ -458,7 +458,7 @@ impl ApiClient {
     ) -> Result<amzn_codewhisperer_client::operation::get_usage_limits::GetUsageLimitsOutput, ApiClientError> {
         self.client
             .get_usage_limits()
-            .set_origin(Some(amzn_codewhisperer_client::types::Origin::KiroCli))
+            .set_origin(Some(amzn_codewhisperer_client::types::Origin::AiEditor))
             .send()
             .await
             .map_err(ApiClientError::GetUsageLimitsError)
@@ -543,7 +543,7 @@ impl ApiClient {
             match client
                 .send_message()
                 .conversation_state(conversation_state)
-                .set_source(Some(QDeveloperOrigin::KiroCli))
+                .set_source(Some(QDeveloperOrigin::AiEditor))
                 .send()
                 .await
             {

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

@@ -579,6 +579,8 @@ pub enum ChatResponseStream {
     },
     MeteringEvent {
         usage: Option<f64>,
+        unit: Option<String>,
+        unit_plural: Option<String>,
     },
     SupplementaryWebLinksEvent(()),
     ToolUseEvent {
@@ -665,8 +667,17 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
                 cache_write_input_tokens: token_usage.as_ref().and_then(|t| t.cache_write_input_tokens),
             },
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::MeteringEvent(
-                amzn_codewhisperer_streaming_client::types::MeteringEvent { usage, .. },
-            ) => ChatResponseStream::MeteringEvent { usage },
+                amzn_codewhisperer_streaming_client::types::MeteringEvent {
+                    usage,
+                    unit,
+                    unit_plural,
+                    ..
+                },
+            ) => ChatResponseStream::MeteringEvent {
+                usage,
+                unit,
+                unit_plural,
+            },
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_codewhisperer_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -733,8 +744,17 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
                 cache_write_input_tokens: token_usage.as_ref().and_then(|t| t.cache_write_input_tokens),
             },
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::MeteringEvent(
-                amzn_qdeveloper_streaming_client::types::MeteringEvent { usage, .. },
-            ) => ChatResponseStream::MeteringEvent { usage },
+                amzn_qdeveloper_streaming_client::types::MeteringEvent {
+                    usage,
+                    unit,
+                    unit_plural,
+                    ..
+                },
+            ) => ChatResponseStream::MeteringEvent {
+                usage,
+                unit,
+                unit_plural,
+            },
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_qdeveloper_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -940,7 +960,7 @@ impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::User
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
             .set_model_id(value.model_id)
-            .origin(amzn_codewhisperer_streaming_client::types::Origin::KiroCli)
+            .origin(amzn_codewhisperer_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }
@@ -954,7 +974,7 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
             .set_model_id(value.model_id)
-            .origin(amzn_qdeveloper_streaming_client::types::Origin::KiroCli)
+            .origin(amzn_qdeveloper_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }

crates/chat-cli/src/cli/chat/conversation.rs

@@ -102,11 +102,20 @@ pub struct McpServerInfo {
     pub config: CustomToolConfig,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct MeteringUsageInfo {
+    pub value: f64,
+    pub unit: String,
+    pub unit_plural: String,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct UserTurnMetadata {
     continuation_id: String,
     /// [RequestMetadata] about the ongoing operation.
     requests: Vec<RequestMetadata>,
+    /// Individual usage info consumed in this user turn
+    pub usage_info: Vec<MeteringUsageInfo>,
 }
 
 impl Default for UserTurnMetadata {
@@ -121,6 +130,7 @@ impl UserTurnMetadata {
         Self {
             continuation_id: uuid::Uuid::new_v4().to_string(),
             requests: vec![],
+            usage_info: vec![],
         }
     }
 
@@ -151,6 +161,35 @@ impl UserTurnMetadata {
     pub fn last(&self) -> Option<&RequestMetadata> {
         self.requests.last()
     }
+
+    pub fn add_usage(&mut self, value: f64, unit: String, unit_plural: String) {
+        self.usage_info.push(MeteringUsageInfo {
+            value,
+            unit,
+            unit_plural,
+        });
+    }
+
+    pub fn total_usage(&self) -> Vec<MeteringUsageInfo> {
+        let mut totals = std::collections::HashMap::new();
+        for usage in &self.usage_info {
+            let entry = totals.entry(&usage.unit_plural).or_insert(MeteringUsageInfo {
+                value: 0.0,
+                unit: usage.unit.clone(),
+                unit_plural: usage.unit_plural.clone(),
+            });
+            entry.value += usage.value;
+        }
+        totals.into_values().collect()
+    }
+
+    pub fn total_elapsed_time_ms(&self) -> Option<u64> {
+        if let (Some(first), Some(last)) = (self.requests.first(), self.requests.last()) {
+            Some(last.stream_end_timestamp_ms - first.request_start_timestamp_ms)
+        } else {
+            None
+        }
+    }
 }
 
 /// Tracks state related to an ongoing conversation.

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

@@ -8,13 +8,15 @@ use spinners::{
 use crate::api_client::error::ConverseStreamErrorKind;
 use crate::theme::StyledText;
 use crate::util::ui::should_send_structured_message;
+
 pub mod cli;
 mod consts;
 pub mod context;
 mod conversation;
 mod input_source;
 mod message;
 mod parse;
+mod turn_summary;
 use std::path::MAIN_SEPARATOR;
 pub mod checkpoint;
 mod line_tracker;
@@ -823,6 +825,11 @@ impl ChatSession {
         let mut ctrl_c_stream = self.ctrlc_rx.resubscribe();
         let result = match self.inner.take().expect("state must always be Some") {
             ChatState::PromptUser { skip_printing_tools } => {
+                // Show turn complete when we're back to prompting and no tools pending
+                if self.tool_uses.is_empty() {
+                    turn_summary::display_turn_usage_summary(&mut self.stderr, &self.conversation.user_turn_metadata)?;
+                }
+
                 match (self.interactive, self.tool_uses.is_empty()) {
                     (false, true) => {
                         self.inner = Some(ChatState::Exit);
@@ -2806,6 +2813,15 @@ impl ChatSession {
                     trace!("Consumed: {:?}", msg_event);
 
                     match msg_event {
+                        parser::ResponseEvent::MeteringUsage {
+                            value,
+                            unit,
+                            unit_plural,
+                        } => {
+                            self.conversation
+                                .user_turn_metadata
+                                .add_usage(value, unit.clone(), unit_plural.clone());
+                        },
                         parser::ResponseEvent::ToolUseStart { name } => {
                             // We need to flush the buffer here, otherwise text will not be
                             // printed while we are receiving tool use events.

crates/chat-cli/src/cli/chat/parser.rs

@@ -622,6 +622,23 @@ impl ResponseParser {
                         ChatResponseStream::ToolUseEvent { input, .. } => {
                             self.received_response_size += input.as_ref().map(String::len).unwrap_or_default();
                         },
+                        ChatResponseStream::MeteringEvent {
+                            usage,
+                            unit,
+                            unit_plural,
+                        } => {
+                            info!("GenerateAssistanceResponse - MeteringEvent");
+                            if let (Some(value), Some(unit), Some(unit_plural)) = (usage, unit, unit_plural) {
+                                let _ = self
+                                    .event_tx
+                                    .send(Ok(ResponseEvent::MeteringUsage {
+                                        value: *value,
+                                        unit: unit.clone(),
+                                        unit_plural: unit_plural.clone(),
+                                    }))
+                                    .await;
+                            }
+                        },
                         _ => {
                             warn!(?r, "received unexpected event from the response stream");
                         },
@@ -691,6 +708,12 @@ pub enum ResponseEvent {
         /// Metadata for the request stream.
         request_metadata: RequestMetadata,
     },
+    /// Metering usage consumed for this request
+    MeteringUsage {
+        value: f64,
+        unit: String,
+        unit_plural: String,
+    },
 }
 
 /// Metadata about the sent request and associated response stream.

crates/chat-cli/src/cli/chat/turn_summary.rs

@@ -0,0 +1,65 @@
+use std::io::Write;
+
+use crossterm::{
+    execute,
+    style,
+};
+
+use crate::cli::chat::conversation::UserTurnMetadata;
+
+pub fn format_number(value: f64) -> String {
+    format!("{:.2}", (value * 100.0).floor() / 100.0)
+}
+
+pub fn capitalize_first(s: &str) -> String {
+    let mut chars = s.chars();
+    match chars.next() {
+        None => String::new(),
+        Some(first) => {
+            let mut result = first.to_uppercase().collect::<String>();
+            result.push_str(chars.as_str());
+            result
+        },
+    }
+}
+
+pub fn format_elapsed_time(elapsed_ms: u64) -> String {
+    let elapsed_s = elapsed_ms / 1000;
+    if elapsed_s < 60 {
+        format!("{elapsed_s}s")
+    } else {
+        let minutes = elapsed_s / 60;
+        let seconds = elapsed_s % 60;
+        format!("{minutes}m {seconds}s")
+    }
+}
+
+pub fn display_turn_usage_summary<W: Write>(
+    stderr: &mut W,
+    user_turn_metadata: &UserTurnMetadata,
+) -> Result<(), std::io::Error> {
+    let totals = user_turn_metadata.total_usage();
+    let mut parts = Vec::new();
+
+    // Add usage info
+    for usage in totals {
+        let formatted = format_number(usage.value);
+        let capitalized_unit = capitalize_first(&usage.unit_plural);
+        parts.push(format!("{capitalized_unit} used: {formatted}"));
+    }
+
+    // Add elapsed time
+    if let Some(elapsed_ms) = user_turn_metadata.total_elapsed_time_ms() {
+        parts.push(format!("Elapsed time: {}", format_elapsed_time(elapsed_ms)));
+    }
+
+    if !parts.is_empty() {
+        execute!(
+            stderr,
+            style::SetForegroundColor(style::Color::DarkGrey),
+            style::Print(format!("{}\n", parts.join("  "))),
+            style::ResetColor,
+        )?;
+    }
+    Ok(())
+}