WIP

Author: brandonskiser

crates/chat-cli/build.rs

@@ -122,7 +122,7 @@ fn main() {
 
                 quote::quote!(
                     #[doc = #description]
-                    #[derive(Debug, Clone, PartialEq)]
+                    #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
                     #[non_exhaustive]
                     pub enum #name {
                         #(

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

@@ -49,7 +49,7 @@ pub struct CompactArgs {
 }
 
 impl CompactArgs {
-    pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
+    pub async fn execute(self, os: &mut Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
         let default = CompactStrategy::default();
         let prompt = if self.prompt.is_empty() {
             None

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

@@ -33,6 +33,7 @@ use super::message::{
     ToolUseResult,
     UserMessage,
 };
+use super::parser::RequestMetadata;
 use super::token_counter::{
     CharCount,
     CharCounter,
@@ -65,6 +66,8 @@ use crate::os::Os;
 const CONTEXT_ENTRY_START_HEADER: &str = "--- CONTEXT ENTRY BEGIN ---\n";
 const CONTEXT_ENTRY_END_HEADER: &str = "--- CONTEXT ENTRY END ---\n\n";
 
+pub type HistoryEntry = (UserMessage, AssistantMessage, Option<RequestMetadata>);
+
 /// Tracks state related to an ongoing conversation.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ConversationState {
@@ -73,7 +76,7 @@ pub struct ConversationState {
     /// The next user message to be sent as part of the conversation. Required to be [Some] before
     /// calling [Self::as_sendable_conversation_state].
     next_message: Option<UserMessage>,
-    history: VecDeque<(UserMessage, AssistantMessage)>,
+    history: VecDeque<HistoryEntry>,
     /// The range in the history sendable to the backend (start inclusive, end exclusive).
     valid_history_range: (usize, usize),
     /// Similar to history in that stores user and assistant responses, except that it is not used
@@ -90,7 +93,7 @@ pub struct ConversationState {
     /// Cached value representing the length of the user context message.
     context_message_length: Option<usize>,
     /// Stores the latest conversation summary created by /compact
-    latest_summary: Option<String>,
+    latest_summary: Option<(String, RequestMetadata)>,
     /// Model explicitly selected by the user in this conversation state via `/model`.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub model: Option<String>,
@@ -180,10 +183,10 @@ impl ConversationState {
     }
 
     pub fn latest_summary(&self) -> Option<&str> {
-        self.latest_summary.as_deref()
+        self.latest_summary.as_ref().map(|(s, _)| s.as_str())
     }
 
-    pub fn history(&self) -> &VecDeque<(UserMessage, AssistantMessage)> {
+    pub fn history(&self) -> &VecDeque<HistoryEntry> {
         &self.history
     }
 
@@ -219,7 +222,7 @@ impl ConversationState {
                 let asst = candidate_asst.take().unwrap();
                 let user = candidate_user.take().unwrap();
                 self.append_assistant_transcript(&asst);
-                self.history.push_back((user, asst));
+                self.history.push_back((user, asst, None));
             }
         }
         Some(last_msg.content.to_string())
@@ -251,12 +254,17 @@ impl ConversationState {
     }
 
     /// Sets the response message according to the currently set [Self::next_message].
-    pub fn push_assistant_message(&mut self, os: &mut Os, message: AssistantMessage) {
+    pub fn push_assistant_message(
+        &mut self,
+        os: &mut Os,
+        message: AssistantMessage,
+        request_metadata: Option<RequestMetadata>,
+    ) {
         debug_assert!(self.next_message.is_some(), "next_message should exist");
         let next_user_message = self.next_message.take().expect("next user message should exist");
 
         self.append_assistant_transcript(&message);
-        self.history.push_back((next_user_message, message));
+        self.history.push_back((next_user_message, message, request_metadata));
 
         if let Ok(cwd) = std::env::current_dir() {
             os.database.set_conversation_by_path(cwd, self).ok();
@@ -272,7 +280,21 @@ impl ConversationState {
     ///
     /// This is equivalent to `utterance_id` in the Q API.
     pub fn message_id(&self) -> Option<&str> {
-        self.history.back().and_then(|(_, msg)| msg.message_id())
+        self.history.back().and_then(|(_, msg, _)| msg.message_id())
+    }
+
+    pub fn latest_tool_use_ids(&self) -> Option<String> {
+        self.history
+            .back()
+            .and_then(|(_, assistant, _)| assistant.tool_uses())
+            .map(|tools| (tools.iter().map(|t| t.id.as_str()).collect::<Vec<_>>().join(",")))
+    }
+
+    pub fn latest_tool_use_names(&self) -> Option<String> {
+        self.history
+            .back()
+            .and_then(|(_, assistant, _)| assistant.tool_uses())
+            .map(|tools| (tools.iter().map(|t| t.name.as_str()).collect::<Vec<_>>().join(",")))
     }
 
     /// Updates the history so that, when non-empty, the following invariants are in place:
@@ -476,7 +498,7 @@ impl ConversationState {
                         FILTER OUT CHAT CONVENTIONS (greetings, offers to help, etc).".to_string()
             },
         };
-        if let Some(summary) = &self.latest_summary {
+        if let Some((summary, _)) = &self.latest_summary {
             summary_content.push_str("\n\n");
             summary_content.push_str(CONTEXT_ENTRY_START_HEADER);
             summary_content.push_str("This summary contains ALL relevant information from our previous conversation including tool uses, results, code analysis, and file operations. YOU MUST be sure to include this information when creating your summarization document.\n\n");
@@ -493,7 +515,7 @@ impl ConversationState {
         let mut history = conv_state.history.cloned().collect::<VecDeque<_>>();
         history.drain((history.len().saturating_sub(strategy.messages_to_exclude))..);
         if strategy.truncate_large_messages {
-            for (user_message, _) in &mut history {
+            for (user_message, _, _) in &mut history {
                 user_message.truncate_safe(strategy.max_message_length);
             }
         }
@@ -524,10 +546,15 @@ impl ConversationState {
 
     /// `strategy` - The [CompactStrategy] used for the corresponding
     /// [ConversationState::create_summary_request].
-    pub fn replace_history_with_summary(&mut self, summary: String, strategy: CompactStrategy) {
+    pub fn replace_history_with_summary(
+        &mut self,
+        summary: String,
+        strategy: CompactStrategy,
+        request_metadata: RequestMetadata,
+    ) {
         self.history
             .drain(..(self.history.len().saturating_sub(strategy.messages_to_exclude)));
-        self.latest_summary = Some(summary);
+        self.latest_summary = Some((summary, request_metadata));
     }
 
     pub fn current_profile(&self) -> Option<&str> {
@@ -550,10 +577,10 @@ impl ConversationState {
         &mut self,
         os: &Os,
         conversation_start_context: Option<String>,
-    ) -> (Option<Vec<(UserMessage, AssistantMessage)>>, Vec<(String, String)>) {
+    ) -> (Option<Vec<HistoryEntry>>, Vec<(String, String)>) {
         let mut context_content = String::new();
         let mut dropped_context_files = Vec::new();
-        if let Some(summary) = &self.latest_summary {
+        if let Some((summary, _)) = &self.latest_summary {
             context_content.push_str(CONTEXT_ENTRY_START_HEADER);
             context_content.push_str("This summary contains ALL relevant information from our previous conversation including tool uses, results, code analysis, and file operations. YOU MUST reference this information when answering questions and explicitly acknowledge specific details from the summary when they're relevant to the current question.\n\n");
             context_content.push_str("SUMMARY CONTENT:\n");
@@ -592,7 +619,7 @@ impl ConversationState {
             self.context_message_length = Some(context_content.len());
             let user_msg = UserMessage::new_prompt(context_content);
             let assistant_msg = AssistantMessage::new_response(None, "I will fully incorporate this information when generating my responses, and explicitly acknowledge relevant parts of the summary when answering questions.".into());
-            (Some(vec![(user_msg, assistant_msg)]), dropped_context_files)
+            (Some(vec![(user_msg, assistant_msg, None)]), dropped_context_files)
         } else {
             (None, dropped_context_files)
         }
@@ -647,11 +674,8 @@ impl ConversationState {
 ///
 /// This is intended to provide us ways to accurately assess the exact state that is sent to the
 /// model without having to needlessly clone and mutate [ConversationState] in strange ways.
-pub type BackendConversationState<'a> = BackendConversationStateImpl<
-    'a,
-    std::collections::vec_deque::Iter<'a, (UserMessage, AssistantMessage)>,
-    Option<Vec<(UserMessage, AssistantMessage)>>,
->;
+pub type BackendConversationState<'a> =
+    BackendConversationStateImpl<'a, std::collections::vec_deque::Iter<'a, HistoryEntry>, Option<Vec<HistoryEntry>>>;
 
 /// See [BackendConversationState]
 #[derive(Debug, Clone)]
@@ -665,13 +689,7 @@ pub struct BackendConversationStateImpl<'a, T, U> {
     pub model_id: Option<&'a str>,
 }
 
-impl
-    BackendConversationStateImpl<
-        '_,
-        std::collections::vec_deque::Iter<'_, (UserMessage, AssistantMessage)>,
-        Option<Vec<(UserMessage, AssistantMessage)>>,
-    >
-{
+impl BackendConversationStateImpl<'_, std::collections::vec_deque::Iter<'_, HistoryEntry>, Option<Vec<HistoryEntry>>> {
     fn into_fig_conversation_state(self) -> eyre::Result<FigConversationState> {
         let history = flatten_history(self.context_messages.unwrap_or_default().iter().chain(self.history));
         let user_input_message: UserInputMessage = self
@@ -695,7 +713,7 @@ impl
         // Count the chars used by the messages in the history.
         // this clone is cheap
         let history = self.history.clone();
-        for (user, assistant) in history {
+        for (user, assistant, _) in history {
             user_chars += *user.char_count();
             assistant_chars += *assistant.char_count();
         }
@@ -705,7 +723,7 @@ impl
             .context_messages
             .as_ref()
             .map(|v| {
-                v.iter().fold(0, |acc, (user, assistant)| {
+                v.iter().fold(0, |acc, (user, assistant, _)| {
                     acc + *user.char_count() + *assistant.char_count()
                 })
             })
@@ -730,9 +748,9 @@ pub struct ConversationSize {
 /// Converts a list of user/assistant message pairs into a flattened list of ChatMessage.
 fn flatten_history<'a, T>(history: T) -> Vec<ChatMessage>
 where
-    T: Iterator<Item = &'a (UserMessage, AssistantMessage)>,
+    T: Iterator<Item = &'a HistoryEntry>,
 {
-    history.fold(Vec::new(), |mut acc, (user, assistant)| {
+    history.fold(Vec::new(), |mut acc, (user, assistant, _)| {
         acc.push(ChatMessage::UserInputMessage(user.clone().into_history_entry()));
         acc.push(ChatMessage::AssistantResponseMessage(assistant.clone().into()));
         acc
@@ -774,7 +792,7 @@ fn format_hook_context<'a>(hook_results: impl IntoIterator<Item = &'a (Hook, Str
 }
 
 fn enforce_conversation_invariants(
-    history: &mut VecDeque<(UserMessage, AssistantMessage)>,
+    history: &mut VecDeque<HistoryEntry>,
     next_message: &mut Option<UserMessage>,
     tools: &HashMap<ToolOrigin, Vec<Tool>>,
 ) -> (usize, usize) {
@@ -791,7 +809,7 @@ fn enforce_conversation_invariants(
             .iter()
             .enumerate()
             .skip(1)
-            .find(|(_, (m, _))| -> bool { !m.has_tool_use_results() })
+            .find(|(_, (m, _, _))| -> bool { !m.has_tool_use_results() })
             .map(|v| v.0)
         {
             Some(i) => {
@@ -815,7 +833,7 @@ fn enforce_conversation_invariants(
 
     // If the first message contains tool results, then we add the results to the content field
     // instead. This is required to avoid validation errors.
-    if let Some((user, _)) = history.front_mut() {
+    if let Some((user, _, _)) = history.front_mut() {
         if user.has_tool_use_results() {
             user.replace_content_with_tool_use_results();
         }
@@ -828,7 +846,7 @@ fn enforce_conversation_invariants(
         history.range(valid_history_range.0..valid_history_range.1).last(),
     ) {
         (Some(next_message), prev_msg) if next_message.has_tool_use_results() => match prev_msg {
-            None | Some((_, AssistantMessage::Response { .. })) => {
+            None | Some((_, AssistantMessage::Response { .. }, _)) => {
                 next_message.replace_content_with_tool_use_results();
             },
             _ => (),
@@ -838,7 +856,7 @@ fn enforce_conversation_invariants(
 
     // If the last message from the assistant contains tool uses AND next_message is set, we need to
     // ensure that next_message contains tool results.
-    if let (Some((_, AssistantMessage::ToolUse { tool_uses, .. })), Some(user_msg)) = (
+    if let (Some((_, AssistantMessage::ToolUse { tool_uses, .. }, _)), Some(user_msg)) = (
         history.range(valid_history_range.0..valid_history_range.1).last(),
         next_message,
     ) {
@@ -859,7 +877,7 @@ fn enforce_conversation_invariants(
 }
 
 fn enforce_tool_use_history_invariants(
-    history: &mut VecDeque<(UserMessage, AssistantMessage)>,
+    history: &mut VecDeque<HistoryEntry>,
     tools: &HashMap<ToolOrigin, Vec<Tool>>,
 ) {
     let tool_names: HashSet<_> = tools
@@ -872,7 +890,7 @@ fn enforce_tool_use_history_invariants(
         .filter(|name| *name != DUMMY_TOOL_NAME)
         .collect();
 
-    for (_, assistant) in history {
+    for (_, assistant, _) in history {
         if let AssistantMessage::ToolUse { tool_uses, .. } = assistant {
             for tool_use in tool_uses {
                 if tool_names.contains(tool_use.name.as_str()) {
@@ -1028,7 +1046,7 @@ mod tests {
                 .await
                 .unwrap();
             assert_conversation_state_invariants(s, i);
-            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()));
+            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()), None);
             conversation.set_next_user_message(i.to_string()).await;
         }
     }
@@ -1065,6 +1083,7 @@ mod tests {
                     args: serde_json::Value::Null,
                     ..Default::default()
                 }]),
+                None,
             );
             conversation.add_tool_results(vec![ToolUseResult {
                 tool_use_id: "tool_id".to_string(),
@@ -1099,14 +1118,15 @@ mod tests {
                         args: serde_json::Value::Null,
                         ..Default::default()
                     }]),
+                    None,
                 );
                 conversation.add_tool_results(vec![ToolUseResult {
                     tool_use_id: "tool_id".to_string(),
                     content: vec![],
                     status: ToolResultStatus::Success,
                 }]);
             } else {
-                conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()));
+                conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()), None);
                 conversation.set_next_user_message(i.to_string()).await;
             }
         }
@@ -1147,7 +1167,7 @@ mod tests {
 
             assert_conversation_state_invariants(s, i);
 
-            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()));
+            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()), None);
             conversation.set_next_user_message(i.to_string()).await;
         }
     }
@@ -1206,7 +1226,7 @@ mod tests {
                 s.user_input_message.content
             );
 
-            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()));
+            conversation.push_assistant_message(&mut os, AssistantMessage::new_response(None, i.to_string()), None);
             conversation.set_next_user_message(i.to_string()).await;
         }
     }