feat: add telemetry for agent config and tool execution timing (#2400)

Author: brandonskiser

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

@@ -20,10 +20,15 @@ use crate::cli::agent::legacy::context::LegacyContextConfig;
 use crate::os::Os;
 use crate::util::directories;
 
-pub async fn migrate(os: &mut Os) -> eyre::Result<Vec<Agent>> {
+/// Performs the migration from legacy profile configuration to agent configuration if it hasn't
+/// already been done.
+///
+/// Returns [Some] with the newly migrated agents if the migration was performed, [None] if the
+/// migration was already done previously.
+pub async fn migrate(os: &mut Os) -> eyre::Result<Option<Vec<Agent>>> {
     let has_migrated = os.database.get_has_migrated()?;
     if has_migrated.is_some_and(|has_migrated| has_migrated) {
-        bail!("Nothing to migrate");
+        return Ok(None);
     }
 
     let legacy_global_context_path = directories::chat_global_context_path(os)?;
@@ -216,5 +221,5 @@ pub async fn migrate(os: &mut Os) -> eyre::Result<Vec<Agent>> {
 
     os.database.set_has_migrated()?;
 
-    Ok(new_agents)
+    Ok(Some(new_agents))
 }

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

@@ -43,6 +43,7 @@ use thiserror::Error;
 use tokio::fs::ReadDir;
 use tracing::{
     error,
+    info,
     warn,
 };
 use wrapper_types::ResourcePath;
@@ -382,12 +383,33 @@ impl Agents {
     ///   agent selection
     /// * `skip_migration` - If true, skips migration of old profiles to new format
     /// * `output` - Writer for outputting warnings, errors, and status messages during loading
-    pub async fn load(os: &mut Os, agent_name: Option<&str>, skip_migration: bool, output: &mut impl Write) -> Self {
+    pub async fn load(
+        os: &mut Os,
+        agent_name: Option<&str>,
+        skip_migration: bool,
+        output: &mut impl Write,
+    ) -> (Self, AgentsLoadMetadata) {
+        // Tracking metadata about the performed load operation.
+        let mut load_metadata = AgentsLoadMetadata {
+            launched_agent: agent_name.map(Into::into),
+            ..Default::default()
+        };
+
         let new_agents = if !skip_migration {
             match legacy::migrate(os).await {
-                Ok(new_agents) => new_agents,
+                Ok(Some(new_agents)) => {
+                    let migrated_count = new_agents.len();
+                    info!(migrated_count, "Profile migration successful");
+                    load_metadata.migration_performed = true;
+                    load_metadata.migrated_count = migrated_count as u32;
+                    new_agents
+                },
+                Ok(None) => {
+                    info!("Migration was not performed");
+                    vec![]
+                },
                 Err(e) => {
-                    warn!("Migration did not happen for the following reason: {e}. This is not necessarily an error");
+                    error!("Migration did not happen for the following reason: {e}");
                     vec![]
                 },
             }
@@ -421,6 +443,7 @@ impl Agents {
                 match result {
                     Ok(agent) => agents.push(agent),
                     Err(e) => {
+                        load_metadata.load_failed_count += 1;
                         let _ = queue!(
                             output,
                             style::SetForegroundColor(Color::Red),
@@ -458,6 +481,7 @@ impl Agents {
                 match result {
                     Ok(agent) => agents.push(agent),
                     Err(e) => {
+                        load_metadata.load_failed_count += 1;
                         let _ = queue!(
                             output,
                             style::SetForegroundColor(Color::Red),
@@ -665,11 +689,14 @@ impl Agents {
             }
         }
 
-        Self {
-            agents,
-            active_idx,
-            ..Default::default()
-        }
+        (
+            Self {
+                agents,
+                active_idx,
+                ..Default::default()
+            },
+            load_metadata,
+        )
     }
 
     /// Returns a label to describe the permission status for a given tool.
@@ -717,6 +744,16 @@ impl Agents {
     }
 }
 
+/// Metadata from the executed [Agents::load] operation.
+#[derive(Debug, Clone, Default)]
+pub struct AgentsLoadMetadata {
+    pub migration_performed: bool,
+    pub migrated_count: u32,
+    pub load_count: u32,
+    pub load_failed_count: u32,
+    pub launched_agent: Option<String>,
+}
+
 async fn load_agents_from_entries(
     mut files: ReadDir,
     os: &Os,

crates/chat-cli/src/cli/agent/root_command_args.rs

@@ -75,7 +75,7 @@ impl AgentArgs {
         let mut stderr = std::io::stderr();
         match self.cmd {
             Some(AgentSubcommands::List) | None => {
-                let agents = Agents::load(os, None, true, &mut stderr).await;
+                let agents = Agents::load(os, None, true, &mut stderr).await.0;
                 let agent_with_path =
                     agents
                         .agents
@@ -100,7 +100,7 @@ impl AgentArgs {
                 writeln!(stderr, "{}", output_str)?;
             },
             Some(AgentSubcommands::Create { name, directory, from }) => {
-                let mut agents = Agents::load(os, None, true, &mut stderr).await;
+                let mut agents = Agents::load(os, None, true, &mut stderr).await.0;
                 let path_with_file_name = create_agent(os, &mut agents, name.clone(), directory, from).await?;
                 let editor_cmd = std::env::var("EDITOR").unwrap_or_else(|_| "vi".to_string());
                 let mut cmd = std::process::Command::new(editor_cmd);
@@ -131,7 +131,7 @@ impl AgentArgs {
                 )?;
             },
             Some(AgentSubcommands::Rename { agent, new_name }) => {
-                let mut agents = Agents::load(os, None, true, &mut stderr).await;
+                let mut agents = Agents::load(os, None, true, &mut stderr).await.0;
                 rename_agent(os, &mut agents, agent.clone(), new_name.clone()).await?;
                 writeln!(stderr, "\n✓ Renamed agent '{}' to '{}'\n", agent, new_name)?;
             },

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

@@ -133,7 +133,7 @@ impl AgentSubcommand {
                     .map_err(|e| ChatError::Custom(format!("Error printing agent schema: {e}").into()))?;
             },
             Self::Create { name, directory, from } => {
-                let mut agents = Agents::load(os, None, true, &mut session.stderr).await;
+                let mut agents = Agents::load(os, None, true, &mut session.stderr).await.0;
                 let path_with_file_name = create_agent(os, &mut agents, name.clone(), directory, from)
                     .await
                     .map_err(|e| ChatError::Custom(Cow::Owned(e.to_string())))?;
@@ -183,7 +183,7 @@ impl AgentSubcommand {
                 )?;
             },
             Self::Rename { agent, new_name } => {
-                let mut agents = Agents::load(os, None, true, &mut session.stderr).await;
+                let mut agents = Agents::load(os, None, true, &mut session.stderr).await.0;
                 rename_agent(os, &mut agents, agent.clone(), new_name.clone())
                     .await
                     .map_err(|e| ChatError::Custom(Cow::Owned(e.to_string())))?;

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

@@ -30,7 +30,10 @@ use std::io::{
 };
 use std::process::ExitCode;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{
+    Duration,
+    Instant,
+};
 
 use amzn_codewhisperer_client::types::SubscriptionStatus;
 use clap::{
@@ -138,6 +141,7 @@ use crate::database::settings::Setting;
 use crate::mcp_client::Prompt;
 use crate::os::Os;
 use crate::telemetry::core::{
+    AgentConfigInitArgs,
     ChatAddedMessageParams,
     ChatConversationType,
     MessageMetaTag,
@@ -240,11 +244,26 @@ impl ChatArgs {
             )?;
         }
 
+        let conversation_id = uuid::Uuid::new_v4().to_string();
+        info!(?conversation_id, "Generated new conversation id");
+
         let agents = {
             let skip_migration = self.no_interactive;
-            let mut agents = Agents::load(os, self.agent.as_deref(), skip_migration, &mut stderr).await;
+            let (mut agents, md) = Agents::load(os, self.agent.as_deref(), skip_migration, &mut stderr).await;
             agents.trust_all_tools = self.trust_all_tools;
 
+            os.telemetry
+                .send_agent_config_init(&os.database, conversation_id.clone(), AgentConfigInitArgs {
+                    agents_loaded_count: md.load_count as i64,
+                    agents_loaded_failed_count: md.load_failed_count as i64,
+                    legacy_profile_migration_executed: md.migration_performed,
+                    legacy_profile_migrated_count: md.migrated_count as i64,
+                    launched_agent: md.launched_agent,
+                })
+                .await
+                .map_err(|err| error!(?err, "failed to send agent config init telemetry"))
+                .ok();
+
             if agents
                 .get_active()
                 .is_some_and(|a| !a.mcp_servers.mcp_servers.is_empty())
@@ -287,8 +306,6 @@ impl ChatArgs {
             None
         };
 
-        let conversation_id = uuid::Uuid::new_v4().to_string();
-        info!(?conversation_id, "Generated new conversation id");
         let (prompt_request_sender, prompt_request_receiver) = std::sync::mpsc::channel::<Option<String>>();
         let (prompt_response_sender, prompt_response_receiver) = std::sync::mpsc::channel::<Vec<String>>();
         let mut tool_manager = ToolManagerBuilder::default()
@@ -495,10 +512,17 @@ pub struct ChatSession {
     spinner: Option<Spinner>,
     /// [ConversationState].
     conversation: ConversationState,
+    /// Tool uses requested by the model that are actively being handled.
     tool_uses: Vec<QueuedTool>,
+    /// An index into [Self::tool_uses] to represent the current tool use being handled.
+    pending_tool_index: Option<usize>,
+    /// The time immediately after having received valid tool uses from the model.
+    ///
+    /// Used to track the time taken from initially prompting the user to tool execute
+    /// completion.
+    tool_turn_start_time: Option<Instant>,
     /// [RequestMetadata] about the ongoing operation.
     user_turn_request_metadata: Vec<RequestMetadata>,
-    pending_tool_index: Option<usize>,
     /// Telemetry events to be sent as part of the conversation. The HashMap key is tool_use_id.
     tool_use_telemetry_events: HashMap<String, ToolUseEventBuilder>,
     /// State used to keep track of tool use relation
@@ -622,6 +646,7 @@ impl ChatSession {
             tool_uses: vec![],
             user_turn_request_metadata: vec![],
             pending_tool_index: None,
+            tool_turn_start_time: None,
             tool_use_telemetry_events: HashMap::new(),
             tool_use_status: ToolUseStatus::Idle,
             failed_request_ids: Vec::new(),
@@ -978,6 +1003,7 @@ impl ChatSession {
         self.conversation.enforce_conversation_invariants();
         self.conversation.reset_next_user_message();
         self.pending_tool_index = None;
+        self.tool_turn_start_time = None;
         self.reset_user_turn();
 
         self.inner = Some(ChatState::PromptUser {
@@ -1806,6 +1832,9 @@ impl ChatSession {
 
             if allowed {
                 tool.accepted = true;
+                self.tool_use_telemetry_events
+                    .entry(tool.id.clone())
+                    .and_modify(|ev| ev.is_trusted = true);
                 continue;
             }
 
@@ -1821,10 +1850,12 @@ impl ChatSession {
         let mut image_blocks: Vec<RichImageBlock> = Vec::new();
 
         for tool in &self.tool_uses {
+            let tool_start = std::time::Instant::now();
             let mut tool_telemetry = self.tool_use_telemetry_events.entry(tool.id.clone());
-            tool_telemetry = tool_telemetry.and_modify(|ev| ev.is_accepted = true);
+            tool_telemetry = tool_telemetry.and_modify(|ev| {
+                ev.is_accepted = true;
+            });
 
-            let tool_start = std::time::Instant::now();
             let invoke_result = tool.tool.invoke(os, &mut self.stdout).await;
 
             if self.spinner.is_some() {
@@ -1837,12 +1868,18 @@ impl ChatSession {
             }
             execute!(self.stdout, style::Print("\n"))?;
 
-            let tool_time = std::time::Instant::now().duration_since(tool_start);
+            let tool_end_time = Instant::now();
+            let tool_time = tool_end_time.duration_since(tool_start);
+            tool_telemetry = tool_telemetry.and_modify(|ev| {
+                ev.execution_duration = Some(tool_time);
+                ev.turn_duration = self.tool_turn_start_time.map(|t| tool_end_time.duration_since(t));
+            });
             if let Tool::Custom(ct) = &tool.tool {
                 tool_telemetry = tool_telemetry.and_modify(|ev| {
+                    ev.is_custom_tool = true;
+                    // legacy fields previously implemented for only MCP tools
                     ev.custom_tool_call_latency = Some(tool_time.as_secs() as usize);
                     ev.input_token_size = Some(ct.get_input_token_size());
-                    ev.is_custom_tool = true;
                 });
             }
             let tool_time = format!("{}.{}", tool_time.as_secs(), tool_time.subsec_millis());
@@ -2232,6 +2269,7 @@ impl ChatSession {
         } else {
             self.tool_uses.clear();
             self.pending_tool_index = None;
+            self.tool_turn_start_time = None;
 
             self.send_chat_telemetry(os, TelemetryResult::Succeeded, None, None, None, true)
                 .await;
@@ -2342,6 +2380,7 @@ impl ChatSession {
 
         self.tool_uses = queued_tools;
         self.pending_tool_index = Some(0);
+        self.tool_turn_start_time = Some(Instant::now());
         Ok(ChatState::ExecuteTools)
     }
 
@@ -2353,6 +2392,7 @@ impl ChatSession {
                 self.conversation.enforce_conversation_invariants();
                 self.conversation.reset_next_user_message();
                 self.pending_tool_index = None;
+                self.tool_turn_start_time = None;
                 return Ok(ChatState::PromptUser {
                     skip_printing_tools: false,
                 });