feat(telemetry): add client application, aws service and operation name (#436)

mukeshsahay · web-flow · commit 48647e1b5272 · 2025-07-17T09:47:52.000-07:00
diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs
@@ -1599,6 +1599,18 @@ impl ChatSession {
             let mut tool_telemetry = self.tool_use_telemetry_events.entry(tool.id.clone());
             tool_telemetry = tool_telemetry.and_modify(|ev| ev.is_accepted = true);
 
+            // Extract AWS service name and operation name if available
+            if let Some(additional_info) = tool.tool.get_additional_info() {
+                if let Some(aws_service_name) = additional_info.get("aws_service_name").and_then(|v| v.as_str()) {
+                    tool_telemetry =
+                        tool_telemetry.and_modify(|ev| ev.aws_service_name = Some(aws_service_name.to_string()));
+                }
+                if let Some(aws_operation_name) = additional_info.get("aws_operation_name").and_then(|v| v.as_str()) {
+                    tool_telemetry =
+                        tool_telemetry.and_modify(|ev| ev.aws_operation_name = Some(aws_operation_name.to_string()));
+                }
+            }
+
             let tool_start = std::time::Instant::now();
             let invoke_result = tool.tool.invoke(os, &mut self.stdout).await;
 
@@ -2179,7 +2191,7 @@ impl ChatSession {
             }
             .map(|v| v.to_string());
 
-            os.telemetry.send_tool_use_suggested(event).ok();
+            os.telemetry.send_tool_use_suggested(&os.database, event).await.ok();
         }
     }
 
diff --git a/crates/chat-cli/src/cli/chat/tool_manager.rs b/crates/chat-cli/src/cli/chat/tool_manager.rs
@@ -86,6 +86,7 @@ use crate::cli::chat::tools::{
     ToolOrigin,
     ToolSpec,
 };
