add modelId in telemetry & persist modelId

Author: evanliu048

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

@@ -105,7 +105,8 @@ pub struct ConversationState {
     latest_summary: Option<String>,
     #[serde(skip)]
     pub updates: Option<SharedWriter>,
-    /// Model explicitly selected by the user in this conversation state via `/model`. (`None` == auto)
+    /// Model explicitly selected by the user in this conversation state via `/model`. (`None` ==
+    /// auto)
     #[serde(skip)]
     pub current_model_id: Option<String>,
 }

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

@@ -567,7 +567,10 @@ impl ChatContext {
         let output_clone = output.clone();
 
         let mut existing_conversation = false;
-        let model_id = database.settings.get_string(Setting::UserDefaultModel);
+        let model_id = match database.get_last_used_model_id() {
+            Ok(Some(id)) => Some(id),
+            Ok(None) | Err(_) => database.settings.get_string(Setting::UserDefaultModel),
+        };
         let conversation_state = if resume_conversation {
             let prior = std::env::current_dir()
                 .ok()
@@ -594,7 +597,7 @@ impl ChatContext {
                     profile,
                     Some(output_clone),
                     tool_manager,
-                    model_id
+                    model_id,
                 )
                 .await
             }
@@ -606,7 +609,7 @@ impl ChatContext {
                 profile,
                 Some(output_clone),
                 tool_manager,
-                model_id
+                model_id,
             )
             .await
         };
