feat: Adds continuation id logic (aws#3114)

* feat: New continuation id logic

* chore: Centralize metadata

* chore: Moves struct to right place

---------

Co-authored-by: Kenneth S. <kennvene@amazon.com>

Author: kennvene-aws

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();

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)]
@@ -542,7 +543,7 @@ impl TryFrom<AssistantResponseMessage> for amzn_qdeveloper_streaming_client::typ
 }
 
 #[non_exhaustive]
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum ChatResponseStream {
     AssistantResponseEvent {
         content: String,
@@ -563,6 +564,16 @@ pub enum ChatResponseStream {
         conversation_id: Option<String>,
         utterance_id: Option<String>,
     },
+    MetadataEvent {
+        total_tokens: Option<i32>,
+        uncached_input_tokens: Option<i32>,
+        output_tokens: Option<i32>,
+        cache_read_input_tokens: Option<i32>,
+        cache_write_input_tokens: Option<i32>,
+    },
+    MeteringEvent {
+        usage: Option<f64>,
+    },
     SupplementaryWebLinksEvent(()),
     ToolUseEvent {
         tool_use_id: String,
@@ -609,6 +620,18 @@ 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, .. },
+            ) => ChatResponseStream::MetadataEvent {
+                total_tokens: token_usage.as_ref().map(|t| t.total_tokens),
+                uncached_input_tokens: token_usage.as_ref().map(|t| t.uncached_input_tokens),
+                output_tokens: token_usage.as_ref().map(|t| t.output_tokens),
+                cache_read_input_tokens: token_usage.as_ref().and_then(|t| t.cache_read_input_tokens),
+                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::ChatResponseStream::ToolUseEvent(
                 amzn_codewhisperer_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -626,7 +649,7 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
             amzn_codewhisperer_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
                 ChatResponseStream::SupplementaryWebLinksEvent(())
             },
-            _ => ChatResponseStream::Unknown,
+            _other => ChatResponseStream::Unknown,
         }
     }
 }
@@ -665,6 +688,18 @@ 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, .. },
+            ) => ChatResponseStream::MetadataEvent {
+                total_tokens: token_usage.as_ref().map(|t| t.total_tokens),
+                uncached_input_tokens: token_usage.as_ref().map(|t| t.uncached_input_tokens),
+                output_tokens: token_usage.as_ref().map(|t| t.output_tokens),
+                cache_read_input_tokens: token_usage.as_ref().and_then(|t| t.cache_read_input_tokens),
+                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::ChatResponseStream::ToolUseEvent(
                 amzn_qdeveloper_streaming_client::types::ToolUseEvent {
                     tool_use_id,
@@ -682,7 +717,7 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
             amzn_qdeveloper_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
                 ChatResponseStream::SupplementaryWebLinksEvent(())
             },
-            _ => ChatResponseStream::Unknown,
+            _other => ChatResponseStream::Unknown,
         }
     }
 }
@@ -870,7 +905,8 @@ 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)
+            //TODO: Setup new origin.
+            .origin(amzn_codewhisperer_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }
@@ -884,7 +920,8 @@ 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)
+            //TODO: Setup new origin.
+            .origin(amzn_qdeveloper_streaming_client::types::Origin::AiEditor)
             .build()
             .expect("Failed to build UserInputMessage")
     }

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