+use crate::database::Database;
 use crate::database::settings::Setting;
 use crate::mcp_client::{
     JsonRpcResponse,
@@ -415,6 +416,7 @@ impl ToolManagerBuilder {
         let pending_clone = pending.clone();
         let (mut msg_rx, messenger_builder) = ServerMessengerBuilder::new(20);
         let telemetry_clone = os.telemetry.clone();
+        let database_clone = os.database.clone();
         let notify = Arc::new(Notify::new());
         let notify_weak = Arc::downgrade(&notify);
         let load_record = Arc::new(Mutex::new(HashMap::<String, Vec<LoadingRecord>>::new()));
@@ -452,7 +454,9 @@ impl ToolManagerBuilder {
                                     &mut sanitized_mapping,
                                     &regex,
                                     &telemetry_clone,
-                                );
+                                    &database_clone,
+                                )
+                                .await;
                                 if let Some(sender) = &loading_status_sender_clone {
                                     // Anomalies here are not considered fatal, thus we shall give
                                     // warnings.
@@ -591,7 +595,8 @@ impl ToolManagerBuilder {
                 Err(e) => {
                     error!("Error initializing mcp client for server {}: {:?}", name, &e);
                     os.telemetry
-                        .send_mcp_server_init(conversation_id.clone(), Some(e.to_string()), 0)
+                        .send_mcp_server_init(&os.database, conversation_id.clone(), Some(e.to_string()), 0)
+                        .await
                         .ok();
                     let _ = messenger.send_tools_list_result(Err(e)).await;
                 },
@@ -1266,13 +1271,14 @@ impl ToolManager {
 }
 
 #[inline]
-fn process_tool_specs(
+async fn process_tool_specs(
     conversation_id: &str,
     server_name: &str,
     specs: &mut Vec<ToolSpec>,
     tn_map: &mut HashMap<String, String>,
     regex: &Regex,
     telemetry: &TelemetryThread,
+    database: &Database,
 ) -> eyre::Result<()> {
     // Each mcp server might have multiple tools.
     // To avoid naming conflicts we are going to namespace it.
@@ -1322,7 +1328,10 @@ fn process_tool_specs(
     specs.retain(|spec| !matches!(spec.tool_origin, ToolOrigin::Native));
     // Send server load success metric datum
     let conversation_id = conversation_id.to_string();
-    let _ = telemetry.send_mcp_server_init(conversation_id, None, number_of_tools);
+    telemetry
+        .send_mcp_server_init(database, conversation_id, None, number_of_tools)
+        .await
+        .ok();
     // Tool name translation. This is beyond of the scope of what is
     // considered a "server load". Reasoning being:
     // - Failures here are not related to server load
diff --git a/crates/chat-cli/src/cli/chat/tools/mod.rs b/crates/chat-cli/src/cli/chat/tools/mod.rs
@@ -129,6 +129,15 @@ impl Tool {
             Tool::Thinking(think) => think.validate(os).await,
         }
     }
+
+    /// Returns additional information about the tool if available
+    pub fn get_additional_info(&self) -> Option<serde_json::Value> {
+        match self {
+            Tool::UseAws(use_aws) => Some(use_aws.get_additional_info()),
+            // Add other tool types here as they implement get_additional_info()
+            _ => None,
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
diff --git a/crates/chat-cli/src/cli/chat/tools/use_aws.rs b/crates/chat-cli/src/cli/chat/tools/use_aws.rs
@@ -173,6 +173,13 @@ impl UseAws {
         Ok(())
     }
 
+    pub fn get_additional_info(&self) -> serde_json::Value {
+        serde_json::json!({
+            "aws_service_name": self.service_name.clone(),
+            "aws_operation_name": self.operation_name.clone()
+        })
+    }
+
     /// Returns the CLI arguments properly formatted as kebab case if parameters is
     /// [Option::Some], otherwise None
     fn cli_parameters(&self) -> Option<Vec<(String, String)>> {
diff --git a/crates/chat-cli/src/cli/mod.rs b/crates/chat-cli/src/cli/mod.rs
@@ -137,7 +137,10 @@ impl RootSubcommand {
 
         // Send executed telemetry.
         if self.valid_for_telemetry() {
-            os.telemetry.send_cli_subcommand_executed(&self).ok();
+            os.telemetry
+                .send_cli_subcommand_executed(&os.database, &self)
+                .await
+                .ok();
         }
 
         match self {
diff --git a/crates/chat-cli/src/telemetry/core.rs b/crates/chat-cli/src/telemetry/core.rs
@@ -42,6 +42,7 @@ pub struct Event {
     pub created_time: Option<SystemTime>,
     pub credential_start_url: Option<String>,
     pub sso_region: Option<String>,
+    pub client_application: Option<String>,
     #[serde(flatten)]
     pub ty: EventType,
 }
@@ -53,6 +54,7 @@ impl Event {
             created_time: Some(SystemTime::now()),
             credential_start_url: None,
             sso_region: None,
+            client_application: None,
         }
     }
 
@@ -64,6 +66,10 @@ impl Event {
         self.sso_region = Some(sso_region);
     }
 
+    pub fn set_client_application(&mut self, client_application: String) {
+        self.client_application = Some(client_application);
+    }
+
     pub fn into_metric_datum(self) -> Option<MetricDatum> {
         match self.ty {
             EventType::UserLoggedIn {} => Some(
@@ -100,6 +106,7 @@ impl Event {
                     credential_start_url: self.credential_start_url.map(Into::into),
                     codewhispererterminal_subcommand: Some(subcommand.into()),
                     codewhispererterminal_in_cloudshell: None,
+                    codewhispererterminal_client_application: self.client_application.map(Into::into),
                 }
                 .into_metric_datum(),
             ),
@@ -152,6 +159,7 @@ impl Event {
                     reason_desc: reason_desc.map(Into::into),
                     status_code: status_code.map(|v| v as i64).map(Into::into),
                     codewhispererterminal_model: model.map(Into::into),
+                    codewhispererterminal_client_application: self.client_application.map(Into::into),
                 }
                 .into_metric_datum(),
             ),
@@ -169,6 +177,8 @@ impl Event {
                 output_token_size,
                 custom_tool_call_latency,
                 model,
+                aws_service_name,
+                aws_operation_name,
             } => Some(
                 CodewhispererterminalToolUseSuggested {
                     create_time: self.created_time,
@@ -190,6 +200,9 @@ impl Event {
                     codewhispererterminal_custom_tool_latency: custom_tool_call_latency
                         .map(|l| CodewhispererterminalCustomToolLatency(l as i64)),
                     codewhispererterminal_model: model.map(Into::into),
+                    codewhispererterminal_client_application: self.client_application.map(Into::into),
+                    codewhispererterminal_aws_service_name: aws_service_name.map(Into::into),
+                    codewhispererterminal_aws_operation_name: aws_operation_name.map(Into::into),
                 }
                 .into_metric_datum(),
             ),
@@ -208,6 +221,7 @@ impl Event {
                     codewhispererterminal_tools_per_mcp_server: Some(CodewhispererterminalToolsPerMcpServer(
                         number_of_tools as i64,
                     )),
+                    codewhispererterminal_client_application: self.client_application.map(Into::into),
                 }
                 .into_metric_datum(),
             ),
@@ -266,6 +280,7 @@ impl Event {
                     reason: reason.map(Into::into),
                     reason_desc: reason_desc.map(Into::into),
                     status_code: status_code.map(|v| v as i64).map(Into::into),
+                    codewhispererterminal_client_application: self.client_application.map(Into::into),
                 }
                 .into_metric_datum(),
             ),
@@ -320,6 +335,8 @@ pub enum EventType {
         output_token_size: Option<usize>,
         custom_tool_call_latency: Option<usize>,
         model: Option<String>,
+        aws_service_name: Option<String>,
+        aws_operation_name: Option<String>,
     },
     McpServerInit {
         conversation_id: String,
@@ -364,6 +381,8 @@ pub struct ToolUseEventBuilder {
     pub output_token_size: Option<usize>,
     pub custom_tool_call_latency: Option<usize>,
     pub model: Option<String>,
+    pub aws_service_name: Option<String>,
+    pub aws_operation_name: Option<String>,
 }
 
 impl ToolUseEventBuilder {
@@ -382,6 +401,8 @@ impl ToolUseEventBuilder {
             output_token_size: None,
             custom_tool_call_latency: None,
             model,
+            aws_service_name: None,
+            aws_operation_name: None,
         }
     }
 
diff --git a/crates/chat-cli/src/telemetry/definitions.rs b/crates/chat-cli/src/telemetry/definitions.rs
@@ -32,6 +32,7 @@ mod tests {
             reason_desc: None,
             status_code: None,
             codewhispererterminal_model: None,
+            codewhispererterminal_client_application: None,
         });
 
         let s = serde_json::to_string_pretty(&metric_datum_init).unwrap();
diff --git a/crates/chat-cli/src/telemetry/mod.rs b/crates/chat-cli/src/telemetry/mod.rs
@@ -66,6 +66,7 @@ pub use crate::telemetry::core::{
     QProfileSwitchIntent,
     TelemetryResult,
 };
+use crate::util::env_var::Q_CLI_CLIENT_APPLICATION;
 use crate::util::system_info::os_version;
 
 #[derive(thiserror::Error, Debug)]
@@ -226,10 +227,17 @@ impl TelemetryThread {
         Ok(self.tx.send(Event::new(EventType::UserLoggedIn {}))?)
     }
 
-    pub fn send_cli_subcommand_executed(&self, subcommand: &RootSubcommand) -> Result<(), TelemetryError> {
-        Ok(self.tx.send(Event::new(EventType::CliSubcommandExecuted {
+    pub async fn send_cli_subcommand_executed(
+        &self,
+        database: &Database,
+        subcommand: &RootSubcommand,
+    ) -> Result<(), TelemetryError> {
+        let mut telemetry_event = Event::new(EventType::CliSubcommandExecuted {
             subcommand: subcommand.to_string(),
-        }))?)
+        });
+        set_event_metadata(database, &mut telemetry_event).await;
+
+        Ok(self.tx.send(telemetry_event)?)
     }
 
     #[allow(clippy::too_many_arguments)] // TODO: Should make a parameters struct.
@@ -246,7 +254,7 @@ impl TelemetryThread {
         status_code: Option<u16>,
         model: Option<String>,
     ) -> Result<(), TelemetryError> {
-        let mut event = Event::new(EventType::ChatAddedMessage {
+        let mut telemetry_event = Event::new(EventType::ChatAddedMessage {
             conversation_id,
             message_id,
             request_id,
@@ -257,13 +265,17 @@ impl TelemetryThread {
             status_code,
             model,
         });
-        set_start_url_and_region(database, &mut event).await;
+        set_event_metadata(database, &mut telemetry_event).await;
 
-        Ok(self.tx.send(event)?)
+        Ok(self.tx.send(telemetry_event)?)
     }
 
-    pub fn send_tool_use_suggested(&self, event: ToolUseEventBuilder) -> Result<(), TelemetryError> {
-        Ok(self.tx.send(Event::new(EventType::ToolUseSuggested {
+    pub async fn send_tool_use_suggested(
+        &self,
+        database: &Database,
+        event: ToolUseEventBuilder,
+    ) -> Result<(), TelemetryError> {
+        let mut telemetry_event = Event::new(EventType::ToolUseSuggested {
             conversation_id: event.conversation_id,
             utterance_id: event.utterance_id,
             user_input_id: event.user_input_id,
@@ -277,20 +289,29 @@ impl TelemetryThread {
             output_token_size: event.output_token_size,
             custom_tool_call_latency: event.custom_tool_call_latency,
             model: event.model,
-        }))?)
+            aws_service_name: event.aws_service_name,
+            aws_operation_name: event.aws_operation_name,
+        });
+        set_event_metadata(database, &mut telemetry_event).await;
+
+        Ok(self.tx.send(telemetry_event)?)
     }
 
-    pub fn send_mcp_server_init(
+    pub async fn send_mcp_server_init(
         &self,
+        database: &Database,
         conversation_id: String,
         init_failure_reason: Option<String>,
         number_of_tools: usize,
     ) -> Result<(), TelemetryError> {
-        Ok(self.tx.send(Event::new(crate::telemetry::EventType::McpServerInit {
+        let mut telemetry_event = Event::new(crate::telemetry::EventType::McpServerInit {
             conversation_id,
             init_failure_reason,
             number_of_tools,
-        }))?)
+        });
+        set_event_metadata(database, &mut telemetry_event).await;
+
+        Ok(self.tx.send(telemetry_event)?)
     }
 
     pub fn send_did_select_profile(
@@ -336,28 +357,33 @@ impl TelemetryThread {
         reason_desc: Option<String>,
         status_code: Option<u16>,
     ) -> Result<(), TelemetryError> {
-        let mut event = Event::new(EventType::MessageResponseError {
+        let mut telemetry_event = Event::new(EventType::MessageResponseError {
             result,
             reason,
             reason_desc,
             status_code,
             conversation_id,
             context_file_length,
         });
-        set_start_url_and_region(database, &mut event).await;
+        set_event_metadata(database, &mut telemetry_event).await;
 
-        Ok(self.tx.send(event)?)
+        Ok(self.tx.send(telemetry_event)?)
     }
 }
 
-async fn set_start_url_and_region(database: &Database, event: &mut Event) {
+async fn set_event_metadata(database: &Database, event: &mut Event) {
     let (start_url, region) = get_start_url_and_region(database).await;
     if let Some(start_url) = start_url {
         event.set_start_url(start_url);
     }
     if let Some(region) = region {
         event.set_sso_region(region);
     }
+
+    // Set the client application from environment variable
+    if let Ok(client_app) = std::env::var(Q_CLI_CLIENT_APPLICATION) {
+        event.set_client_application(client_app);
+    }
 }
 
 #[derive(Debug)]
@@ -627,7 +653,8 @@ mod test {
 
         thread.send_user_logged_in().ok();
         thread
-            .send_cli_subcommand_executed(&RootSubcommand::Version { changelog: None })
+            .send_cli_subcommand_executed(&database, &RootSubcommand::Version { changelog: None })
+            .await
             .ok();
         thread
             .send_chat_added_message(
diff --git a/crates/chat-cli/src/util/consts.rs b/crates/chat-cli/src/util/consts.rs
diff --git a/crates/chat-cli/telemetry_definitions.json b/crates/chat-cli/telemetry_definitions.json

Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,15 @@ impl Tool {`
`129`	`129`	`Tool::Thinking(think) => think.validate(os).await,`
`130`	`130`	`}`
`131`	`131`	`}`
	`132`	`+`
	`133`	`+ /// Returns additional information about the tool if available`
	`134`	`+ pub fn get_additional_info(&self) -> Option<serde_json::Value> {`
	`135`	`+ match self {`
	`136`	`+ Tool::UseAws(use_aws) => Some(use_aws.get_additional_info()),`
	`137`	`+ // Add other tool types here as they implement get_additional_info()`
	`138`	`+ _ => None,`
	`139`	`+ }`
	`140`	`+ }`
`132`	`141`	`}`
`133`	`142`
`134`	`143`	`#[derive(Debug, Clone)]`