chore: Updates runtime models

kennvene-aws · kennvene-aws · commit 040dd567ac8b · 2025-10-02T18:51:27.000-07:00
diff --git a/crates/chat-cli/src/api_client/mod.rs b/crates/chat-cli/src/api_client/mod.rs
@@ -382,6 +382,7 @@ impl ApiClient {
             conversation_id,
             user_input_message,
             history,
+            agent_continuation_id,
         } = conversation;
 
         let model_id_opt: Option<String> = user_input_message.model_id.clone();
@@ -400,6 +401,8 @@ impl ApiClient {
                         .map(|v| v.into_iter().map(|i| i.try_into()).collect::<Result<Vec<_>, _>>())
                         .transpose()?,
                 )
+                .set_agent_continuation_id(agent_continuation_id)
+                .agent_task_type(amzn_codewhisperer_streaming_client::types::AgentTaskType::Vibe)
                 .build()
                 .expect("building conversation should not fail");
 
@@ -744,6 +747,7 @@ mod tests {
                     model_id: Some("model".to_owned()),
                 },
                 history: None,
+                agent_continuation_id: None,
             })
             .await
             .unwrap();
diff --git a/crates/chat-cli/src/api_client/model.rs b/crates/chat-cli/src/api_client/model.rs
@@ -97,6 +97,7 @@ pub struct ConversationState {
     pub conversation_id: Option<String>,
     pub user_input_message: UserInputMessage,
     pub history: Option<Vec<ChatMessage>>,
+    pub agent_continuation_id: Option<String>,
 }
 
 #[derive(Debug, Clone)]
@@ -563,6 +564,12 @@ pub enum ChatResponseStream {
         conversation_id: Option<String>,
         utterance_id: Option<String>,
     },
+    MetadataEvent {
+        total_tokens: Option<i32>,
+    },
+    MeteringEvent {
+        usage: Option<f64>,
+    },
     SupplementaryWebLinksEvent(()),
     ToolUseEvent {
         tool_use_id: String,
@@ -609,6 +616,23 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
                 conversation_id,
                 utterance_id,
             },
+            amzn_codewhisperer_streaming_client::types::ChatResponseStream::MetadataEvent(
+                amzn_codewhisperer_streaming_client::types::MetadataEvent { token_usage, .. },
+            ) => {
+                println!(
+                    "DEBUG: Codewhisperer MetadataEvent matched! token_usage: {:?}",
+                    token_usage
+                );
+                ChatResponseStream::MetadataEvent {
+                    total_tokens: token_usage.map(|t| t.total_tokens),
+                }
+            },
+            amzn_codewhisperer_streaming_client::types::ChatResponseStream::MeteringEvent(
+                amzn_codewhisperer_streaming_client::types::MeteringEvent { usage, .. },
+            ) => {
+                println!("DEBUG: Codewhisperer MeteringEvent matched! usage: {:?}", usage);
+                ChatResponseStream::MeteringEvent { usage }
+            },
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_codewhisperer_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -626,7 +650,10 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
                 ChatResponseStream::SupplementaryWebLinksEvent(())
             },
-            _ => ChatResponseStream::Unknown,
+            other => {
+                println!("DEBUG: Codewhisperer catch-all - Unknown event: {:?}", other);
+                ChatResponseStream::Unknown
+            },
         }
     }
 }
@@ -665,6 +692,23 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
                 conversation_id,
                 utterance_id,
             },
+            amzn_qdeveloper_streaming_client::types::ChatResponseStream::MetadataEvent(
+                amzn_qdeveloper_streaming_client::types::MetadataEvent { token_usage, .. },
+            ) => {
+                println!(
+                    "DEBUG: QDeveloper MetadataEvent matched! token_usage: {:?}",
+                    token_usage
+                );
+                ChatResponseStream::MetadataEvent {
+                    total_tokens: token_usage.map(|t| t.total_tokens),
+                }
+            },
+            amzn_qdeveloper_streaming_client::types::ChatResponseStream::MeteringEvent(
+                amzn_qdeveloper_streaming_client::types::MeteringEvent { usage, .. },
+            ) => {
+                println!("DEBUG: QDeveloper MeteringEvent matched! usage: {:?}", usage);
+                ChatResponseStream::MeteringEvent { usage }
+            },
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::ToolUseEvent(
                 amzn_qdeveloper_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -682,7 +726,10 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
                 ChatResponseStream::SupplementaryWebLinksEvent(())
             },
-            _ => ChatResponseStream::Unknown,
+            other => {
+                println!("DEBUG: Codewhisperer catch-all - Unknown event: {:?}", other);
+                ChatResponseStream::Unknown
+            },
         }
     }
 }