@@ -378,14 +378,12 @@ async fn process_path(
 async fn add_file_to_context(os: &Os, path: &Path, context_files: &mut Vec<(String, String)>) -> Result<()> {
     let filename = path.to_string_lossy().to_string();
     let content = os.fs.read_to_string(path).await?;
-    
+
     // Check if this is a steering file that needs front matter filtering
-    if filename.contains(".kiro/steering") && filename.ends_with(".md") {
-        if !should_include_steering_file(&content)? {
-            return Ok(());
-        }
+    if filename.contains(".kiro/steering") && filename.ends_with(".md") && !should_include_steering_file(&content)? {
+        return Ok(());
     }
-    
+
     context_files.push((filename, content));
     Ok(())
 }
@@ -402,30 +400,30 @@ fn should_include_steering_file(content: &str) -> Result<bool> {
         // No front matter - include the file
         return Ok(true);
     }
-    
+
     // Find the end of the front matter
     let lines: Vec<&str> = content.lines().collect();
     let mut end_index = None;
-    
+
     for (i, line) in lines.iter().enumerate().skip(1) {
         if line.trim() == "---" {
             end_index = Some(i);
             break;
         }
     }
-    
+
     let end_index = match end_index {
         Some(idx) => idx,
         None => {
             // Malformed front matter - include the file
             return Ok(true);
-        }
+        },
     };
-    
+
     // Extract and parse the front matter
     let front_matter_lines = &lines[1..end_index];
     let front_matter_yaml = front_matter_lines.join("\n");
-    
+
     match serde_yaml::from_str::<FrontMatter>(&front_matter_yaml) {
         Ok(front_matter) => {
             match front_matter.inclusion.as_deref() {
@@ -435,11 +433,11 @@ fn should_include_steering_file(content: &str) -> Result<bool> {
                 None => Ok(true),               // No inclusion field - include
                 Some(_) => Ok(true),            // Unknown inclusion value - include
             }
-        }
+        },
         Err(_) => {
             // Failed to parse front matter - include the file
             Ok(true)
-        }
+        },
     }
 }
 

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

@@ -102,6 +102,57 @@ pub struct McpServerInfo {
     pub config: CustomToolConfig,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct UserTurnMetadata {
+    continuation_id: String,
+    /// [RequestMetadata] about the ongoing operation.
+    requests: Vec<RequestMetadata>,
+}
+
+impl Default for UserTurnMetadata {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Enum used to store metadata about user turns
+impl UserTurnMetadata {
+    pub fn new() -> Self {
+        Self {
+            continuation_id: uuid::Uuid::new_v4().to_string(),
+            requests: vec![],
+        }
+    }
+
+    pub fn continuation_id(&self) -> &str {
+        &self.continuation_id
+    }
+
+    pub fn add_request(&mut self, request: RequestMetadata) {
+        self.requests.push(request);
+    }
+
+    pub fn first_request(&self) -> Option<RequestMetadata> {
+        self.requests.first().cloned()
+    }
+
+    pub fn last_request(&self) -> Option<RequestMetadata> {
+        self.requests.last().cloned()
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &RequestMetadata> {
+        self.requests.iter()
+    }
+
+    pub fn first(&self) -> Option<&RequestMetadata> {
+        self.requests.first()
+    }
+
+    pub fn last(&self) -> Option<&RequestMetadata> {
+        self.requests.last()
+    }
+}
+
 /// Tracks state related to an ongoing conversation.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ConversationState {
@@ -149,6 +200,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>,
+    /// Metadata about the ongoing user turn operation
+    #[serde(default)]
+    pub user_turn_metadata: UserTurnMetadata,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -212,6 +266,7 @@ impl ConversationState {
             checkpoint_manager: None,
             mcp_enabled,
             tangent_state: None,
+            user_turn_metadata: UserTurnMetadata::new(),
         }
     }
 
@@ -378,6 +433,10 @@ impl ConversationState {
         self.next_message = None;
     }
 
+    pub fn current_continuation_id(&self) -> &str {
+        self.user_turn_metadata.continuation_id()
+    }
+
     pub async fn set_next_user_message(&mut self, input: String) {
         debug_assert!(self.next_message.is_none(), "next_message should not exist");
         if let Some(next_message) = self.next_message.as_ref() {
@@ -614,6 +673,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: Some(self.user_turn_metadata.continuation_id()),
         })
     }
 
@@ -719,6 +779,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: Some(self.user_turn_metadata.continuation_id().to_string()),
         })
     }
 
@@ -777,6 +838,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: Some(self.user_turn_metadata.continuation_id().to_string()),
         })
     }
 
@@ -992,6 +1054,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 +1070,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),
         })
     }