Add token counts, timestamps, and model to rollouts

Author: bl-ue

codex-rs/core/src/chat_completions.rs

@@ -43,7 +43,7 @@ pub(crate) async fn stream_chat_completions(
 
     for item in &prompt.input {
         match item {
-            ResponseItem::Message { role, content } => {
+            ResponseItem::Message { role, content, .. } => {
                 let mut text = String::new();
                 for c in content {
                     match c {
@@ -255,6 +255,8 @@ where
                     content: vec![ContentItem::OutputText {
                         text: content.to_string(),
                     }],
+                    token_usage: None,
+                    timestamp: Some(crate::models::generate_timestamp()),
                 };
 
                 let _ = tx_event.send(Ok(ResponseEvent::OutputItemDone(item))).await;
@@ -402,6 +404,8 @@ where
                             content: vec![crate::models::ContentItem::OutputText {
                                 text: std::mem::take(&mut this.cumulative),
                             }],
+                            token_usage: token_usage.clone(),
+                            timestamp: Some(crate::models::generate_timestamp()),
                         };
 
                         // Buffer Completed so it is returned *after* the aggregated message.

codex-rs/core/src/codex.rs

@@ -82,6 +82,7 @@ use crate::protocol::SandboxPolicy;
 use crate::protocol::SessionConfiguredEvent;
 use crate::protocol::Submission;
 use crate::protocol::TaskCompleteEvent;
+use crate::protocol::TokenUsage;
 use crate::rollout::RolloutRecorder;
 use crate::safety::SafetyCheck;
 use crate::safety::assess_command_safety;
@@ -1159,6 +1160,9 @@ async fn try_run_turn(
                 token_usage,
             } => {
                 if let Some(token_usage) = token_usage {
+                    // Attach token usage to the last assistant message in this turn
+                    attach_token_usage_to_last_assistant_message(&mut output, token_usage.clone());
+
                     sess.tx_event
                         .send(Event {
                             id: sub_id.to_string(),
@@ -2003,7 +2007,7 @@ fn format_exec_output(output: &str, exit_code: i32, duration: Duration) -> Strin
 
 fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<String> {
     responses.iter().rev().find_map(|item| {
-        if let ResponseItem::Message { role, content } = item {
+        if let ResponseItem::Message { role, content, .. } = item {
             if role == "assistant" {
                 content.iter().rev().find_map(|ci| {
                     if let ContentItem::OutputText { text } = ci {
@@ -2021,6 +2025,23 @@ fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<St
     })
 }
 
+fn attach_token_usage_to_last_assistant_message(
+    items: &mut [ProcessedResponseItem],
+    usage: TokenUsage,
+) {
+    for processed_item in items.iter_mut().rev() {
+        if let ResponseItem::Message {
+            role, token_usage, ..
+        } = &mut processed_item.item
+        {
+            if role == "assistant" && token_usage.is_none() {
+                *token_usage = Some(usage);
+                break;
+            }
+        }
+    }
+}
+
 /// See [`ConversationHistory`] for details.
 fn record_conversation_history(disable_response_storage: bool, wire_api: WireApi) -> bool {
     if disable_response_storage {

codex-rs/core/src/models.rs

@@ -5,8 +5,12 @@ use mcp_types::CallToolResult;
 use serde::Deserialize;
 use serde::Serialize;
 use serde::ser::Serializer;
+use time::OffsetDateTime;
+use time::format_description::FormatItem;
+use time::macros::format_description;
 
 use crate::protocol::InputItem;
+use crate::protocol::TokenUsage;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type", rename_all = "snake_case")]
@@ -39,6 +43,10 @@ pub enum ResponseItem {
     Message {
         role: String,
         content: Vec<ContentItem>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        token_usage: Option<TokenUsage>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        timestamp: Option<String>,
     },
     Reasoning {
         id: String,
@@ -78,7 +86,12 @@ pub enum ResponseItem {
 impl From<ResponseInputItem> for ResponseItem {
     fn from(item: ResponseInputItem) -> Self {
         match item {
-            ResponseInputItem::Message { role, content } => Self::Message { role, content },
+            ResponseInputItem::Message { role, content } => Self::Message {
+                role,
+                content,
+                token_usage: None,
+                timestamp: Some(generate_timestamp()),
+            },
             ResponseInputItem::FunctionCallOutput { call_id, output } => {
                 Self::FunctionCallOutput { call_id, output }
             }
@@ -222,6 +235,16 @@ impl std::ops::Deref for FunctionCallOutputPayload {
     }
 }
 
+/// Generate a timestamp string in the same format as session timestamps.
+/// Format: "YYYY-MM-DDTHH:MM:SS.sssZ" (ISO 8601 with millisecond precision in UTC)
+pub fn generate_timestamp() -> String {
+    let timestamp_format: &[FormatItem] =
+        format_description!("[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z");
+    OffsetDateTime::now_utc()
+        .format(timestamp_format)
+        .unwrap_or_else(|_| "1970-01-01T00:00:00.000Z".to_string())
+}
+
 #[cfg(test)]
 mod tests {
     #![allow(clippy::unwrap_used)]
@@ -260,6 +283,120 @@ mod tests {
         assert_eq!(v.get("output").unwrap().as_str().unwrap(), "bad");
     }
 
+    #[test]
+    fn message_with_token_usage_and_timestamp() {
+        let usage = TokenUsage {
+            input_tokens: 100,
+            output_tokens: 50,
+            total_tokens: 150,
+            cached_input_tokens: Some(25),
+            reasoning_output_tokens: None,
+        };
+
+        let timestamp = "2025-07-15T10:30:45.123Z".to_string();
+
+        let message = ResponseItem::Message {
+            role: "assistant".to_string(),
+            content: vec![ContentItem::OutputText {
+                text: "Hello".to_string(),
+            }],
+            token_usage: Some(usage.clone()),
+            timestamp: Some(timestamp.clone()),
+        };
+
+        // Test serialization
+        let json = serde_json::to_string(&message).unwrap();
+        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
+
+        assert_eq!(parsed["type"], "message");
+        assert_eq!(parsed["role"], "assistant");
+        assert_eq!(parsed["content"][0]["text"], "Hello");
+        assert_eq!(parsed["token_usage"]["input_tokens"], 100);
+        assert_eq!(parsed["token_usage"]["output_tokens"], 50);
+        assert_eq!(parsed["token_usage"]["total_tokens"], 150);
+        assert_eq!(parsed["timestamp"], timestamp);
+
+        // Test deserialization
+        let deserialized: ResponseItem = serde_json::from_str(&json).unwrap();
+        if let ResponseItem::Message {
+            role,
+            content,
+            token_usage: Some(token_usage),
+            timestamp: Some(ts),
+            ..
+        } = deserialized
+        {
+            assert_eq!(role, "assistant");
+            assert_eq!(content.len(), 1);
+            assert_eq!(token_usage.input_tokens, 100);
+            assert_eq!(token_usage.output_tokens, 50);
+            assert_eq!(ts, timestamp);
+        } else {
+            panic!("Expected Message with token_usage and timestamp");
+        }
+    }
+
+    #[test]
+    fn message_without_optional_fields() {
+        let message = ResponseItem::Message {
+            role: "user".to_string(),
+            content: vec![ContentItem::InputText {
+                text: "Hi".to_string(),
+            }],
+            token_usage: None,
+            timestamp: None,
+        };
+
+        // Test serialization - optional fields should be omitted
+        let json = serde_json::to_string(&message).unwrap();
+        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
+
+        assert_eq!(parsed["type"], "message");
+        assert_eq!(parsed["role"], "user");
+        assert!(parsed.get("token_usage").is_none());
+        assert!(parsed.get("timestamp").is_none());
+
+        // Test deserialization - should work with missing fields
+        let deserialized: ResponseItem = serde_json::from_str(&json).unwrap();
+        if let ResponseItem::Message {
+            role,
+            token_usage,
+            timestamp,
+            ..
+        } = deserialized
+        {
+            assert_eq!(role, "user");
+            assert!(token_usage.is_none());
+            assert!(timestamp.is_none());
+        } else {
+            panic!("Expected Message without optional fields");
+        }
+    }
+
+    #[test]
+    fn generate_timestamp_format() {
+        let timestamp = generate_timestamp();
+
+        // Should be valid ISO 8601 format: YYYY-MM-DDTHH:MM:SS.sssZ
+        let parts: Vec<&str> = timestamp.split('T').collect();
+        assert_eq!(parts.len(), 2);
+
+        let date_part = parts[0];
+        let time_part = parts[1];
+
+        // Date should be YYYY-MM-DD format
+        assert_eq!(date_part.len(), 10);
+        assert!(date_part.contains('-'));
+
+        // Time should end with Z and have milliseconds
+        assert!(time_part.ends_with('Z'));
+        assert!(time_part.contains('.'));
+
+        // Should be able to parse as a valid timestamp
+        assert!(timestamp.len() >= 20); // Minimum ISO format length
+        assert!(timestamp.len() <= 30); // Maximum reasonable length
+    }
+
     #[test]
     fn deserialize_shell_tool_call_params() {
         let json = r#"{

codex-rs/core/src/rollout.rs

@@ -28,6 +28,7 @@ struct SessionMeta {
     timestamp: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     instructions: Option<String>,
+    model: String,
 }
 
 /// Records all [`ResponseItem`]s for a session and flushes them to disk after
@@ -71,6 +72,7 @@ impl RolloutRecorder {
             timestamp,
             id: session_id.to_string(),
             instructions,
+            model: config.model.to_string(),
         };
 
         // A reasonably-sized bounded channel. If the buffer fills up the send