@@ -870,7 +917,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::Cli)
+            .origin(amzn_codewhisperer_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }
@@ -884,7 +931,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::Cli)
+            .origin(amzn_qdeveloper_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }
diff --git a/crates/chat-cli/src/cli/chat/cli/prompts.rs b/crates/chat-cli/src/cli/chat/cli/prompts.rs
@@ -2490,23 +2490,23 @@ mod tests {
             prompt_get: prompt1,
         };
 
-        let bundles = vec![&bundle1, &bundle2];
+        let bundles = [&bundle1, &bundle2];
 
         // Test filtering by server
         let filtered: Vec<&PromptBundle> = bundles.iter().filter(|b| b.server_name == "server1").copied().collect();
         assert_eq!(filtered.len(), 1);
         assert_eq!(filtered[0].server_name, "server1");
 
         // Test no filtering (all bundles)
-        let all: Vec<&PromptBundle> = bundles.iter().copied().collect();
+        let all: Vec<&PromptBundle> = bundles.to_vec();
         assert_eq!(all.len(), 2);
     }
 
     #[test]
     fn test_ambiguous_prompt_message_generation() {
         // Test generating disambiguation message
         let prompt_name = "test_prompt";
-        let server_names = vec!["server1", "server2", "server3"];
+        let server_names = ["server1", "server2", "server3"];
 
         let alt_names: Vec<String> = server_names
             .iter()
diff --git a/crates/chat-cli/src/cli/chat/conversation.rs b/crates/chat-cli/src/cli/chat/conversation.rs
@@ -27,6 +27,7 @@ use tracing::{
     debug,
     warn,
 };
+use uuid;
 
 use super::cli::compact::CompactStrategy;
 use super::cli::hooks::HookOutput;
@@ -149,6 +150,9 @@ pub struct ConversationState {
     /// Tangent mode checkpoint - stores main conversation when in tangent mode
     #[serde(default, skip_serializing_if = "Option::is_none")]
     tangent_state: Option<ConversationCheckpoint>,
+    /// Current continuation ID for billing tracking - generated per conversation turn
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    current_continuation_id: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -212,6 +216,7 @@ impl ConversationState {
             checkpoint_manager: None,
             mcp_enabled,
             tangent_state: None,
+            current_continuation_id: None,
         }
     }
 
@@ -376,6 +381,12 @@ impl ConversationState {
 
     pub fn reset_next_user_message(&mut self) {
         self.next_message = None;
+        // Don't clear continuation_id here - it should persist for the entire turn
+    }
+
+    /// Returns the current continuation ID for billing tracking
+    pub fn continuation_id(&self) -> Option<&str> {
+        self.current_continuation_id.as_deref()
     }
 
     pub async fn set_next_user_message(&mut self, input: String) {
@@ -391,6 +402,10 @@ impl ConversationState {
             input
         };
 
+        // Generate new continuation ID for this turn
+        let new_continuation_id = uuid::Uuid::new_v4().to_string();
+        self.current_continuation_id = Some(new_continuation_id);
+
         let msg = UserMessage::new_prompt(input, Some(Local::now().fixed_offset()));
         self.next_message = Some(msg);
     }
@@ -614,6 +629,7 @@ impl ConversationState {
             dropped_context_files,
             tools: &self.tools,
             model_id: self.model_info.as_ref().map(|m| m.model_id.as_str()),
+            continuation_id: self.current_continuation_id.as_deref(),
         })
     }
 
@@ -719,6 +735,7 @@ impl ConversationState {
                 .unwrap_or(UserMessage::new_prompt(summary_content, None)) // should not happen
                 .into_user_input_message(self.model_info.as_ref().map(|m| m.model_id.clone()), &tools),
             history: Some(flatten_history(history.iter())),
+            agent_continuation_id: self.current_continuation_id.clone(),
         })
     }
 
@@ -777,6 +794,7 @@ Return only the JSON configuration, no additional text.",
             conversation_id: Some(self.conversation_id.clone()),
             user_input_message: generation_message.into_user_input_message(self.model.clone(), &tools),
             history: Some(flatten_history(history.iter())),
+            agent_continuation_id: self.current_continuation_id.clone(),
         })
     }
 
@@ -957,6 +975,9 @@ Return only the JSON configuration, no additional text.",
     }
 }
 
+#[cfg(test)]
+mod continuation_tests;
+
 pub fn format_tool_spec(tool_spec: HashMap<String, ToolSpec>) -> HashMap<ToolOrigin, Vec<Tool>> {
     tool_spec
         .into_values()
@@ -992,6 +1013,7 @@ pub struct BackendConversationStateImpl<'a, T, U> {
     pub dropped_context_files: Vec<(String, String)>,
     pub tools: &'a HashMap<ToolOrigin, Vec<Tool>>,
     pub model_id: Option<&'a str>,
+    pub continuation_id: Option<&'a str>,
 }
 
 impl BackendConversationStateImpl<'_, std::collections::vec_deque::Iter<'_, HistoryEntry>, Option<Vec<HistoryEntry>>> {
@@ -1007,6 +1029,7 @@ impl BackendConversationStateImpl<'_, std::collections::vec_deque::Iter<'_, Hist
             conversation_id: Some(self.conversation_id.to_string()),
             user_input_message,
             history: Some(history),
+            agent_continuation_id: self.continuation_id.map(str::to_string),
         })
     }
 