@@ -3042,32 +3045,67 @@ impl ChatContext {
             },
             Command::Model => {
                 queue!(self.output, style::Print("\n"))?;
-                let labels: Vec<&str> = MODEL_OPTIONS.iter().map(|(l, _)| *l).collect();
+                let active_model_id = self.conversation_state.current_model_id.as_deref();
+                let labels: Vec<String> = MODEL_OPTIONS
+                    .iter()
+                    .map(|(label, model_id)| {
+                        if (model_id.is_empty() && active_model_id.is_none()) || Some(*model_id) == active_model_id {
+                            format!("{} (active)", label)
+                        } else {
+                            label.to_string()
+                        }
+                    })
+                    .collect();
+                let default_index = MODEL_OPTIONS
+                    .iter()
+                    .position(|(_, model_id)| Some(*model_id) == active_model_id)
+                    .unwrap_or(0);
                 let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
                     .with_prompt("Select a model for this chat session")
                     .items(&labels)
-                    .default(0)
+                    .default(default_index)
                     .interact_on_opt(&dialoguer::console::Term::stdout())
                 {
-                    Ok(sel) => sel,
+                    Ok(sel) => {
+                        let _ = crossterm::execute!(
+                            std::io::stdout(),
+                            crossterm::style::SetForegroundColor(crossterm::style::Color::Magenta)
+                        );
+                        sel
+                    },
                     // Ctrl‑C -> Err(Interrupted)
                     Err(DError::IO(ref e)) if e.kind() == io::ErrorKind::Interrupted => None,
                     Err(e) => return Err(ChatError::Custom(format!("Failed to choose model: {e}").into())),
                 };
 
+                queue!(self.output, style::ResetColor)?;
+
                 if let Some(index) = selection {
-                    let (label, model_opt) = MODEL_OPTIONS[index];
+                    let (label, model_id) = MODEL_OPTIONS[index];
+                    let model_id_str = model_id.to_string();
 
-                    self.conversation_state.current_model_id = Some(model_opt.to_string());
-                    telemetry.update_model_id(self.conversation_state.current_model_id.clone());
+                    if model_id == "" {
+                        self.conversation_state.current_model_id = None;
+                        telemetry.update_model_id(None);
+                        let _ = database.unset_last_used_model_id();
+                    } else {
+                        self.conversation_state.current_model_id = Some(model_id_str.clone());
+                        telemetry.update_model_id(Some(model_id_str.clone()));
+                        let _ = database.set_last_used_model_id(model_id_str);
+                    }
 
-                    use crossterm::{
-                        queue,
-                        style,
-                    };
-                    queue!(self.output, style::Print("\n"), style::Print(format!(" Swtiched model to {}\n\n", label)))?;
+                    queue!(
+                        self.output,
+                        style::Print("\n"),
+                        style::Print(format!(" Switched model to {}\n\n", label)),
+                        style::ResetColor,
+                        style::SetForegroundColor(Color::Reset),
+                        style::SetBackgroundColor(Color::Reset),
+                    )?;
                 }
-                
+
+                queue!(self.output, style::ResetColor)?;
+                self.output.flush()?;
                 ChatState::PromptUser {
                     tool_uses: None,
                     pending_tool_index: None,

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

@@ -60,6 +60,7 @@ const IDC_REGION_KEY: &str = "auth.idc.region";
 // We include this key to remove for backwards compatibility
 const CUSTOMIZATION_STATE_KEY: &str = "api.selectedCustomization";
 const ROTATING_TIP_KEY: &str = "chat.greeting.rotating_tips_current_index";
+const LAST_USED_MODEL_ID: &str = "lastUsedModelId";
 
 const MIGRATIONS: &[Migration] = migrations![
     "000_migration_table",
@@ -307,6 +308,21 @@ impl Database {
         Ok(tip)
     }
 
+    /// Get the model id used for last conversation state.
+    pub fn get_last_used_model_id(&self) -> Result<Option<String>, DatabaseError> {
+        self.get_json_entry::<String>(Table::State, LAST_USED_MODEL_ID)
+    }
+
+    /// Set the model id used for last conversation state.
+    pub fn set_last_used_model_id(&mut self, last_used_model_id: String) -> Result<usize, DatabaseError> {
+        self.set_json_entry(Table::State, LAST_USED_MODEL_ID, last_used_model_id)
+    }
+
+    /// UnsSet the model id used for last conversation state.
+    pub fn unset_last_used_model_id(&mut self) -> Result<(), DatabaseError> {
+        self.delete_entry(Table::State, LAST_USED_MODEL_ID)
+    }
+
     /// Get a chat conversation given a path to the conversation.
     pub fn get_conversation_by_path(
         &mut self,

crates/chat-cli/src/database/settings.rs

@@ -202,14 +202,12 @@ mod test {
         assert_eq!(settings.get(Setting::McpLoadedBefore), None);
         assert_eq!(settings.get(Setting::UserDefaultModel), None);
 
-
         settings.set(Setting::TelemetryEnabled, true).await.unwrap();
         settings.set(Setting::OldClientId, "test").await.unwrap();
         settings.set(Setting::ShareCodeWhispererContent, false).await.unwrap();
         settings.set(Setting::McpLoadedBefore, true).await.unwrap();
         settings.set(Setting::UserDefaultModel, "model 1").await.unwrap();
 
-
         assert_eq!(settings.get(Setting::TelemetryEnabled), Some(&Value::Bool(true)));
         assert_eq!(
             settings.get(Setting::OldClientId),

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

@@ -6,7 +6,10 @@ mod install_method;
 
 use core::ToolUseEventBuilder;
 use std::str::FromStr;
-use std::sync::{Arc, RwLock};
+use std::sync::{
+    Arc,
+    RwLock,
+};
 
 use amzn_codewhisperer_client::types::{
     ChatAddMessageEvent,
@@ -148,8 +151,11 @@ impl Clone for TelemetryThread {
 
 impl TelemetryThread {
     pub async fn new(env: &Env, database: &mut Database) -> Result<Self, TelemetryError> {
-        //todo yifan change to last used modelid
-        let current_model_id = Arc::new(RwLock::new(None));
+        let model_id = match database.get_last_used_model_id() {
+            Ok(Some(id)) => Some(id),
+            Ok(None) | Err(_) => database.settings.get_string(Setting::UserDefaultModel),
+        };
+        let current_model_id = Arc::new(RwLock::new(model_id));
         let telemetry_client = TelemetryClient::new(env, database, current_model_id.clone()).await?;
         let (tx, mut rx) = mpsc::unbounded_channel();
         let handle = tokio::spawn(async move {
@@ -330,7 +336,11 @@ struct TelemetryClient {
 }
 
 impl TelemetryClient {
-    async fn new(env: &Env, database: &mut Database, current_model_id: Arc<RwLock<Option<String>>>) -> Result<Self, TelemetryError> {
+    async fn new(
+        env: &Env,
+        database: &mut Database,
+        current_model_id: Arc<RwLock<Option<String>>>,
+    ) -> Result<Self, TelemetryError> {
         let telemetry_enabled = !cfg!(test)
             && env.get_os("Q_DISABLE_TELEMETRY").is_none()
             && database.settings.get_bool(Setting::TelemetryEnabled).unwrap_or(true);
@@ -506,7 +516,13 @@ mod test {
     #[tokio::test]
     async fn client_context() {
         let mut database = Database::new().await.unwrap();
-        let client = TelemetryClient::new(&Env::new(), &mut database, Arc::new(RwLock::new(Some("model".to_owned())))).await.unwrap();
+        let client = TelemetryClient::new(
+            &Env::new(),
+            &mut database,
+            Arc::new(RwLock::new(Some("model".to_owned()))),
+        )
+        .await
+        .unwrap();
         let context = client.user_context().unwrap();
 
         assert_eq!(context.ide_category, IdeCategory::Cli);
@@ -575,7 +591,13 @@ mod test {
     #[ignore = "needs auth which is not in CI"]
     async fn test_without_optout() {
         let mut database = Database::new().await.unwrap();
-        let client = TelemetryClient::new(&Env::new(), &mut database, Arc::new(RwLock::new(Some("model".to_owned())))).await.unwrap();
+        let client = TelemetryClient::new(
+            &Env::new(),
+            &mut database,
+            Arc::new(RwLock::new(Some("model".to_owned()))),
+        )
+        .await
+        .unwrap();
         client
             .codewhisperer_client
             .send_telemetry_event(

crates/fig_api_client/src/clients/client.rs

@@ -12,7 +12,6 @@ use amzn_codewhisperer_client::types::{
     TelemetryEvent,
     UserContext,
 };
-use crate::model::ConversationState;
 use amzn_consolas_client::Client as ConsolasClient;
 use amzn_consolas_client::error::DisplayErrorContext;
 use amzn_consolas_client::operation::generate_recommendations::GenerateRecommendationsError;
@@ -224,7 +223,6 @@ impl Client {
         Ok(customizations)
     }
 
-    //(todo) yifan : how to propgate modelID
     // .telemetry_event(TelemetryEvent::UserTriggerDecisionEvent(user_trigger_decision_event))
     // .user_context(user_context)
     // .opt_out_preference(opt_out_preference)

crates/fig_api_client/src/model.rs

@@ -646,7 +646,7 @@ pub struct UserInputMessage {
     pub user_input_message_context: Option<UserInputMessageContext>,
     pub user_intent: Option<UserIntent>,
     pub images: Option<Vec<ImageBlock>>,
-    pub model_id: Option<String>
+    pub model_id: Option<String>,
 }
 
 impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::UserInputMessage {