[codex] Preserve message metadata in session_raw transcripts (#1231)

Co-authored-by: Jwalin Shah <jshah1331@gmail.com>

src/openhuman/agent/harness/session/transcript.rs

@@ -46,7 +46,9 @@
 //! ```
 //!
 //! Only `role` and `content` are required. All other fields are optional.
-//! Unknown fields on read are ignored (forward-compat).
+//! UI-visible rows may also carry a stable `id` and `extra_metadata` so
+//! the session transcript can eventually replace the separate thread
+//! message log without losing message-level addressing.
 
 use crate::openhuman::providers::ChatMessage;
 use anyhow::{Context, Result};
@@ -138,8 +140,12 @@ struct MetaPayload {
 /// forward-compatibility.
 #[derive(Serialize, Deserialize)]
 struct MessageLine {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    id: Option<String>,
     role: String,
     content: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    extra_metadata: Option<serde_json::Value>,
     #[serde(skip_serializing_if = "Option::is_none")]
     model: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -203,16 +209,20 @@ pub fn write_transcript(
         // options together so there's no separate unwrap.
         let line = match (last_assistant_idx, last_assistant_turn_usage) {
             (Some(idx), Some(tu)) if idx == i => MessageLine {
+                id: msg.id.clone(),
                 role: msg.role.clone(),
                 content: msg.content.clone(),
+                extra_metadata: msg.extra_metadata.clone(),
                 model: Some(tu.model.clone()),
                 usage: Some(tu.usage.clone()),
                 ts: Some(tu.ts.clone()),
                 _extra: HashMap::new(),
             },
             _ => MessageLine {
+                id: msg.id.clone(),
                 role: msg.role.clone(),
                 content: msg.content.clone(),
+                extra_metadata: msg.extra_metadata.clone(),
                 model: None,
                 usage: None,
                 ts: None,
@@ -357,8 +367,10 @@ fn read_transcript_jsonl(path: &Path) -> Result<SessionTranscript> {
         match serde_json::from_str::<MessageLine>(line) {
             Ok(ml) => {
                 messages.push(ChatMessage {
+                    id: ml.id,
                     role: ml.role,
                     content: ml.content,
+                    extra_metadata: ml.extra_metadata,
                 });
             }
             Err(err) => {
@@ -387,6 +399,48 @@ fn read_transcript_jsonl(path: &Path) -> Result<SessionTranscript> {
     Ok(SessionTranscript { meta, messages })
 }
 
+/// Find the newest root `session_raw/*.jsonl` transcript whose metadata
+/// declares `thread_id`.
+///
+/// Root transcripts live directly under `session_raw/` and do not carry
+/// the `__` separator used for sub-agent siblings. This helper is the
+/// bridge PR-2 can use to route UI thread reads to the canonical root
+/// transcript without accidentally folding delegated worker transcripts
+/// into the main chat timeline.
+pub fn find_root_transcript_for_thread(workspace_dir: &Path, thread_id: &str) -> Option<PathBuf> {
+    let thread_id = thread_id.trim();
+    if thread_id.is_empty() {
+        return None;
+    }
+
+    let raw_dir = raw_session_dir(workspace_dir);
+    let entries = fs::read_dir(&raw_dir).ok()?;
+    let mut matches: Vec<PathBuf> = entries
+        .flatten()
+        .map(|entry| entry.path())
+        .filter(|path| {
+            path.extension().and_then(|s| s.to_str()) == Some("jsonl")
+                && path
+                    .file_stem()
+                    .and_then(|s| s.to_str())
+                    .is_some_and(|stem| !stem.contains("__"))
+        })
+        .filter(|path| match read_transcript(path) {
+            Ok(transcript) => transcript.meta.thread_id.as_deref() == Some(thread_id),
+            Err(err) => {
+                log::warn!(
+                    "[transcript] skipping unreadable root transcript candidate {}: {err}",
+                    path.display()
+                );
+                false
+            }
+        })
+        .collect();
+
+    matches.sort();
+    matches.pop()
+}
+
 // ── Path resolution ──────────────────────────────────────────────────
 
 /// Resolve a transcript path under `session_raw/{stem}.jsonl` — a
@@ -663,8 +717,10 @@ fn parse_legacy_messages(raw: &str) -> Result<Vec<ChatMessage>> {
             };
             let content = &raw[content_start..content_start + content_end_rel];
             messages.push(ChatMessage {
+                id: None,
                 role,
                 content: content.replace(LEGACY_MSG_CLOSE_ESCAPED, LEGACY_MSG_CLOSE),
+                extra_metadata: None,
             });
             search_from = content_start + content_end_rel + LEGACY_MSG_CLOSE.len();
             continue;
@@ -672,8 +728,10 @@ fn parse_legacy_messages(raw: &str) -> Result<Vec<ChatMessage>> {
 
         let content = &raw[content_start..content_start + content_end_rel];
         messages.push(ChatMessage {
+            id: None,
             role,
             content: content.replace(LEGACY_MSG_CLOSE_ESCAPED, LEGACY_MSG_CLOSE),
+            extra_metadata: None,
         });
 
         search_from = content_start + content_end_rel + close_tag.len();

src/openhuman/agent/harness/session/transcript_tests.rs

@@ -53,15 +53,56 @@ fn round_trip_produces_byte_identical_messages() {
 
     assert_eq!(loaded.messages.len(), messages.len());
     for (original, loaded) in messages.iter().zip(loaded.messages.iter()) {
+        assert_eq!(original.id, loaded.id, "id mismatch");
         assert_eq!(original.role, loaded.role, "role mismatch");
         assert_eq!(
             original.content, loaded.content,
             "content mismatch for role={}",
             original.role
         );
+        assert_eq!(
+            original.extra_metadata, loaded.extra_metadata,
+            "extra metadata mismatch for role={}",
+            original.role
+        );
     }
 }
 
+#[test]
+fn message_id_and_extra_metadata_round_trip() {
+    let dir = TempDir::new().unwrap();
+    let path = dir.path().join("message_identity.jsonl");
+    let mut messages = sample_messages();
+    messages[1].id = Some("msg_user_123".into());
+    messages[1].extra_metadata = Some(serde_json::json!({
+        "citations": [{"id": "mem-1", "label": "Memory"}],
+        "tool_call_id": "call-1"
+    }));
+    let meta = sample_meta();
+
+    write_transcript(&path, &messages, &meta, None).unwrap();
+
+    let loaded = read_transcript(&path).unwrap();
+    assert_eq!(loaded.messages[1].id.as_deref(), Some("msg_user_123"));
+    assert_eq!(
+        loaded.messages[1].extra_metadata,
+        Some(serde_json::json!({
+            "citations": [{"id": "mem-1", "label": "Memory"}],
+            "tool_call_id": "call-1"
+        }))
+    );
+
+    let raw = fs::read_to_string(&path).unwrap();
+    assert!(
+        raw.contains("\"id\":\"msg_user_123\""),
+        "message id should be persisted in JSONL"
+    );
+    assert!(
+        raw.contains("\"extra_metadata\""),
+        "extra metadata should be persisted in JSONL"
+    );
+}
+
 /// JSON encoding handles any delimiter natively, making the old
 /// HTML-comment escaping unnecessary. This test verifies that content
 /// containing the legacy closing delimiter round-trips correctly via
@@ -230,6 +271,48 @@ fn find_latest_picks_newest_keyed_stem_in_flat_dir() {
     assert!(latest.to_string_lossy().ends_with("1714999999_main.jsonl"));
 }
 
+#[test]
+fn find_root_transcript_for_thread_skips_subagent_siblings() {
+    let dir = TempDir::new().unwrap();
+    let raw_dir = dir.path().join("session_raw");
+    fs::create_dir_all(&raw_dir).unwrap();
+
+    let mut root_meta = sample_meta();
+    root_meta.thread_id = Some("thread-abc".into());
+    write_transcript(
+        &raw_dir.join("1714000000_orchestrator_thread-abc.jsonl"),
+        &sample_messages(),
+        &root_meta,
+        None,
+    )
+    .unwrap();
+
+    let mut newer_other_meta = sample_meta();
+    newer_other_meta.thread_id = Some("thread-other".into());
+    write_transcript(
+        &raw_dir.join("1714999999_orchestrator_thread-other.jsonl"),
+        &sample_messages(),
+        &newer_other_meta,
+        None,
+    )
+    .unwrap();
+
+    let mut subagent_meta = sample_meta();
+    subagent_meta.thread_id = Some("thread-abc".into());
+    write_transcript(
+        &raw_dir.join("1715000000_orchestrator_thread-abc__1715000100_worker.jsonl"),
+        &sample_messages(),
+        &subagent_meta,
+        None,
+    )
+    .unwrap();
+
+    let found = find_root_transcript_for_thread(dir.path(), "thread-abc").unwrap();
+    assert!(found
+        .to_string_lossy()
+        .ends_with("1714000000_orchestrator_thread-abc.jsonl"));
+}
+
 #[test]
 fn find_latest_falls_back_to_legacy_ddmmyyyy_raw_dir() {
     // Pre-migration transcript at session_raw/DDMMYYYY/main_*.jsonl

src/openhuman/agent/harness/subagent_runner/extract_tool.rs

@@ -441,16 +441,22 @@ fn write_extract_transcript(
 
     let messages = vec![
         ChatMessage {
+            id: None,
             role: "system".into(),
             content: system_prompt.to_string(),
+            extra_metadata: None,
         },
         ChatMessage {
+            id: None,
             role: "user".into(),
             content: user_prompt.to_string(),
+            extra_metadata: None,
         },
         ChatMessage {
+            id: None,
             role: "assistant".into(),
             content: assistant_text,
+            extra_metadata: None,
         },
     ];
 

src/openhuman/agent/multimodal.rs

@@ -159,8 +159,10 @@ pub async fn prepare_messages_for_provider(
 
         let content = compose_multimodal_message(&cleaned_text, &normalized_refs);
         normalized_messages.push(ChatMessage {
+            id: message.id.clone(),
             role: message.role.clone(),
             content,
+            extra_metadata: message.extra_metadata.clone(),
         });
     }
 

src/openhuman/providers/compatible_tests.rs

@@ -568,6 +568,35 @@ fn convert_messages_for_native_maps_tool_result_payload() {
     assert_eq!(converted[0].content.as_deref(), Some("done"));
 }
 
+#[test]
+fn chat_message_identity_metadata_is_not_provider_wire_payload() {
+    let message = ChatMessage {
+        id: Some("msg_123".to_string()),
+        role: "user".to_string(),
+        content: "hello".to_string(),
+        extra_metadata: Some(serde_json::json!({"citation": "mem-1"})),
+    };
+
+    let serialized = serde_json::to_value(&message).unwrap();
+
+    assert_eq!(
+        serialized.get("role").and_then(|v| v.as_str()),
+        Some("user")
+    );
+    assert_eq!(
+        serialized.get("content").and_then(|v| v.as_str()),
+        Some("hello")
+    );
+    assert!(
+        serialized.get("id").is_none(),
+        "provider ChatMessage serialization must not leak UI message ids"
+    );
+    assert!(
+        serialized.get("extra_metadata").is_none(),
+        "provider ChatMessage serialization must not leak UI metadata"
+    );
+}
+
 #[test]
 fn flatten_system_messages_merges_into_first_user() {
     let input = vec![
@@ -892,8 +921,10 @@ fn response_with_multiple_tool_calls() {
 async fn chat_with_tools_fails_without_key() {
     let p = make_provider("TestProvider", "https://example.com", None);
     let messages = vec![ChatMessage {
+        id: None,
         role: "user".to_string(),
         content: "hello".to_string(),
+        extra_metadata: None,
     }];
     let tools = vec![serde_json::json!({
         "type": "function",

src/openhuman/providers/router.rs

@@ -416,8 +416,10 @@ mod tests {
         );
 
         let messages = vec![ChatMessage {
+            id: None,
             role: "user".to_string(),
             content: "use tools".to_string(),
+            extra_metadata: None,
         }];
         let tools = vec![serde_json::json!({
             "type": "function",
@@ -447,8 +449,10 @@ mod tests {
         );
 
         let messages = vec![ChatMessage {
+            id: None,
             role: "user".to_string(),
             content: "reason about this".to_string(),
+            extra_metadata: None,
         }];
         let tools = vec![serde_json::json!({"type": "function", "function": {"name": "test"}})];
 

src/openhuman/providers/traits.rs

@@ -7,36 +7,48 @@ use std::fmt::Write;
 /// A single message in a conversation.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ChatMessage {
+    #[serde(default, skip_serializing)]
+    pub id: Option<String>,
     pub role: String,
     pub content: String,
+    #[serde(default, skip_serializing)]
+    pub extra_metadata: Option<serde_json::Value>,
 }
 
 impl ChatMessage {
     pub fn system(content: impl Into<String>) -> Self {
         Self {
+            id: None,
             role: "system".into(),
             content: content.into(),
+            extra_metadata: None,
         }
     }
 
     pub fn user(content: impl Into<String>) -> Self {
         Self {
+            id: None,
             role: "user".into(),
             content: content.into(),
+            extra_metadata: None,
         }
     }
 
     pub fn assistant(content: impl Into<String>) -> Self {
         Self {
+            id: None,
             role: "assistant".into(),
             content: content.into(),
+            extra_metadata: None,
         }
     }
 
     pub fn tool(content: impl Into<String>) -> Self {
         Self {
+            id: None,
             role: "tool".into(),
             content: content.into(),
+            extra_metadata: None,
         }
     }
 }