diff --git a/crates/chat-cli/src/cli/chat/conversation/continuation_tests.rs b/crates/chat-cli/src/cli/chat/conversation/continuation_tests.rs
@@ -0,0 +1,136 @@
+// ABOUTME: Tests for continuation ID billing functionality
+// ABOUTME: Ensures continuation IDs are generated and propagated correctly for billing consistency
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use uuid;
+
+    use super::super::ConversationState;
+    use crate::cli::agent::Agents;
+    use crate::cli::chat::tool_manager::ToolManager;
+    use crate::os::Os;
+
+    #[tokio::test]
+    async fn test_continuation_id_generation() {
+        // Setup
+        let os = Os::new().await.unwrap();
+        let agents = Agents::default();
+        let tool_manager = ToolManager::default();
+        let mut conversation = ConversationState::new(
+            "test-conversation",
+            agents,
+            HashMap::new(),
+            tool_manager,
+            None,
+            &os,
+            true,
+        )
+        .await;
+
+        // Initially no continuation ID
+        assert!(conversation.continuation_id().is_none());
+
+        // Set user message should generate continuation ID
+        conversation.set_next_user_message("test message".to_string()).await;
+
+        let continuation_id = conversation.continuation_id();
+        assert!(continuation_id.is_some());
+        assert!(!continuation_id.unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_continuation_id_uniqueness() {
+        // Setup
+        let os = Os::new().await.unwrap();
+        let agents = Agents::default();
+        let tool_manager = ToolManager::default();
+        let mut conversation = ConversationState::new(
+            "test-conversation",
+            agents,
+            HashMap::new(),
+            tool_manager,
+            None,
+            &os,
+            true,
+        )
+        .await;
+
+        // Generate first continuation ID
+        conversation.set_next_user_message("first message".to_string()).await;
+        let first_id = conversation.continuation_id().unwrap().to_string();
+
+        // Reset and generate second continuation ID
+        conversation.reset_next_user_message();
+        conversation.set_next_user_message("second message".to_string()).await;
+        let second_id = conversation.continuation_id().unwrap().to_string();
+
+        // Should be different
+        assert_ne!(first_id, second_id);
+    }
+
+    #[tokio::test]
+    async fn test_continuation_id_format() {
+        // Setup
+        let os = Os::new().await.unwrap();
+        let agents = Agents::default();
+        let tool_manager = ToolManager::default();
+        let mut conversation = ConversationState::new(
+            "test-conversation",
+            agents,
+            HashMap::new(),
+            tool_manager,
+            None,
+            &os,
+            true,
+        )
+        .await;
+
+        // Generate continuation ID
+        conversation.set_next_user_message("test message".to_string()).await;
+        let continuation_id = conversation.continuation_id().unwrap();
+
+        // Should be a valid UUID format (36 characters with hyphens)
+        assert_eq!(continuation_id.len(), 36);
+        assert!(continuation_id.chars().filter(|&c| c == '-').count() == 4);
+
+        // Should be parseable as UUID
+        assert!(uuid::Uuid::parse_str(continuation_id).is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_continuation_id_in_api_request() {
+        // Setup
+        let os = Os::new().await.unwrap();
+        let agents = Agents::default();
+        let tool_manager = ToolManager::default();
+        let mut conversation = ConversationState::new(
+            "test-conversation",
+            agents,
+            HashMap::new(),
+            tool_manager,
+            None,
+            &os,
+            true,
+        )
+        .await;
+
+        // Set user message to generate continuation ID
+        conversation.set_next_user_message("test message".to_string()).await;
+        let expected_continuation_id = conversation.continuation_id().unwrap().to_string();
+
+        // Create sendable conversation state (this is what gets sent to API)
+        let mut stderr = std::io::stderr();
+        let sendable_state = conversation
+            .as_sendable_conversation_state(&os, &mut stderr, false)
+            .await
+            .unwrap();
+
+        // Verify continuation ID is included in the API request
+        assert_eq!(
+            sendable_state.agent_continuation_id.as_deref(),
+            Some(expected_continuation_id.as_str())
+        );
+    }
+}
diff --git a/crates/chat-cli/src/cli/chat/mod.rs b/crates/chat-cli/src/cli/chat/mod.rs
diff --git a/crates/chat-cli/src/cli/chat/parser.rs b/crates/chat-cli/src/cli/chat/parser.rs
