From 54d9506d5136278bc0dc1bbca268904458b20f83 Mon Sep 17 00:00:00 2001 From: Jwalin Shah Date: Tue, 5 May 2026 11:09:41 -0700 Subject: [PATCH] [codex] Preserve message metadata in session_raw transcripts (#1231) Co-authored-by: Jwalin Shah --- .../agent/harness/session/transcript.rs | 60 +++++++++++++- .../agent/harness/session/transcript_tests.rs | 83 +++++++++++++++++++ .../harness/subagent_runner/extract_tool.rs | 6 ++ src/openhuman/agent/multimodal.rs | 2 + src/openhuman/providers/compatible_tests.rs | 31 +++++++ src/openhuman/providers/router.rs | 4 + src/openhuman/providers/traits.rs | 12 +++ 7 files changed, 197 insertions(+), 1 deletion(-) diff --git a/src/openhuman/agent/harness/session/transcript.rs b/src/openhuman/agent/harness/session/transcript.rs index 9fca5dbff..410875782 100644 --- a/src/openhuman/agent/harness/session/transcript.rs +++ b/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, role: String, content: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + extra_metadata: Option, #[serde(skip_serializing_if = "Option::is_none")] model: Option, #[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 { match serde_json::from_str::(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 { 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 { + 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 = 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> { }; 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> { 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(); diff --git a/src/openhuman/agent/harness/session/transcript_tests.rs b/src/openhuman/agent/harness/session/transcript_tests.rs index 21c586a3a..a1d58d32e 100644 --- a/src/openhuman/agent/harness/session/transcript_tests.rs +++ b/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 diff --git a/src/openhuman/agent/harness/subagent_runner/extract_tool.rs b/src/openhuman/agent/harness/subagent_runner/extract_tool.rs index a4a89431a..fa2c327cd 100644 --- a/src/openhuman/agent/harness/subagent_runner/extract_tool.rs +++ b/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, }, ]; diff --git a/src/openhuman/agent/multimodal.rs b/src/openhuman/agent/multimodal.rs index c28c466a5..c9195efbf 100644 --- a/src/openhuman/agent/multimodal.rs +++ b/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(), }); } diff --git a/src/openhuman/providers/compatible_tests.rs b/src/openhuman/providers/compatible_tests.rs index 4fbd894dc..0e2cb3677 100644 --- a/src/openhuman/providers/compatible_tests.rs +++ b/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", diff --git a/src/openhuman/providers/router.rs b/src/openhuman/providers/router.rs index b12bd5205..c61416f4d 100644 --- a/src/openhuman/providers/router.rs +++ b/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"}})]; diff --git a/src/openhuman/providers/traits.rs b/src/openhuman/providers/traits.rs index ceb5394d4..e453a11a2 100644 --- a/src/openhuman/providers/traits.rs +++ b/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, pub role: String, pub content: String, + #[serde(default, skip_serializing)] + pub extra_metadata: Option, } impl ChatMessage { pub fn system(content: impl Into) -> Self { Self { + id: None, role: "system".into(), content: content.into(), + extra_metadata: None, } } pub fn user(content: impl Into) -> Self { Self { + id: None, role: "user".into(), content: content.into(), + extra_metadata: None, } } pub fn assistant(content: impl Into) -> Self { Self { + id: None, role: "assistant".into(), content: content.into(), + extra_metadata: None, } } pub fn tool(content: impl Into) -> Self { Self { + id: None, role: "tool".into(), content: content.into(), + extra_metadata: None, } } }