fix: prompt caching with timestamp included (#2560)

Author: brandonskiser

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

@@ -6,6 +6,7 @@ use std::collections::{
 use std::io::Write;
 use std::sync::atomic::Ordering;
 
+use chrono::Utc;
 use crossterm::style::Color;
 use crossterm::{
     execute,
@@ -188,7 +189,7 @@ impl ConversationState {
             let Prompt { role, content } = prompt;
             match role {
                 crate::mcp_client::Role::User => {
-                    let user_msg = UserMessage::new_prompt(content.to_string());
+                    let user_msg = UserMessage::new_prompt(content.to_string(), None);
                     candidate_user.replace(user_msg);
                 },
                 crate::mcp_client::Role::Assistant => {
@@ -231,7 +232,7 @@ impl ConversationState {
             input
         };
 
-        let msg = UserMessage::new_prompt(input);
+        let msg = UserMessage::new_prompt(input, Some(Utc::now()));
         self.next_message = Some(msg);
     }
 
@@ -320,14 +321,19 @@ impl ConversationState {
 
     pub fn add_tool_results_with_images(&mut self, tool_results: Vec<ToolUseResult>, images: Vec<ImageBlock>) {
         debug_assert!(self.next_message.is_none());
-        self.next_message = Some(UserMessage::new_tool_use_results_with_images(tool_results, images));
+        self.next_message = Some(UserMessage::new_tool_use_results_with_images(
+            tool_results,
+            images,
+            Some(Utc::now()),
+        ));
     }
 
     /// Sets the next user message with "cancelled" tool results.
     pub fn abandon_tool_use(&mut self, tools_to_be_abandoned: &[QueuedTool], deny_input: String) {
         self.next_message = Some(UserMessage::new_cancelled_tool_uses(
             Some(deny_input),
             tools_to_be_abandoned.iter().map(|t| t.id.as_str()),
+            Some(Utc::now()),
         ));
     }
 
@@ -502,7 +508,7 @@ impl ConversationState {
         }
 
         let conv_state = self.backend_conversation_state(os, false, &mut vec![]).await?;
-        let mut summary_message = Some(UserMessage::new_prompt(summary_content.clone()));
+        let mut summary_message = Some(UserMessage::new_prompt(summary_content.clone(), None));
 
         // Create the history according to the passed compact strategy.
         let mut history = conv_state.history.cloned().collect::<VecDeque<_>>();
@@ -531,7 +537,7 @@ impl ConversationState {
         Ok(FigConversationState {
             conversation_id: Some(self.conversation_id.clone()),
             user_input_message: summary_message
-                .unwrap_or(UserMessage::new_prompt(summary_content)) // should not happen
+                .unwrap_or(UserMessage::new_prompt(summary_content, None)) // should not happen
                 .into_user_input_message(self.model.clone(), &tools),
             history: Some(flatten_history(history.iter())),
         })
@@ -614,7 +620,7 @@ impl ConversationState {
 
         if !context_content.is_empty() {
             self.context_message_length = Some(context_content.len());
-            let user = UserMessage::new_prompt(context_content);
+            let user = UserMessage::new_prompt(context_content, None);
             let assistant = 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![HistoryEntry {
@@ -840,6 +846,7 @@ fn enforce_conversation_invariants(
                     debug!("abandoning tool results");
                     *next_message = Some(UserMessage::new_prompt(
                         "The conversation history has overflowed, clearing state".to_string(),
+                        None,
                     ));
                 }
             },
@@ -893,6 +900,7 @@ fn enforce_conversation_invariants(
             *user_msg = UserMessage::new_cancelled_tool_uses(
                 user_msg.prompt().map(|p| p.to_string()),
                 tool_uses.iter().map(|t| t.id.as_str()),
+                None,
             );
         }
     }

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

@@ -54,7 +54,7 @@ pub struct UserMessage {
     pub additional_context: String,
     pub env_context: UserEnvContext,
     pub content: UserMessageContent,
-    pub timestamp: DateTime<Utc>,
+    pub timestamp: Option<DateTime<Utc>>,
     pub images: Option<Vec<ImageBlock>>,
 }
 
@@ -107,20 +107,24 @@ impl UserMessageContent {
 impl UserMessage {
     /// Creates a new [UserMessage::Prompt], automatically detecting and adding the user's
     /// environment [UserEnvContext].
-    pub fn new_prompt(prompt: String) -> Self {
+    pub fn new_prompt(prompt: String, timestamp: Option<DateTime<Utc>>) -> Self {
         Self {
             images: None,
-            timestamp: Utc::now(),
+            timestamp,
             additional_context: String::new(),
             env_context: UserEnvContext::generate_new(),
             content: UserMessageContent::Prompt { prompt },
         }
     }
 
-    pub fn new_cancelled_tool_uses<'a>(prompt: Option<String>, tool_use_ids: impl Iterator<Item = &'a str>) -> Self {
+    pub fn new_cancelled_tool_uses<'a>(
+        prompt: Option<String>,
+        tool_use_ids: impl Iterator<Item = &'a str>,
+        timestamp: Option<DateTime<Utc>>,
+    ) -> Self {
         Self {
             images: None,
-            timestamp: Utc::now(),
+            timestamp,
             additional_context: String::new(),
             env_context: UserEnvContext::generate_new(),
             content: UserMessageContent::CancelledToolUses {
@@ -141,7 +145,7 @@ impl UserMessage {
     pub fn new_tool_use_results(results: Vec<ToolUseResult>) -> Self {
         Self {
             additional_context: String::new(),
-            timestamp: Utc::now(),
+            timestamp: None,
             env_context: UserEnvContext::generate_new(),
             content: UserMessageContent::ToolUseResults {
                 tool_use_results: results,
@@ -150,10 +154,14 @@ impl UserMessage {
         }
     }
 
-    pub fn new_tool_use_results_with_images(results: Vec<ToolUseResult>, images: Vec<ImageBlock>) -> Self {
+    pub fn new_tool_use_results_with_images(
+        results: Vec<ToolUseResult>,
+        images: Vec<ImageBlock>,
+        timestamp: Option<DateTime<Utc>>,
+    ) -> Self {
         Self {
             additional_context: String::new(),
-            timestamp: Utc::now(),
+            timestamp,
             env_context: UserEnvContext::generate_new(),
             content: UserMessageContent::ToolUseResults {
                 tool_use_results: results,
@@ -278,31 +286,37 @@ impl UserMessage {
         }
     }
 
-    /// Returns a formatted [String] containing [Self::additional_context] and [Self::prompt].
+    /// Returns a formatted [String] containing [Self::additional_context], [Self::timestamp], and
+    /// [Self::prompt].
     fn content_with_context(&self) -> String {
-        // Format the time with iso8601 format using Z, e.g. 2025-08-08T17:43:28.672Z
-        let timestamp = self.timestamp.to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
+        let mut content = String::new();
 
-        let prompt_with_timestamp = self.prompt().map(|p| {
-            format!(
-                "{}Current UTC time: {}{}{}{}{}",
+        if let Some(ts) = self.timestamp {
+            content.push_str(&format!(
+                "{}Current UTC time: {}{}\n",
                 CONTEXT_ENTRY_START_HEADER,
-                timestamp,
+                // Format the time with iso8601 format using Z, e.g. 2025-08-08T17:43:28.672Z
+                ts.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
                 CONTEXT_ENTRY_END_HEADER,
-                USER_ENTRY_START_HEADER,
-                p,
-                USER_ENTRY_END_HEADER
-            )
-        });
-
-        match (self.additional_context.is_empty(), prompt_with_timestamp) {
-            // Only add special delimiters if we have both a prompt and additional context
-            (false, Some(prompt)) => format!("{}\n{}", self.additional_context, prompt),
-            (true, Some(prompt)) => prompt,
-            _ => self.additional_context.clone(),
+            ));
+        }
+
+        if !self.additional_context.is_empty() {
+            content.push_str(&self.additional_context);
+            content.push('\n');
         }
-        .trim()
-        .to_string()
+
+        // Only add special delimiters around the user's prompt if there is no timestamp or
+        // additional context to add.
+        match (content.is_empty(), self.prompt()) {
+            (false, Some(p)) => {
+                content.push_str(&format!("{}{}{}", USER_ENTRY_START_HEADER, p, USER_ENTRY_END_HEADER));
+            },
+            (true, Some(p)) => content.push_str(p),
+            _ => (),
+        };
+
+        content.trim().to_string()
     }
 }
 
@@ -549,20 +563,44 @@ mod tests {
 
     #[test]
     fn test_user_input_message_timestamp_formatting() {
-        let msg = UserMessage::new_prompt("hello world".to_string());
+        const USER_PROMPT: &str = "hello world";
+
+        let msg = UserMessage::new_prompt(USER_PROMPT.to_string(), Some(Utc::now()));
+
+        let msgs = [
+            msg.clone().into_user_input_message(None, &HashMap::new()),
+            msg.clone().into_history_entry(),
+        ];
+
+        for m in msgs {
+            println!("checking {:?}", m);
+            assert!(m.content.contains(CONTEXT_ENTRY_START_HEADER));
+            assert!(m.content.contains("Current UTC time"));
+            assert!(m.content.contains(CONTEXT_ENTRY_END_HEADER));
+            assert!(m.content.contains(USER_ENTRY_START_HEADER));
+            assert!(m.content.contains(USER_PROMPT));
+            assert!(m.content.contains(USER_ENTRY_END_HEADER.trim()));
+        }
+    }
+
+    #[test]
+    fn test_user_input_message_without_context() {
+        const USER_PROMPT: &str = "hello world";
+
+        let msg = UserMessage::new_prompt(USER_PROMPT.to_string(), None);
 
         let msgs = [
             msg.clone().into_user_input_message(None, &HashMap::new()),
             msg.clone().into_history_entry(),
         ];
 
         for m in msgs {
-            m.content.contains(CONTEXT_ENTRY_START_HEADER);
-            m.content.contains("Current UTC time");
-            m.content.contains(CONTEXT_ENTRY_END_HEADER);
-            m.content.contains(USER_ENTRY_START_HEADER);
-            m.content.contains("hello world");
-            m.content.contains(USER_ENTRY_END_HEADER);
+            assert!(!m.content.contains(CONTEXT_ENTRY_START_HEADER));
+            assert!(!m.content.contains("Current UTC time"));
+            assert!(!m.content.contains(CONTEXT_ENTRY_END_HEADER));
+            assert!(!m.content.contains(USER_ENTRY_START_HEADER));
+            assert!(m.content.contains(USER_PROMPT));
+            assert!(!m.content.contains(USER_ENTRY_END_HEADER.trim()));
         }
     }
 }