diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 459ee8a..aa61eec 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -41,6 +41,7 @@ pub use translate::{ AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, }; +use anyhow::Context; use braintrust_sdk_rust::SpanComponents; use clap::{Args, ValueEnum}; use std::ffi::OsString; @@ -126,18 +127,27 @@ pub struct ImportArgs { /// Agent that produced the session. #[arg(value_enum)] pub source: ImportSource, - /// Codex or Claude Code session id shown by the agent's resume command. - pub session_id: String, + /// One or more session ids shown by the agent's resume command. + #[arg( + value_name = "SESSION_ID", + num_args = 1.., + required_unless_present = "all", + conflicts_with = "all" + )] + pub session_ids: Vec, + /// Import every locally discoverable session for this agent. + #[arg(long, conflicts_with = "session_ids")] + pub all: bool, /// Destination object reference, such as `project_logs:` or /// `experiment:`. - #[arg(value_name = "DESTINATION", conflicts_with = "parent")] + #[arg(long, value_name = "DESTINATION", conflicts_with = "parent")] pub destination: Option, /// Attach the imported session below an exported Braintrust span. #[arg(long, value_name = "SPAN_COMPONENTS", conflicts_with = "destination")] pub parent: Option, /// Keep following the transcript until Ctrl-C, importing new turns as the /// coding-agent session grows. - #[arg(long)] + #[arg(long, conflicts_with = "all")] pub attach: bool, } @@ -416,13 +426,27 @@ pub async fn run_import( opts: ServeOptions, mut config: Option, ) -> anyhow::Result<()> { + validate_import_selection(&args)?; let destination = args .parent .map(|components| wire::TraceDestination::ParentSpan { components }) .or(args.destination); apply_import_destination(&mut config, destination)?; - let file = transcript_import::resolve_transcript(&args.session_id, args.source)?; - import_transcript(&file, args.source, opts, config, args.attach).await + let files = transcript_import::resolve_transcripts(&args.session_ids, args.all, args.source)?; + if args.attach { + return import_transcript(&files[0], args.source, opts, config, true).await; + } + import_transcripts(&files, args.source, opts, config).await +} + +fn validate_import_selection(args: &ImportArgs) -> anyhow::Result<()> { + if args.all != args.session_ids.is_empty() { + anyhow::bail!("provide explicit session ids or use --all, but not both"); + } + if args.attach && args.session_ids.len() != 1 { + anyhow::bail!("--attach requires exactly one session id"); + } + Ok(()) } /// Launch a coding agent with inherited stdio and inject Braintrust hooks for @@ -705,6 +729,27 @@ pub async fn import_transcript( processor.finish().await } +/// Import multiple completed native transcripts through one processor. +/// +/// This shares translator and sink setup while keeping each native session's +/// correlation state isolated by session id. +pub async fn import_transcripts( + files: &[PathBuf], + source: ImportSource, + opts: ServeOptions, + config: Option, +) -> anyhow::Result<()> { + let mut processor = ImportProcessor::new(opts, config); + for file in files { + let mut tail = transcript_import::TranscriptTail::new(file.clone(), source); + let entries = tail + .poll(true) + .with_context(|| format!("import transcript {}", file.display()))?; + processor.process(entries).await?; + } + processor.finish().await +} + struct ImportLive { translator: Box, sink: Box, @@ -846,6 +891,55 @@ fn json_str_field(payload: &serde_json::Value, field: &str) -> Option { #[cfg(test)] mod tests { use super::*; + use clap::Parser; + + #[derive(Debug, Parser)] + struct ImportCli { + #[command(flatten)] + args: ImportArgs, + } + + #[test] + fn import_args_accept_multiple_sessions_or_all() { + let explicit = ImportCli::try_parse_from([ + "test", + "codex", + "session-a", + "session-b", + "--destination", + "project_logs:project-id", + ]) + .unwrap() + .args; + assert_eq!(explicit.session_ids, ["session-a", "session-b"]); + assert!(!explicit.all); + assert!(explicit.destination.is_some()); + + let all = ImportCli::try_parse_from(["test", "claude", "--all"]) + .unwrap() + .args; + assert!(all.session_ids.is_empty()); + assert!(all.all); + + assert!(ImportCli::try_parse_from(["test", "codex"]).is_err()); + assert!(ImportCli::try_parse_from(["test", "codex", "session-a", "--all"]).is_err()); + } + + #[test] + fn attach_requires_one_explicit_session() { + let args = ImportArgs { + source: ImportSource::Codex, + session_ids: vec!["one".into(), "two".into()], + all: false, + destination: None, + parent: None, + attach: true, + }; + assert!(validate_import_selection(&args) + .unwrap_err() + .to_string() + .contains("exactly one")); + } #[test] fn json_string_fields_accept_strings_and_numbers_only() { diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 474127a..5881c39 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -347,7 +347,8 @@ mod tests { let services = Arc::new(RecordingHost::new(None, Some("stop before lookup"))); let args = ImportArgs { source: ImportSource::Codex, - session_id: "00000000-0000-0000-0000-000000000000".into(), + session_ids: vec!["00000000-0000-0000-0000-000000000000".into()], + all: false, destination, parent: None, attach: false, diff --git a/bt-daemon/src/transcript_import.rs b/bt-daemon/src/transcript_import.rs deleted file mode 100644 index 648eda0..0000000 --- a/bt-daemon/src/transcript_import.rs +++ /dev/null @@ -1,794 +0,0 @@ -use crate::wire::Envelope; -use crate::ImportSource; -use anyhow::{bail, Context}; -use serde_json::{json, Value}; -use std::path::{Path, PathBuf}; - -pub(crate) fn resolve_transcript( - session_id: &str, - source: ImportSource, -) -> anyhow::Result { - validate_session_id(session_id)?; - resolve_transcript_in(session_id, source, &transcript_roots(source)) -} - -fn transcript_roots(source: ImportSource) -> Vec { - let home = std::env::var_os("HOME") - .or_else(|| std::env::var_os("USERPROFILE")) - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(".")); - match source { - ImportSource::Codex => { - let codex_home = std::env::var_os("CODEX_HOME") - .map(PathBuf::from) - .unwrap_or_else(|| home.join(".codex")); - vec![ - codex_home.join("sessions"), - codex_home.join("archived_sessions"), - ] - } - ImportSource::Claude => { - let claude_home = std::env::var_os("CLAUDE_CONFIG_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| home.join(".claude")); - vec![claude_home.join("projects")] - } - } -} - -fn resolve_transcript_in( - session_id: &str, - source: ImportSource, - roots: &[PathBuf], -) -> anyhow::Result { - validate_session_id(session_id)?; - let expected = match source { - ImportSource::Codex => format!("{session_id}.jsonl"), - ImportSource::Claude => format!("{session_id}.jsonl"), - }; - let mut matches = Vec::new(); - for root in roots { - find_matching_files(root, &expected, source, &mut matches); - } - matches.sort(); - matches.dedup(); - match matches.as_slice() { - [path] => Ok(path.clone()), - [] => { - let locations = roots - .iter() - .map(|root| root.display().to_string()) - .collect::>() - .join(", "); - bail!( - "no {} transcript found for session {session_id}; searched {locations}", - source_name(source) - ) - } - paths => { - let locations = paths - .iter() - .map(|path| path.display().to_string()) - .collect::>() - .join(", "); - bail!( - "multiple {} transcripts found for session {session_id}: {locations}", - source_name(source) - ) - } - } -} - -fn find_matching_files( - directory: &Path, - expected_suffix: &str, - source: ImportSource, - matches: &mut Vec, -) { - let Ok(entries) = std::fs::read_dir(directory) else { - return; - }; - for entry in entries.flatten() { - let Ok(file_type) = entry.file_type() else { - continue; - }; - if file_type.is_symlink() { - continue; - } - let path = entry.path(); - if file_type.is_dir() { - find_matching_files(&path, expected_suffix, source, matches); - continue; - } - let Some(name) = path.file_name().and_then(|name| name.to_str()) else { - continue; - }; - let is_match = match source { - // Codex prefixes rollout files with their timestamp. Claude names - // the main transcript exactly after the session id. - ImportSource::Codex => name.ends_with(expected_suffix), - ImportSource::Claude => name == expected_suffix, - }; - if is_match { - matches.push(path); - } - } -} - -fn validate_session_id(session_id: &str) -> anyhow::Result<()> { - if session_id.is_empty() - || !session_id - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) - { - bail!("invalid session id {session_id:?}"); - } - Ok(()) -} - -fn source_name(source: ImportSource) -> &'static str { - match source { - ImportSource::Codex => "Codex", - ImportSource::Claude => "Claude Code", - } -} - -pub(crate) fn transcript_envelopes( - path: &Path, - source: ImportSource, -) -> anyhow::Result> { - let contents = std::fs::read_to_string(path) - .with_context(|| format!("read transcript {}", path.display()))?; - let mut records = Vec::new(); - let mut record_end_offsets = Vec::new(); - let mut offset = 0_u64; - for (index, line) in contents.split_inclusive('\n').enumerate() { - offset += line.len() as u64; - if line.trim().is_empty() { - continue; - } - records.push( - serde_json::from_str(line).with_context(|| { - format!("parse transcript {} line {}", path.display(), index + 1) - })?, - ); - record_end_offsets.push(offset); - } - if records.is_empty() { - bail!("transcript {} is empty", path.display()); - } - match source { - ImportSource::Codex => codex_envelopes(path, &records), - ImportSource::Claude => { - claude_envelopes(path, &records, &record_end_offsets, contents.len() as u64) - } - } -} - -/// Incrementally converts a growing native transcript into synthetic hook -/// events for one persistent translator. The final poll closes the active -/// turn/session; ordinary polls keep the newest turn open. -pub(crate) struct TranscriptTail { - path: PathBuf, - source: ImportSource, - started: bool, - completed_turns: usize, - active_turn: Option, - codex_checkpoints: usize, - last_len: u64, -} - -impl TranscriptTail { - pub(crate) fn new(path: PathBuf, source: ImportSource) -> Self { - Self { - path, - source, - started: false, - completed_turns: 0, - active_turn: None, - codex_checkpoints: 0, - last_len: 0, - } - } - - pub(crate) fn poll(&mut self, finalize: bool) -> anyhow::Result> { - let events = match transcript_envelopes(&self.path, self.source) { - Ok(events) => events, - Err(_) if !finalize => return Ok(Vec::new()), - Err(error) => return Err(error), - }; - let len = std::fs::metadata(&self.path)?.len(); - match self.source { - ImportSource::Codex => self.poll_codex(events, len, finalize), - ImportSource::Claude => self.poll_claude(events, len, finalize), - } - } - - fn poll_codex( - &mut self, - events: Vec, - len: u64, - finalize: bool, - ) -> anyhow::Result> { - if events.len() < 2 { - bail!("Codex import did not produce session boundary events"); - } - let mut out = Vec::new(); - if !self.started { - out.push(events[0].clone()); - self.started = true; - } - let checkpoints = &events[1..events.len() - 1]; - out.extend(checkpoints.iter().skip(self.codex_checkpoints).cloned()); - self.codex_checkpoints = checkpoints.len(); - let mut tail = events.last().cloned().unwrap(); - if finalize { - out.push(tail); - } else if len != self.last_len { - tail.event = "ImportCheckpoint".into(); - if let Some(payload) = tail.payload.as_object_mut() { - payload.insert("hook_event_name".into(), json!("ImportCheckpoint")); - } - out.push(tail); - } - self.last_len = len; - Ok(out) - } - - fn poll_claude( - &mut self, - events: Vec, - len: u64, - finalize: bool, - ) -> anyhow::Result> { - if events.len() < 2 || !(events.len() - 2).is_multiple_of(2) { - bail!("Claude import did not produce turn boundary pairs"); - } - let mut out = Vec::new(); - if !self.started { - out.push(events[0].clone()); - self.started = true; - } - let turn_count = (events.len() - 2) / 2; - let completed_target = if finalize { - turn_count - } else { - turn_count.saturating_sub(1) - }; - while self.completed_turns < completed_target { - let turn = self.completed_turns; - if self.active_turn != Some(turn) { - out.push(events[1 + turn * 2].clone()); - } - out.push(events[2 + turn * 2].clone()); - self.completed_turns += 1; - self.active_turn = None; - } - if !finalize && turn_count > 0 { - let active = turn_count - 1; - if self.active_turn != Some(active) { - out.push(events[1 + active * 2].clone()); - self.active_turn = Some(active); - } - if len != self.last_len { - let mut checkpoint = events.last().cloned().unwrap(); - checkpoint.event = "ImportCheckpoint".into(); - if let Some(payload) = checkpoint.payload.as_object_mut() { - payload.insert("hook_event_name".into(), json!("ImportCheckpoint")); - } - out.push(checkpoint); - } - } - if finalize { - out.push(events.last().cloned().unwrap()); - } - self.last_len = len; - Ok(out) - } -} - -fn codex_envelopes(path: &Path, records: &[Value]) -> anyhow::Result> { - let meta = records - .iter() - .find(|record| record.get("type").and_then(Value::as_str) == Some("session_meta")); - let session_id = meta - .and_then(|record| record.pointer("/payload/id")) - .and_then(Value::as_str) - .map(str::to_owned) - .unwrap_or_else(|| file_session_id(path)); - let source_version = meta - .and_then(|record| record.pointer("/payload/cli_version")) - .and_then(Value::as_str) - .map(str::to_owned); - let (start_ms, end_ms) = timestamp_bounds(records); - let transcript_path = path.to_string_lossy(); - let last_message = records.iter().rev().find_map(|record| { - (record.pointer("/payload/type").and_then(Value::as_str) == Some("task_complete")) - .then(|| record.pointer("/payload/last_agent_message").cloned()) - .flatten() - }); - let mut events = vec![envelope( - "codex", - source_version.clone(), - &session_id, - "SessionStart", - start_ms, - json!({ - "session_id": session_id, - "hook_event_name": "SessionStart", - "transcript_path": transcript_path, - "source": "import", - "_bt_import_through_ms": start_ms - }), - )]; - let mut checkpoints = records - .iter() - .filter(|record| { - matches!( - record.pointer("/payload/type").and_then(Value::as_str), - Some("task_started" | "task_complete" | "turn_aborted") - ) - }) - .filter_map(timestamp_ms) - .collect::>(); - checkpoints.sort_unstable(); - checkpoints.dedup(); - for checkpoint_ms in checkpoints { - if checkpoint_ms <= start_ms || checkpoint_ms >= end_ms { - continue; - } - events.push(envelope( - "codex", - source_version.clone(), - &session_id, - "ImportCheckpoint", - checkpoint_ms, - json!({ - "session_id": session_id, - "hook_event_name": "ImportCheckpoint", - "transcript_path": transcript_path, - "_bt_import_through_ms": checkpoint_ms - }), - )); - } - events.push(envelope( - "codex", - source_version, - &session_id, - "Stop", - end_ms, - json!({ - "session_id": session_id, - "hook_event_name": "Stop", - "transcript_path": transcript_path, - "last_agent_message": last_message, - "_bt_import_through_ms": end_ms - }), - )); - Ok(events) -} - -fn claude_envelopes( - path: &Path, - records: &[Value], - record_end_offsets: &[u64], - transcript_len: u64, -) -> anyhow::Result> { - if records.len() != record_end_offsets.len() { - bail!("Claude transcript record offsets do not match parsed records"); - } - let session_id = records - .iter() - .find_map(|record| string_at(record, "/sessionId")) - .unwrap_or_else(|| file_session_id(path)); - let source_version = records - .iter() - .find_map(|record| string_at(record, "/version")); - let cwd = records.iter().find_map(|record| string_at(record, "/cwd")); - let (start_ms, end_ms) = timestamp_bounds(records); - let transcript_path = path.to_string_lossy().into_owned(); - let mut events = vec![claude_import_envelope( - source_version.clone(), - &session_id, - "SessionStart", - start_ms.saturating_sub(1), - json!({ - "session_id": session_id, - "hook_event_name": "SessionStart", - "transcript_path": transcript_path, - "cwd": cwd, - "source": "import" - }), - 0, - )]; - - let user_indexes: Vec = records - .iter() - .enumerate() - .filter_map(|(index, record)| is_real_user(record).then_some(index)) - .collect(); - for (turn, &index) in user_indexes.iter().enumerate() { - let next = user_indexes.get(turn + 1).copied().unwrap_or(records.len()); - let segment = &records[index..next]; - let turn_start = timestamp_ms(&records[index]).unwrap_or(start_ms); - // Claude can append queue bookkeeping ahead of older conversation - // records. Only message rows define the native turn's duration. - let turn_end = segment - .iter() - .filter(|record| { - matches!( - record.get("type").and_then(Value::as_str), - Some("user" | "assistant") - ) - }) - .filter_map(timestamp_ms) - .max() - .unwrap_or(turn_start); - let turn_cwd = segment - .iter() - .find_map(|record| string_at(record, "/cwd")) - .or_else(|| cwd.clone()); - let prompt = records[index] - .pointer("/message/content") - .cloned() - .unwrap_or(Value::Null); - events.push(claude_import_envelope( - source_version.clone(), - &session_id, - "UserPromptSubmit", - turn_start, - json!({ - "session_id": session_id, - "hook_event_name": "UserPromptSubmit", - "transcript_path": transcript_path, - "cwd": turn_cwd, - "prompt": prompt - }), - record_end_offsets[index], - )); - let error = last_assistant_error(segment); - let stop_event = if error.is_some() { - "StopFailure" - } else { - "Stop" - }; - events.push(claude_import_envelope( - source_version.clone(), - &session_id, - stop_event, - turn_end, - json!({ - "session_id": session_id, - "hook_event_name": stop_event, - "transcript_path": transcript_path, - "cwd": turn_cwd, - "last_assistant_message": last_assistant_text(segment), - "error": error - }), - record_end_offsets[next.saturating_sub(1)], - )); - } - events.push(claude_import_envelope( - source_version, - &session_id, - "SessionEnd", - end_ms.saturating_add(1), - json!({ - "session_id": session_id, - "hook_event_name": "SessionEnd", - "transcript_path": transcript_path, - "cwd": cwd, - "reason": "transcript_import" - }), - transcript_len, - )); - Ok(events) -} - -fn claude_import_envelope( - source_version: Option, - session_id: &str, - event: &str, - ts_ms: i64, - mut payload: Value, - through_offset: u64, -) -> Envelope { - if let Some(payload) = payload.as_object_mut() { - payload.insert("_bt_import_through_offset".into(), json!(through_offset)); - } - envelope( - "claude-code", - source_version, - session_id, - event, - ts_ms, - payload, - ) -} - -fn envelope( - source: &str, - source_version: Option, - session_id: &str, - event: &str, - ts_ms: i64, - payload: Value, -) -> Envelope { - Envelope { - source: source.into(), - source_version, - plugin_version: None, - session_id: session_id.into(), - event: event.into(), - ts_ms, - managed_run_id: None, - payload, - route: None, - config: None, - } -} - -fn is_real_user(record: &Value) -> bool { - if record.get("type").and_then(Value::as_str) != Some("user") { - return false; - } - !record - .pointer("/message/content") - .and_then(Value::as_array) - .is_some_and(|blocks| { - blocks - .iter() - .any(|block| block.get("type").and_then(Value::as_str) == Some("tool_result")) - }) -} - -fn last_assistant_text(records: &[Value]) -> Option { - records.iter().rev().find_map(|record| { - if record.get("type").and_then(Value::as_str) != Some("assistant") { - return None; - } - let content = record.pointer("/message/content")?; - if let Some(text) = content.as_str() { - return Some(json!(text)); - } - let text = content - .as_array()? - .iter() - .filter_map(|block| { - (block.get("type").and_then(Value::as_str) == Some("text")) - .then(|| block.get("text").and_then(Value::as_str)) - .flatten() - }) - .collect::>() - .join("\n"); - (!text.is_empty()).then(|| json!(text)) - }) -} - -fn last_assistant_error(records: &[Value]) -> Option { - records.iter().rev().find_map(|record| { - if record.get("type").and_then(Value::as_str) != Some("assistant") - || record.get("isApiErrorMessage").and_then(Value::as_bool) != Some(true) - { - return None; - } - string_at(record, "/error") - .or_else(|| { - last_assistant_text(std::slice::from_ref(record))? - .as_str() - .map(str::to_owned) - }) - .or_else(|| Some("Claude API error".into())) - }) -} - -fn timestamp_bounds(records: &[Value]) -> (i64, i64) { - let mut timestamps = records.iter().filter_map(timestamp_ms); - let Some(first) = timestamps.next() else { - return (0, 0); - }; - timestamps.fold((first, first), |(min, max), value| { - (min.min(value), max.max(value)) - }) -} - -fn timestamp_ms(record: &Value) -> Option { - chrono::DateTime::parse_from_rfc3339(record.get("timestamp")?.as_str()?) - .ok() - .map(|timestamp| timestamp.timestamp_millis()) -} - -fn string_at(record: &Value, pointer: &str) -> Option { - record.pointer(pointer)?.as_str().map(str::to_owned) -} - -fn file_session_id(path: &Path) -> String { - path.file_stem() - .and_then(|stem| stem.to_str()) - .filter(|stem| !stem.is_empty()) - .unwrap_or("imported-session") - .to_owned() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn finds_codex_rollout_by_session_suffix() { - let temp = tempfile::tempdir().unwrap(); - let root = temp.path().join("sessions"); - let transcript = root - .join("2026/07/31") - .join("rollout-2026-07-31T12-00-00-session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - std::fs::write(&transcript, "{}\n").unwrap(); - - assert_eq!( - resolve_transcript_in("session-123", ImportSource::Codex, &[root]).unwrap(), - transcript - ); - } - - #[test] - fn finds_only_exact_claude_session_filename() { - let temp = tempfile::tempdir().unwrap(); - let root = temp.path().join("projects"); - let project = root.join("-tmp-project"); - std::fs::create_dir_all(&project).unwrap(); - std::fs::write(project.join("prefix-session-123.jsonl"), "{}\n").unwrap(); - let transcript = project.join("session-123.jsonl"); - std::fs::write(&transcript, "{}\n").unwrap(); - - assert_eq!( - resolve_transcript_in("session-123", ImportSource::Claude, &[root]).unwrap(), - transcript - ); - } - - #[test] - fn rejects_unsafe_session_ids() { - let error = resolve_transcript_in( - "../session", - ImportSource::Codex, - &[PathBuf::from("unused")], - ) - .unwrap_err(); - assert!(error.to_string().contains("invalid session id")); - } - - #[test] - fn reports_missing_and_ambiguous_sessions() { - let temp = tempfile::tempdir().unwrap(); - let first = temp.path().join("first"); - let second = temp.path().join("second"); - std::fs::create_dir_all(&first).unwrap(); - std::fs::create_dir_all(&second).unwrap(); - - let missing = resolve_transcript_in( - "missing", - ImportSource::Claude, - &[first.clone(), second.clone()], - ) - .unwrap_err(); - assert!(missing.to_string().contains("no Claude Code transcript")); - - std::fs::write(first.join("duplicate.jsonl"), "{}\n").unwrap(); - std::fs::write(second.join("duplicate.jsonl"), "{}\n").unwrap(); - let ambiguous = - resolve_transcript_in("duplicate", ImportSource::Claude, &[first, second]).unwrap_err(); - assert!(ambiguous - .to_string() - .contains("multiple Claude Code transcripts")); - } - - #[test] - fn codex_import_adds_native_turn_checkpoints() { - let records = vec![ - json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"session-123"}}), - json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}}), - json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1"}}), - json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-2"}}), - json!({"timestamp":"2026-01-01T00:00:05Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-2"}}), - ]; - let events = codex_envelopes(Path::new("rollout.jsonl"), &records).unwrap(); - - assert_eq!(events.first().unwrap().event, "SessionStart"); - assert_eq!(events.last().unwrap().event, "Stop"); - assert_eq!( - events - .iter() - .filter(|event| event.event == "ImportCheckpoint") - .count(), - 3 - ); - } - - #[test] - fn codex_tail_keeps_session_open_until_final_poll() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("session-123.jsonl"); - let mut records = vec![ - json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"session-123"}}), - json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}}), - json!({"timestamp":"2026-01-01T00:00:03Z","type":"response_item","payload":{"type":"message","role":"assistant"}}), - ]; - let write = |records: &[Value]| { - std::fs::write( - &path, - records - .iter() - .map(Value::to_string) - .collect::>() - .join("\n"), - ) - .unwrap(); - }; - write(&records); - let mut tail = TranscriptTail::new(path.clone(), ImportSource::Codex); - let first = tail.poll(false).unwrap(); - assert_eq!(first.first().unwrap().event, "SessionStart"); - assert_eq!(first.last().unwrap().event, "ImportCheckpoint"); - assert!(first.iter().all(|event| event.event != "Stop")); - - records.push(json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1"}})); - write(&records); - let second = tail.poll(false).unwrap(); - assert!(second.iter().all(|event| event.event != "SessionStart")); - assert_eq!(second.last().unwrap().event, "ImportCheckpoint"); - assert_eq!(tail.poll(true).unwrap().last().unwrap().event, "Stop"); - } - - #[test] - fn claude_tail_closes_only_completed_turns() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("session-123.jsonl"); - let mut records = vec![ - json!({"type":"user","sessionId":"session-123","timestamp":"2026-01-01T00:00:01Z","message":{"content":"one"}}), - json!({"type":"assistant","sessionId":"session-123","timestamp":"2026-01-01T00:00:02Z","message":{"content":"answer one"}}), - ]; - let write = |records: &[Value]| { - std::fs::write( - &path, - records - .iter() - .map(Value::to_string) - .collect::>() - .join("\n"), - ) - .unwrap(); - }; - write(&records); - let mut tail = TranscriptTail::new(path.clone(), ImportSource::Claude); - let first = tail.poll(false).unwrap(); - assert_eq!( - first - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["SessionStart", "UserPromptSubmit", "ImportCheckpoint"] - ); - - records.extend([ - json!({"type":"user","sessionId":"session-123","timestamp":"2026-01-01T00:00:03Z","message":{"content":"two"}}), - json!({"type":"assistant","sessionId":"session-123","timestamp":"2026-01-01T00:00:04Z","message":{"content":"answer two"}}), - ]); - write(&records); - let second = tail.poll(false).unwrap(); - assert_eq!( - second - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["Stop", "UserPromptSubmit", "ImportCheckpoint"] - ); - assert_eq!( - tail.poll(true) - .unwrap() - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["Stop", "SessionEnd"] - ); - } -} diff --git a/bt-daemon/src/transcript_import/claude.rs b/bt-daemon/src/transcript_import/claude.rs new file mode 100644 index 0000000..6132f6b --- /dev/null +++ b/bt-daemon/src/transcript_import/claude.rs @@ -0,0 +1,434 @@ +use super::{ + envelope, file_session_id, find_jsonl_files, read_jsonl_records, string_at, timestamp_bounds, + timestamp_ms, validate_session_id, +}; +use crate::wire::Envelope; +use anyhow::bail; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::io::BufRead; +use std::path::{Path, PathBuf}; + +#[derive(Default)] +pub(super) struct Tail { + started: bool, + completed_turns: usize, + active_turn: Option, + last_len: u64, +} + +impl Tail { + pub(super) fn poll( + &mut self, + events: Vec, + len: u64, + finalize: bool, + ) -> anyhow::Result> { + if events.len() < 2 { + bail!("Claude import did not produce session boundary events"); + } + let mut out = Vec::new(); + if !self.started { + out.push(events[0].clone()); + self.started = true; + } + let middle = &events[1..events.len() - 1]; + let turn_starts = middle + .iter() + .enumerate() + .filter_map(|(index, event)| (event.event == "UserPromptSubmit").then_some(index)) + .collect::>(); + let turn_count = turn_starts.len(); + let completed_target = if finalize { + turn_count + } else { + turn_count.saturating_sub(1) + }; + while self.completed_turns < completed_target { + let turn = self.completed_turns; + let start = turn_starts[turn]; + let end = turn_starts.get(turn + 1).copied().unwrap_or(middle.len()); + let skip = usize::from(self.active_turn == Some(turn)); + out.extend(middle[start + skip..end].iter().cloned()); + self.completed_turns += 1; + self.active_turn = None; + } + if !finalize && turn_count > 0 { + let active = turn_count - 1; + if self.active_turn != Some(active) { + out.push(middle[turn_starts[active]].clone()); + self.active_turn = Some(active); + } + if len != self.last_len { + let mut checkpoint = events.last().cloned().unwrap(); + checkpoint.event = "ImportCheckpoint".into(); + if let Some(payload) = checkpoint.payload.as_object_mut() { + payload.insert("hook_event_name".into(), json!("ImportCheckpoint")); + } + out.push(checkpoint); + } + } + if finalize { + out.push(events.last().cloned().unwrap()); + } + self.last_len = len; + Ok(out) + } +} + +pub(super) fn transcript_session_id(path: &Path) -> Option { + let file = std::fs::File::open(path).ok()?; + for line in std::io::BufReader::new(file).lines().map_while(Result::ok) { + let Ok(record) = serde_json::from_str::(&line) else { + continue; + }; + let Some(session_id) = record.get("sessionId").and_then(Value::as_str) else { + continue; + }; + if validate_session_id(session_id).is_err() { + return None; + } + return filename_matches(path, session_id).then(|| session_id.to_owned()); + } + None +} + +pub(super) fn roots(home: &Path) -> Vec { + let claude_home = std::env::var_os("CLAUDE_CONFIG_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| home.join(".claude")); + vec![claude_home.join("projects")] +} + +pub(super) fn filename_matches(path: &Path, session_id: &str) -> bool { + path.file_name().and_then(|name| name.to_str()) == Some(&format!("{session_id}.jsonl")) +} + +pub(super) fn envelopes( + path: &Path, + records: &[Value], + record_end_offsets: &[u64], + transcript_len: u64, +) -> anyhow::Result> { + if records.len() != record_end_offsets.len() { + bail!("Claude transcript record offsets do not match parsed records"); + } + let session_id = records + .iter() + .find_map(|record| string_at(record, "/sessionId")) + .unwrap_or_else(|| file_session_id(path)); + let source_version = records + .iter() + .find_map(|record| string_at(record, "/version")); + let cwd = records.iter().find_map(|record| string_at(record, "/cwd")); + let (start_ms, end_ms) = timestamp_bounds(records); + let transcript_path = path.to_string_lossy().into_owned(); + let subagents = subagents(path, records)?; + let mut events = vec![import_envelope( + source_version.clone(), + &session_id, + "SessionStart", + start_ms.saturating_sub(1), + json!({ + "session_id": session_id, + "hook_event_name": "SessionStart", + "transcript_path": transcript_path, + "cwd": cwd, + "source": "import" + }), + 0, + )]; + + let user_indexes: Vec = records + .iter() + .enumerate() + .filter_map(|(index, record)| is_real_user(record).then_some(index)) + .collect(); + for (turn, &index) in user_indexes.iter().enumerate() { + let next = user_indexes.get(turn + 1).copied().unwrap_or(records.len()); + let segment = &records[index..next]; + let turn_start = timestamp_ms(&records[index]).unwrap_or(start_ms); + // Claude can append queue bookkeeping ahead of older conversation + // records. Only message rows define the native turn's duration. + let turn_end = segment + .iter() + .filter(|record| { + matches!( + record.get("type").and_then(Value::as_str), + Some("user" | "assistant") + ) + }) + .filter_map(timestamp_ms) + .max() + .unwrap_or(turn_start); + let turn_cwd = segment + .iter() + .find_map(|record| string_at(record, "/cwd")) + .or_else(|| cwd.clone()); + let prompt = records[index] + .pointer("/message/content") + .cloned() + .unwrap_or(Value::Null); + events.push(import_envelope( + source_version.clone(), + &session_id, + "UserPromptSubmit", + turn_start, + json!({ + "session_id": session_id, + "hook_event_name": "UserPromptSubmit", + "transcript_path": transcript_path, + "cwd": turn_cwd, + "prompt": prompt + }), + record_end_offsets[index], + )); + for subagent in subagents + .iter() + .filter(|subagent| subagent.record_index >= index && subagent.record_index < next) + { + events.push(import_envelope( + source_version.clone(), + &session_id, + "SubagentStart", + subagent.start_ms, + json!({ + "session_id": session_id, + "hook_event_name": "SubagentStart", + "transcript_path": transcript_path, + "agent_id": subagent.agent_id, + "agent_type": subagent.agent_type, + "agent_transcript_path": subagent.path + }), + record_end_offsets[subagent.start_index], + )); + events.push(import_envelope( + source_version.clone(), + &session_id, + "SubagentStop", + subagent.end_ms, + json!({ + "session_id": session_id, + "hook_event_name": "SubagentStop", + "transcript_path": transcript_path, + "agent_id": subagent.agent_id, + "agent_type": subagent.agent_type, + "agent_transcript_path": subagent.path, + "last_assistant_message": subagent.last_assistant_message + }), + record_end_offsets[subagent.record_index], + )); + } + let error = last_assistant_error(segment); + let stop_event = if error.is_some() { + "StopFailure" + } else { + "Stop" + }; + events.push(import_envelope( + source_version.clone(), + &session_id, + stop_event, + turn_end, + json!({ + "session_id": session_id, + "hook_event_name": stop_event, + "transcript_path": transcript_path, + "cwd": turn_cwd, + "last_assistant_message": last_assistant_text(segment), + "error": error + }), + record_end_offsets[next.saturating_sub(1)], + )); + } + events.push(import_envelope( + source_version, + &session_id, + "SessionEnd", + end_ms.saturating_add(1), + json!({ + "session_id": session_id, + "hook_event_name": "SessionEnd", + "transcript_path": transcript_path, + "cwd": cwd, + "reason": "transcript_import" + }), + transcript_len, + )); + Ok(events) +} + +struct Subagent { + agent_id: String, + agent_type: Option, + path: String, + start_index: usize, + record_index: usize, + start_ms: i64, + end_ms: i64, + last_assistant_message: Option, +} + +fn subagents(path: &Path, records: &[Value]) -> anyhow::Result> { + let Some(stem) = path.file_stem() else { + return Ok(Vec::new()); + }; + let directory = path.with_file_name(stem).join("subagents"); + let mut paths = Vec::new(); + find_jsonl_files(&directory, &mut paths); + paths.sort(); + + let mut result_records = HashMap::)>::new(); + let mut calls = HashMap::)>::new(); + for (index, record) in records.iter().enumerate() { + if let Some(agent_id) = string_at(record, "/toolUseResult/agentId") { + let call_id = record + .pointer("/message/content") + .and_then(Value::as_array) + .and_then(|blocks| { + blocks.iter().find_map(|block| { + (block.get("type").and_then(Value::as_str) == Some("tool_result")) + .then(|| string_at(block, "/tool_use_id")) + .flatten() + }) + }); + result_records.insert(agent_id, (index, call_id)); + } + let Some(blocks) = record.pointer("/message/content").and_then(Value::as_array) else { + continue; + }; + for block in blocks { + if block.get("type").and_then(Value::as_str) != Some("tool_use") + || block.get("name").and_then(Value::as_str) != Some("Agent") + { + continue; + } + let Some(call_id) = string_at(block, "/id") else { + continue; + }; + calls.insert(call_id, (index, string_at(block, "/input/subagent_type"))); + } + } + + let fallback_index = records.len().saturating_sub(1); + paths + .into_iter() + .filter_map(|subagent_path| { + let name = subagent_path.file_stem()?.to_str()?; + let agent_id = name.strip_prefix("agent-")?.to_owned(); + let (record_index, call_id) = result_records + .get(&agent_id) + .cloned() + .unwrap_or((fallback_index, None)); + let (start_index, agent_type) = call_id + .as_ref() + .and_then(|call_id| calls.get(call_id)) + .cloned() + .unwrap_or((record_index, None)); + Some(( + subagent_path, + agent_id, + start_index, + record_index, + agent_type, + )) + }) + .map( + |(subagent_path, agent_id, start_index, record_index, agent_type)| { + let subagent_records = read_jsonl_records(&subagent_path)?; + let (child_start, child_end) = timestamp_bounds(&subagent_records); + let start_ms = timestamp_ms(&records[start_index]) + .unwrap_or(child_start) + .min(child_start); + let end_ms = timestamp_ms(&records[record_index]) + .unwrap_or(child_end) + .max(child_end); + Ok(Subagent { + agent_id, + agent_type, + path: subagent_path.to_string_lossy().into_owned(), + start_index, + record_index, + start_ms, + end_ms, + last_assistant_message: last_assistant_text(&subagent_records), + }) + }, + ) + .collect() +} + +fn import_envelope( + source_version: Option, + session_id: &str, + event: &str, + ts_ms: i64, + mut payload: Value, + through_offset: u64, +) -> Envelope { + if let Some(payload) = payload.as_object_mut() { + payload.insert("_bt_import_through_offset".into(), json!(through_offset)); + } + envelope( + "claude-code", + source_version, + session_id, + event, + ts_ms, + payload, + ) +} + +fn is_real_user(record: &Value) -> bool { + if record.get("type").and_then(Value::as_str) != Some("user") { + return false; + } + !record + .pointer("/message/content") + .and_then(Value::as_array) + .is_some_and(|blocks| { + blocks + .iter() + .any(|block| block.get("type").and_then(Value::as_str) == Some("tool_result")) + }) +} + +fn last_assistant_text(records: &[Value]) -> Option { + records.iter().rev().find_map(|record| { + if record.get("type").and_then(Value::as_str) != Some("assistant") { + return None; + } + let content = record.pointer("/message/content")?; + if let Some(text) = content.as_str() { + return Some(json!(text)); + } + let text = content + .as_array()? + .iter() + .filter_map(|block| { + (block.get("type").and_then(Value::as_str) == Some("text")) + .then(|| block.get("text").and_then(Value::as_str)) + .flatten() + }) + .collect::>() + .join("\n"); + (!text.is_empty()).then(|| json!(text)) + }) +} + +fn last_assistant_error(records: &[Value]) -> Option { + records.iter().rev().find_map(|record| { + if record.get("type").and_then(Value::as_str) != Some("assistant") + || record.get("isApiErrorMessage").and_then(Value::as_bool) != Some(true) + { + return None; + } + string_at(record, "/error") + .or_else(|| { + last_assistant_text(std::slice::from_ref(record))? + .as_str() + .map(str::to_owned) + }) + .or_else(|| Some("Claude API error".into())) + }) +} diff --git a/bt-daemon/src/transcript_import/codex.rs b/bt-daemon/src/transcript_import/codex.rs new file mode 100644 index 0000000..09f5895 --- /dev/null +++ b/bt-daemon/src/transcript_import/codex.rs @@ -0,0 +1,371 @@ +use super::{ + envelope, file_session_id, find_jsonl_files, read_jsonl_records, string_at, timestamp_bounds, + timestamp_ms, validate_session_id, +}; +use crate::wire::Envelope; +use anyhow::bail; +use serde_json::{json, Value}; +use std::collections::{HashMap, HashSet}; +use std::io::BufRead; +use std::path::{Path, PathBuf}; + +#[derive(Default)] +pub(super) struct Tail { + started: bool, + checkpoints: usize, + last_len: u64, +} + +impl Tail { + pub(super) fn poll( + &mut self, + events: Vec, + len: u64, + finalize: bool, + ) -> anyhow::Result> { + if events.len() < 2 { + bail!("Codex import did not produce session boundary events"); + } + let mut out = Vec::new(); + if !self.started { + out.push(events[0].clone()); + self.started = true; + } + let middle = &events[1..events.len() - 1]; + let checkpoints = middle + .iter() + .filter(|event| event.event == "ImportCheckpoint") + .collect::>(); + out.extend( + checkpoints + .iter() + .skip(self.checkpoints) + .map(|event| (*event).clone()), + ); + self.checkpoints = checkpoints.len(); + let mut tail = events.last().cloned().unwrap(); + if finalize { + out.extend( + middle + .iter() + .filter(|event| event.event != "ImportCheckpoint") + .cloned(), + ); + out.push(tail); + } else if len != self.last_len { + tail.event = "ImportCheckpoint".into(); + if let Some(payload) = tail.payload.as_object_mut() { + payload.insert("hook_event_name".into(), json!("ImportCheckpoint")); + } + out.push(tail); + } + self.last_len = len; + Ok(out) + } +} + +pub(super) fn transcript_session_id(path: &Path) -> Option { + let file = std::fs::File::open(path).ok()?; + for line in std::io::BufReader::new(file).lines().map_while(Result::ok) { + let Ok(record) = serde_json::from_str::(&line) else { + continue; + }; + if record.get("type").and_then(Value::as_str) != Some("session_meta") { + continue; + } + if record.pointer("/payload/source/subagent").is_some() { + return None; + } + let session_id = record.pointer("/payload/id").and_then(Value::as_str)?; + if validate_session_id(session_id).is_err() { + return None; + } + return filename_matches(path, session_id).then(|| session_id.to_owned()); + } + None +} + +pub(super) fn roots(home: &Path) -> Vec { + let codex_home = std::env::var_os("CODEX_HOME") + .map(PathBuf::from) + .unwrap_or_else(|| home.join(".codex")); + vec![ + codex_home.join("sessions"), + codex_home.join("archived_sessions"), + ] +} + +pub(super) fn filename_matches(path: &Path, session_id: &str) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(&format!("{session_id}.jsonl"))) +} + +pub(super) fn envelopes(path: &Path, records: &[Value]) -> anyhow::Result> { + let meta = records + .iter() + .find(|record| record.get("type").and_then(Value::as_str) == Some("session_meta")); + let session_id = meta + .and_then(|record| record.pointer("/payload/id")) + .and_then(Value::as_str) + .map(str::to_owned) + .unwrap_or_else(|| file_session_id(path)); + let source_version = meta + .and_then(|record| record.pointer("/payload/cli_version")) + .and_then(Value::as_str) + .map(str::to_owned); + let (start_ms, end_ms) = timestamp_bounds(records); + let transcript_path = path.to_string_lossy(); + let last_message = last_message(records); + let mut events = vec![envelope( + "codex", + source_version.clone(), + &session_id, + "SessionStart", + start_ms, + json!({ + "session_id": session_id, + "hook_event_name": "SessionStart", + "transcript_path": transcript_path, + "source": "import", + "_bt_import_through_ms": start_ms + }), + )]; + let mut checkpoints = records + .iter() + .filter(|record| { + matches!( + record.pointer("/payload/type").and_then(Value::as_str), + Some("task_started" | "task_complete" | "turn_aborted") + ) + }) + .filter_map(timestamp_ms) + .collect::>(); + checkpoints.sort_unstable(); + checkpoints.dedup(); + for checkpoint_ms in checkpoints { + if checkpoint_ms <= start_ms || checkpoint_ms >= end_ms { + continue; + } + events.push(envelope( + "codex", + source_version.clone(), + &session_id, + "ImportCheckpoint", + checkpoint_ms, + json!({ + "session_id": session_id, + "hook_event_name": "ImportCheckpoint", + "transcript_path": transcript_path, + "_bt_import_through_ms": checkpoint_ms + }), + )); + } + let mut visited = HashSet::new(); + visited.insert(session_id.clone()); + append_subagent_events( + &mut events, + path, + records, + None, + &session_id, + source_version.clone(), + &mut visited, + )?; + events.push(envelope( + "codex", + source_version, + &session_id, + "Stop", + end_ms, + json!({ + "session_id": session_id, + "hook_event_name": "Stop", + "transcript_path": transcript_path, + "last_agent_message": last_message, + "_bt_import_through_ms": end_ms + }), + )); + Ok(events) +} + +fn append_subagent_events( + events: &mut Vec, + parent_path: &Path, + parent_records: &[Value], + parent_agent_id: Option<&str>, + root_session_id: &str, + source_version: Option, + visited: &mut HashSet, +) -> anyhow::Result<()> { + let calls = spawn_calls(parent_records); + if calls.is_empty() { + return Ok(()); + } + let search_root = transcript_search_root(parent_path); + for call in calls { + if !visited.insert(call.agent_id.clone()) { + continue; + } + let Some(child_path) = find_transcript_by_id(&search_root, &call.agent_id) else { + continue; + }; + let child_records = read_jsonl_records(&child_path)?; + let (child_start, child_end) = timestamp_bounds(&child_records); + let parent_transcript_path = parent_path.to_string_lossy().into_owned(); + let child_transcript_path = child_path.to_string_lossy().into_owned(); + let mut post_payload = json!({ + "session_id": root_session_id, + "hook_event_name": "PostToolUse", + "transcript_path": parent_transcript_path, + "tool_name": "spawn_agent", + "tool_use_id": call.call_id, + "tool_response": { "agent_id": call.agent_id }, + "_bt_import_through_ms": call.result_ms + }); + if let (Some(parent_agent_id), Some(payload)) = + (parent_agent_id, post_payload.as_object_mut()) + { + payload.insert("agent_id".into(), json!(parent_agent_id)); + } + events.push(envelope( + "codex", + source_version.clone(), + root_session_id, + "PostToolUse", + call.result_ms, + post_payload, + )); + events.push(envelope( + "codex", + source_version.clone(), + root_session_id, + "SubagentStart", + child_start.min(call.start_ms), + json!({ + "session_id": root_session_id, + "hook_event_name": "SubagentStart", + "agent_id": call.agent_id, + "agent_type": call.agent_type, + "transcript_path": child_transcript_path, + "_bt_import_through_ms": child_start + }), + )); + append_subagent_events( + events, + &child_path, + &child_records, + Some(&call.agent_id), + root_session_id, + source_version.clone(), + visited, + )?; + events.push(envelope( + "codex", + source_version.clone(), + root_session_id, + "SubagentStop", + child_end.max(call.result_ms), + json!({ + "session_id": root_session_id, + "hook_event_name": "SubagentStop", + "agent_id": call.agent_id, + "agent_transcript_path": child_transcript_path, + "last_agent_message": last_message(&child_records), + "_bt_import_through_ms": child_end + }), + )); + } + Ok(()) +} + +struct SpawnCall { + call_id: String, + agent_id: String, + agent_type: Option, + start_ms: i64, + result_ms: i64, +} + +fn spawn_calls(records: &[Value]) -> Vec { + let mut calls = HashMap::, i64)>::new(); + let mut spawned = Vec::new(); + for record in records { + let Some(payload) = record.get("payload") else { + continue; + }; + let Some(call_id) = payload.get("call_id").and_then(Value::as_str) else { + continue; + }; + if payload.get("type").and_then(Value::as_str) == Some("function_call") + && payload.get("name").and_then(Value::as_str) == Some("spawn_agent") + { + let arguments = payload + .get("arguments") + .and_then(Value::as_str) + .and_then(|value| serde_json::from_str::(value).ok()); + let agent_type = arguments + .as_ref() + .and_then(|value| string_at(value, "/agent_type")); + calls.insert( + call_id.to_owned(), + (agent_type, timestamp_ms(record).unwrap_or(0)), + ); + continue; + } + if payload.get("type").and_then(Value::as_str) != Some("function_call_output") { + continue; + } + let Some((agent_type, start_ms)) = calls.get(call_id).cloned() else { + continue; + }; + let output = payload + .get("output") + .and_then(Value::as_str) + .and_then(|value| serde_json::from_str::(value).ok()); + let Some(agent_id) = output + .as_ref() + .and_then(|value| string_at(value, "/agent_id")) + else { + continue; + }; + spawned.push(SpawnCall { + call_id: call_id.to_owned(), + agent_id, + agent_type, + start_ms, + result_ms: timestamp_ms(record).unwrap_or(start_ms), + }); + } + spawned +} + +fn transcript_search_root(path: &Path) -> PathBuf { + path.ancestors() + .find(|ancestor| { + matches!( + ancestor.file_name().and_then(|name| name.to_str()), + Some("sessions" | "archived_sessions") + ) + }) + .map(Path::to_path_buf) + .or_else(|| path.parent().map(Path::to_path_buf)) + .unwrap_or_else(|| PathBuf::from(".")) +} + +fn find_transcript_by_id(root: &Path, session_id: &str) -> Option { + let mut candidates = Vec::new(); + find_jsonl_files(root, &mut candidates); + candidates.sort(); + candidates + .into_iter() + .find(|path| filename_matches(path, session_id)) +} + +fn last_message(records: &[Value]) -> Option { + records.iter().rev().find_map(|record| { + (record.pointer("/payload/type").and_then(Value::as_str) == Some("task_complete")) + .then(|| record.pointer("/payload/last_agent_message").cloned()) + .flatten() + }) +} diff --git a/bt-daemon/src/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs new file mode 100644 index 0000000..e43f5e1 --- /dev/null +++ b/bt-daemon/src/transcript_import/mod.rs @@ -0,0 +1,659 @@ +use crate::wire::Envelope; +use crate::ImportSource; +use anyhow::{bail, Context}; +use serde_json::Value; +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +mod claude; +mod codex; + +pub(crate) fn resolve_transcripts( + session_ids: &[String], + all: bool, + source: ImportSource, +) -> anyhow::Result> { + match (all, session_ids.is_empty()) { + (true, true) => discover_transcripts(source), + (false, false) => session_ids + .iter() + .map(|session_id| resolve_transcript(session_id, source)) + .collect(), + (true, false) => bail!("--all cannot be combined with explicit session ids"), + (false, true) => bail!("provide at least one session id or use --all"), + } +} + +pub(crate) fn resolve_transcript( + session_id: &str, + source: ImportSource, +) -> anyhow::Result { + validate_session_id(session_id)?; + resolve_transcript_in(session_id, source, &transcript_roots(source)) +} + +fn transcript_roots(source: ImportSource) -> Vec { + let home = std::env::var_os("HOME") + .or_else(|| std::env::var_os("USERPROFILE")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(".")); + match source { + ImportSource::Codex => codex::roots(&home), + ImportSource::Claude => claude::roots(&home), + } +} + +fn discover_transcripts(source: ImportSource) -> anyhow::Result> { + let roots = transcript_roots(source); + discover_transcripts_in(source, &roots) +} + +fn discover_transcripts_in( + source: ImportSource, + roots: &[PathBuf], +) -> anyhow::Result> { + let mut candidates = Vec::new(); + for root in roots { + find_jsonl_files(root, &mut candidates); + } + candidates.sort(); + candidates.dedup(); + + let mut sessions = BTreeMap::>::new(); + for path in candidates { + let Some(session_id) = transcript_session_id(&path, source) else { + continue; + }; + sessions.entry(session_id).or_default().push(path); + } + if sessions.is_empty() { + let locations = roots + .iter() + .map(|root| root.display().to_string()) + .collect::>() + .join(", "); + bail!( + "no {} transcripts found; searched {locations}", + source_name(source) + ); + } + + let mut resolved = Vec::with_capacity(sessions.len()); + for (session_id, paths) in sessions { + match paths.as_slice() { + [path] => resolved.push(path.clone()), + paths => { + let locations = paths + .iter() + .map(|path| path.display().to_string()) + .collect::>() + .join(", "); + bail!( + "multiple {} transcripts found for session {session_id}: {locations}", + source_name(source) + ); + } + } + } + Ok(resolved) +} + +fn find_jsonl_files(directory: &Path, matches: &mut Vec) { + let Ok(entries) = std::fs::read_dir(directory) else { + return; + }; + for entry in entries.flatten() { + let Ok(file_type) = entry.file_type() else { + continue; + }; + if file_type.is_symlink() { + continue; + } + let path = entry.path(); + if file_type.is_dir() { + find_jsonl_files(&path, matches); + } else if path.extension().and_then(|extension| extension.to_str()) == Some("jsonl") { + matches.push(path); + } + } +} + +fn transcript_session_id(path: &Path, source: ImportSource) -> Option { + match source { + ImportSource::Codex => codex::transcript_session_id(path), + ImportSource::Claude => claude::transcript_session_id(path), + } +} + +fn resolve_transcript_in( + session_id: &str, + source: ImportSource, + roots: &[PathBuf], +) -> anyhow::Result { + validate_session_id(session_id)?; + let mut matches = Vec::new(); + for root in roots { + let mut candidates = Vec::new(); + find_jsonl_files(root, &mut candidates); + matches.extend(candidates.into_iter().filter(|path| match source { + ImportSource::Codex => codex::filename_matches(path, session_id), + ImportSource::Claude => claude::filename_matches(path, session_id), + })); + } + matches.sort(); + matches.dedup(); + match matches.as_slice() { + [path] => Ok(path.clone()), + [] => { + let locations = roots + .iter() + .map(|root| root.display().to_string()) + .collect::>() + .join(", "); + bail!( + "no {} transcript found for session {session_id}; searched {locations}", + source_name(source) + ) + } + paths => { + let locations = paths + .iter() + .map(|path| path.display().to_string()) + .collect::>() + .join(", "); + bail!( + "multiple {} transcripts found for session {session_id}: {locations}", + source_name(source) + ) + } + } +} + +fn validate_session_id(session_id: &str) -> anyhow::Result<()> { + if session_id.is_empty() + || !session_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + { + bail!("invalid session id {session_id:?}"); + } + Ok(()) +} + +fn source_name(source: ImportSource) -> &'static str { + match source { + ImportSource::Codex => "Codex", + ImportSource::Claude => "Claude Code", + } +} + +pub(crate) fn transcript_envelopes( + path: &Path, + source: ImportSource, +) -> anyhow::Result> { + let contents = std::fs::read_to_string(path) + .with_context(|| format!("read transcript {}", path.display()))?; + let mut records = Vec::new(); + let mut record_end_offsets = Vec::new(); + let mut offset = 0_u64; + for (index, line) in contents.split_inclusive('\n').enumerate() { + offset += line.len() as u64; + if line.trim().is_empty() { + continue; + } + records.push( + serde_json::from_str(line).with_context(|| { + format!("parse transcript {} line {}", path.display(), index + 1) + })?, + ); + record_end_offsets.push(offset); + } + if records.is_empty() { + bail!("transcript {} is empty", path.display()); + } + match source { + ImportSource::Codex => codex::envelopes(path, &records), + ImportSource::Claude => { + claude::envelopes(path, &records, &record_end_offsets, contents.len() as u64) + } + } +} + +/// Incrementally converts a growing native transcript into synthetic hook +/// events for one persistent translator. The final poll closes the active +/// turn/session; ordinary polls keep the newest turn open. +pub(crate) struct TranscriptTail { + path: PathBuf, + source: ImportSource, + state: TailState, +} + +enum TailState { + Codex(codex::Tail), + Claude(claude::Tail), +} + +impl TranscriptTail { + pub(crate) fn new(path: PathBuf, source: ImportSource) -> Self { + Self { + path, + source, + state: match source { + ImportSource::Codex => TailState::Codex(codex::Tail::default()), + ImportSource::Claude => TailState::Claude(claude::Tail::default()), + }, + } + } + + pub(crate) fn poll(&mut self, finalize: bool) -> anyhow::Result> { + let events = match transcript_envelopes(&self.path, self.source) { + Ok(events) => events, + Err(_) if !finalize => return Ok(Vec::new()), + Err(error) => return Err(error), + }; + let len = std::fs::metadata(&self.path)?.len(); + match &mut self.state { + TailState::Codex(state) => state.poll(events, len, finalize), + TailState::Claude(state) => state.poll(events, len, finalize), + } + } +} + +fn read_jsonl_records(path: &Path) -> anyhow::Result> { + let contents = std::fs::read_to_string(path) + .with_context(|| format!("read transcript {}", path.display()))?; + contents + .lines() + .enumerate() + .filter(|(_, line)| !line.trim().is_empty()) + .map(|(index, line)| { + serde_json::from_str(line) + .with_context(|| format!("parse transcript {} line {}", path.display(), index + 1)) + }) + .collect() +} + +fn envelope( + source: &str, + source_version: Option, + session_id: &str, + event: &str, + ts_ms: i64, + payload: Value, +) -> Envelope { + Envelope { + source: source.into(), + source_version, + plugin_version: None, + session_id: session_id.into(), + event: event.into(), + ts_ms, + managed_run_id: None, + payload, + route: None, + config: None, + } +} + +fn timestamp_bounds(records: &[Value]) -> (i64, i64) { + let mut timestamps = records.iter().filter_map(timestamp_ms); + let Some(first) = timestamps.next() else { + return (0, 0); + }; + timestamps.fold((first, first), |(min, max), value| { + (min.min(value), max.max(value)) + }) +} + +fn timestamp_ms(record: &Value) -> Option { + chrono::DateTime::parse_from_rfc3339(record.get("timestamp")?.as_str()?) + .ok() + .map(|timestamp| timestamp.timestamp_millis()) +} + +fn string_at(record: &Value, pointer: &str) -> Option { + record.pointer(pointer)?.as_str().map(str::to_owned) +} + +fn file_session_id(path: &Path) -> String { + path.file_stem() + .and_then(|stem| stem.to_str()) + .filter(|stem| !stem.is_empty()) + .unwrap_or("imported-session") + .to_owned() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn finds_codex_rollout_by_session_suffix() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("sessions"); + let transcript = root + .join("2026/07/31") + .join("rollout-2026-07-31T12-00-00-session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + std::fs::write(&transcript, "{}\n").unwrap(); + + assert_eq!( + resolve_transcript_in("session-123", ImportSource::Codex, &[root]).unwrap(), + transcript + ); + } + + #[test] + fn finds_only_exact_claude_session_filename() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("projects"); + let project = root.join("-tmp-project"); + std::fs::create_dir_all(&project).unwrap(); + std::fs::write(project.join("prefix-session-123.jsonl"), "{}\n").unwrap(); + let transcript = project.join("session-123.jsonl"); + std::fs::write(&transcript, "{}\n").unwrap(); + + assert_eq!( + resolve_transcript_in("session-123", ImportSource::Claude, &[root]).unwrap(), + transcript + ); + } + + #[test] + fn rejects_unsafe_session_ids() { + let error = resolve_transcript_in( + "../session", + ImportSource::Codex, + &[PathBuf::from("unused")], + ) + .unwrap_err(); + assert!(error.to_string().contains("invalid session id")); + } + + #[test] + fn reports_missing_and_ambiguous_sessions() { + let temp = tempfile::tempdir().unwrap(); + let first = temp.path().join("first"); + let second = temp.path().join("second"); + std::fs::create_dir_all(&first).unwrap(); + std::fs::create_dir_all(&second).unwrap(); + + let missing = resolve_transcript_in( + "missing", + ImportSource::Claude, + &[first.clone(), second.clone()], + ) + .unwrap_err(); + assert!(missing.to_string().contains("no Claude Code transcript")); + + std::fs::write(first.join("duplicate.jsonl"), "{}\n").unwrap(); + std::fs::write(second.join("duplicate.jsonl"), "{}\n").unwrap(); + let ambiguous = + resolve_transcript_in("duplicate", ImportSource::Claude, &[first, second]).unwrap_err(); + assert!(ambiguous + .to_string() + .contains("multiple Claude Code transcripts")); + } + + #[test] + fn discovers_only_top_level_codex_transcripts_in_stable_order() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("sessions"); + let first = root.join("2026/01/01/rollout-first-session-a.jsonl"); + let second = root.join("2026/01/02/rollout-second-session-b.jsonl"); + std::fs::create_dir_all(first.parent().unwrap()).unwrap(); + std::fs::create_dir_all(second.parent().unwrap()).unwrap(); + std::fs::write( + &first, + "{\"type\":\"event_msg\"}\n{\"type\":\"session_meta\",\"payload\":{\"id\":\"session-a\"}}\n", + ) + .unwrap(); + std::fs::write( + &second, + "{\"type\":\"session_meta\",\"payload\":{\"id\":\"session-b\"}}\n", + ) + .unwrap(); + let subagent = root.join("2026/01/02/rollout-agent-c.jsonl"); + std::fs::write( + subagent, + "{\"type\":\"session_meta\",\"payload\":{\"id\":\"agent-c\",\"source\":{\"subagent\":{\"thread_spawn\":{\"parent_thread_id\":\"session-b\"}}}}}\n", + ) + .unwrap(); + std::fs::write(root.join("not-a-transcript.jsonl"), "{}\n").unwrap(); + + assert_eq!( + discover_transcripts_in(ImportSource::Codex, &[root]).unwrap(), + vec![first, second] + ); + } + + #[test] + fn claude_all_selects_parent_while_subagents_are_related_content() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("projects"); + let project = root.join("-tmp-project"); + let transcript = project.join("session-a.jsonl"); + let subagent = project.join("session-a/subagents/agent-1.jsonl"); + std::fs::create_dir_all(subagent.parent().unwrap()).unwrap(); + std::fs::write( + &transcript, + "{\"type\":\"user\",\"sessionId\":\"session-a\"}\n", + ) + .unwrap(); + std::fs::write( + &subagent, + "{\"type\":\"assistant\",\"sessionId\":\"session-a\"}\n", + ) + .unwrap(); + std::fs::write(project.join("notes.jsonl"), "{}\n").unwrap(); + + assert_eq!( + discover_transcripts_in(ImportSource::Claude, &[root]).unwrap(), + vec![transcript] + ); + } + + #[test] + fn claude_import_emits_related_subagent_lifecycle_inside_parent_turn() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp.path().join("session-a.jsonl"); + let subagent = temp.path().join("session-a/subagents/agent-child-a.jsonl"); + std::fs::create_dir_all(subagent.parent().unwrap()).unwrap(); + let records = [ + json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:01Z","message":{"content":"delegate"}}), + json!({"type":"assistant","sessionId":"session-a","timestamp":"2026-01-01T00:00:02Z","message":{"content":[{"type":"tool_use","id":"call-a","name":"Agent","input":{"subagent_type":"reviewer"}}]}}), + json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:05Z","toolUseResult":{"agentId":"child-a"},"message":{"content":[{"type":"tool_result","tool_use_id":"call-a","content":"done"}]}}), + json!({"type":"assistant","sessionId":"session-a","timestamp":"2026-01-01T00:00:06Z","message":{"content":[{"type":"text","text":"complete"}]}}), + ]; + std::fs::write( + &transcript, + records + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"), + ) + .unwrap(); + std::fs::write( + &subagent, + json!({"type":"assistant","sessionId":"session-a","timestamp":"2026-01-01T00:00:04Z","message":{"id":"sub-message","content":[{"type":"text","text":"reviewed"}]}}).to_string(), + ) + .unwrap(); + + let events = transcript_envelopes(&transcript, ImportSource::Claude).unwrap(); + let names = events + .iter() + .map(|event| event.event.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "SessionStart", + "UserPromptSubmit", + "SubagentStart", + "SubagentStop", + "Stop", + "SessionEnd" + ] + ); + assert_eq!(events[2].payload["agent_id"], json!("child-a")); + assert_eq!(events[2].payload["agent_type"], json!("reviewer")); + assert_eq!( + Path::new(events[3].payload["agent_transcript_path"].as_str().unwrap()), + subagent + ); + } + + #[test] + fn codex_import_discovers_spawned_rollout_and_emits_lifecycle() { + let temp = tempfile::tempdir().unwrap(); + let parent = temp.path().join("rollout-parent.jsonl"); + let child = temp.path().join("rollout-child-a.jsonl"); + let records = vec![ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"parent"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"response_item","payload":{"type":"function_call","call_id":"call-a","name":"spawn_agent","arguments":"{\"agent_type\":\"reviewer\"}"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call-a","output":"{\"agent_id\":\"child-a\"}"}}), + ]; + std::fs::write( + &parent, + records + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"), + ) + .unwrap(); + std::fs::write( + &child, + [ + json!({"timestamp":"2026-01-01T00:00:02Z","type":"session_meta","payload":{"id":"child-a","source":{"subagent":{"thread_spawn":{"parent_thread_id":"parent"}}}}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"reviewed"}}), + ] + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"), + ) + .unwrap(); + + let events = codex::envelopes(&parent, &records).unwrap(); + assert!(events.iter().any(|event| { + event.event == "SubagentStart" && event.payload["agent_id"] == json!("child-a") + })); + assert!(events.iter().any(|event| { + event.event == "SubagentStop" + && event.payload["agent_transcript_path"] == json!(child.to_string_lossy()) + })); + } + + #[test] + fn codex_import_adds_native_turn_checkpoints() { + let records = vec![ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"session-123"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1"}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-2"}}), + json!({"timestamp":"2026-01-01T00:00:05Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-2"}}), + ]; + let events = codex::envelopes(Path::new("rollout.jsonl"), &records).unwrap(); + + assert_eq!(events.first().unwrap().event, "SessionStart"); + assert_eq!(events.last().unwrap().event, "Stop"); + assert_eq!( + events + .iter() + .filter(|event| event.event == "ImportCheckpoint") + .count(), + 3 + ); + } + + #[test] + fn codex_tail_keeps_session_open_until_final_poll() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("session-123.jsonl"); + let mut records = vec![ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"session-123"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"response_item","payload":{"type":"message","role":"assistant"}}), + ]; + let write = |records: &[Value]| { + std::fs::write( + &path, + records + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"), + ) + .unwrap(); + }; + write(&records); + let mut tail = TranscriptTail::new(path.clone(), ImportSource::Codex); + let first = tail.poll(false).unwrap(); + assert_eq!(first.first().unwrap().event, "SessionStart"); + assert_eq!(first.last().unwrap().event, "ImportCheckpoint"); + assert!(first.iter().all(|event| event.event != "Stop")); + + records.push(json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1"}})); + write(&records); + let second = tail.poll(false).unwrap(); + assert!(second.iter().all(|event| event.event != "SessionStart")); + assert_eq!(second.last().unwrap().event, "ImportCheckpoint"); + assert_eq!(tail.poll(true).unwrap().last().unwrap().event, "Stop"); + } + + #[test] + fn claude_tail_closes_only_completed_turns() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("session-123.jsonl"); + let mut records = vec![ + json!({"type":"user","sessionId":"session-123","timestamp":"2026-01-01T00:00:01Z","message":{"content":"one"}}), + json!({"type":"assistant","sessionId":"session-123","timestamp":"2026-01-01T00:00:02Z","message":{"content":"answer one"}}), + ]; + let write = |records: &[Value]| { + std::fs::write( + &path, + records + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"), + ) + .unwrap(); + }; + write(&records); + let mut tail = TranscriptTail::new(path.clone(), ImportSource::Claude); + let first = tail.poll(false).unwrap(); + assert_eq!( + first + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["SessionStart", "UserPromptSubmit", "ImportCheckpoint"] + ); + + records.extend([ + json!({"type":"user","sessionId":"session-123","timestamp":"2026-01-01T00:00:03Z","message":{"content":"two"}}), + json!({"type":"assistant","sessionId":"session-123","timestamp":"2026-01-01T00:00:04Z","message":{"content":"answer two"}}), + ]); + write(&records); + let second = tail.poll(false).unwrap(); + assert_eq!( + second + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["Stop", "UserPromptSubmit", "ImportCheckpoint"] + ); + assert_eq!( + tail.poll(true) + .unwrap() + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["Stop", "SessionEnd"] + ); + } +} diff --git a/bt-daemon/src/translate/codex.rs b/bt-daemon/src/translate/codex.rs index 187bb2b..31c97e3 100644 --- a/bt-daemon/src/translate/codex.rs +++ b/bt-daemon/src/translate/codex.rs @@ -184,7 +184,7 @@ impl AgentTranslator for CodexTranslator { match event.event.as_str() { // Catch up first: this same hook may be the first observation of // the spawn_agent transcript record that establishes call -> turn. - "PostToolUse" if agent_id.is_none() => self.record_spawned_agent(payload), + "PostToolUse" => self.record_spawned_agent(payload), "SubagentStop" => { if let Some(p) = str_field(payload, "agent_transcript_path") { self.close_subagent(&p, event.ts_ms, &mut ops); @@ -695,8 +695,9 @@ impl CodexTranslator { let span_id = ids::span_id(&self.session_id, &format!("tool:{call_id}")); // spawn_agent: remember which turn ran it, so a later SubagentStart can - // nest the subagent root under this turn (main scope only). - if scope.kind == ScopeKind::Main && tool_name == SPAWN_AGENT_TOOL { + // nest the subagent root under this turn in either the main or a nested + // agent scope. + if tool_name == SPAWN_AGENT_TOOL { self.spawn_turn_by_call_id .insert(call_id.clone(), turn_span.clone()); } diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 5ac4c02..1c199c2 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -1,4 +1,6 @@ -use bt_daemon::{import_transcript, DebugSinkFactory, ImportSource, Registry, ServeOptions}; +use bt_daemon::{ + import_transcript, import_transcripts, DebugSinkFactory, ImportSource, Registry, ServeOptions, +}; use serde_json::{json, Value}; use std::io::Write; use std::sync::Arc; @@ -33,6 +35,164 @@ fn inserted(rows: &[Value], span_type: &str) -> usize { .count() } +#[tokio::test] +async fn imports_multiple_transcripts_in_one_invocation() { + let tmp = tempfile::tempdir().unwrap(); + let first = tmp.path().join("first.jsonl"); + let second = tmp.path().join("second.jsonl"); + for (path, session_id) in [(&first, "claude-first"), (&second, "claude-second")] { + write_jsonl( + path, + &[ + json!({"type":"user","timestamp":"2026-01-01T00:00:01Z","sessionId":session_id,"message":{"role":"user","content":"hello"}}), + json!({"type":"assistant","timestamp":"2026-01-01T00:00:02Z","sessionId":session_id,"message":{"id":format!("message-{session_id}"),"model":"claude-test","role":"assistant","content":[{"type":"text","text":"hi"}],"usage":{"input_tokens":1,"output_tokens":1}}}), + ], + ); + } + + let output = tmp.path().join("spans"); + import_transcripts( + &[first, second], + ImportSource::Claude, + options(&output), + None, + ) + .await + .unwrap(); + + for session_id in ["claude-first", "claude-second"] { + let rows = rows(&output.join(format!("{session_id}.ndjson"))); + assert_eq!(inserted(&rows, "task"), 2, "session and one turn"); + assert_eq!(inserted(&rows, "llm"), 1); + } +} + +#[tokio::test] +async fn importing_claude_parent_also_imports_its_subagent_into_the_same_trace() { + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("claude-parent.jsonl"); + let subagent = tmp + .path() + .join("claude-parent/subagents/agent-child-a.jsonl"); + std::fs::create_dir_all(subagent.parent().unwrap()).unwrap(); + write_jsonl( + &transcript, + &[ + json!({"type":"user","timestamp":"2026-01-01T00:00:01Z","sessionId":"claude-parent","message":{"content":"delegate"}}), + json!({"type":"assistant","timestamp":"2026-01-01T00:00:02Z","sessionId":"claude-parent","message":{"id":"main-call","model":"claude-test","content":[{"type":"tool_use","id":"spawn-call","name":"Agent","input":{"subagent_type":"reviewer"}}],"usage":{"input_tokens":3,"output_tokens":1}}}), + json!({"type":"user","timestamp":"2026-01-01T00:00:05Z","sessionId":"claude-parent","toolUseResult":{"agentId":"child-a"},"message":{"content":[{"type":"tool_result","tool_use_id":"spawn-call","content":"reviewed"}]}}), + json!({"type":"assistant","timestamp":"2026-01-01T00:00:06Z","sessionId":"claude-parent","message":{"id":"main-final","model":"claude-test","content":[{"type":"text","text":"done"}],"usage":{"input_tokens":4,"output_tokens":1}}}), + ], + ); + write_jsonl( + &subagent, + &[ + json!({"type":"assistant","timestamp":"2026-01-01T00:00:04Z","sessionId":"claude-parent","message":{"id":"child-message","model":"claude-test","content":[{"type":"text","text":"reviewed"}],"usage":{"input_tokens":2,"output_tokens":1}}}), + ], + ); + + let output = tmp.path().join("spans"); + import_transcript( + &transcript, + ImportSource::Claude, + options(&output), + None, + false, + ) + .await + .unwrap(); + + let rows = rows(&output.join("claude-parent.ndjson")); + assert_eq!(inserted(&rows, "task"), 3, "session, turn, and subagent"); + assert_eq!(inserted(&rows, "llm"), 3, "parent and child model calls"); + let inserts = rows + .iter() + .filter_map(|op| op.get("Insert")) + .collect::>(); + let root = inserts + .iter() + .find(|row| { + row.pointer("/metadata/session_id").and_then(Value::as_str) == Some("claude-parent") + && row.pointer("/metadata/source").and_then(Value::as_str) == Some("claude-code") + }) + .unwrap(); + let turn = inserts + .iter() + .find(|row| row.get("name").and_then(Value::as_str) == Some("Turn 1")) + .unwrap(); + let child = inserts + .iter() + .find(|row| row.get("name").and_then(Value::as_str) == Some("subagent: reviewer")) + .unwrap(); + assert_eq!(child["root_span_id"], root["span_id"]); + assert_eq!(child["parent_span_ids"][0], turn["span_id"]); +} + +#[tokio::test] +async fn importing_codex_parent_also_imports_its_subagent_into_the_same_trace() { + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("rollout-parent.jsonl"); + let subagent = tmp.path().join("rollout-child-a.jsonl"); + write_jsonl( + &transcript, + &[ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"codex-parent","cwd":"/tmp/demo"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"turn_context","payload":{"model":"gpt-test"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_started","turn_id":"main-turn"}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"response_item","payload":{"type":"function_call","call_id":"spawn-call","name":"spawn_agent","arguments":"{\"agent_type\":\"reviewer\"}"}}), + json!({"timestamp":"2026-01-01T00:00:05Z","type":"response_item","payload":{"type":"function_call_output","call_id":"spawn-call","output":"{\"agent_id\":\"child-a\"}"}}), + json!({"timestamp":"2026-01-01T00:00:10Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"done"}}), + ], + ); + write_jsonl( + &subagent, + &[ + json!({"timestamp":"2026-01-01T00:00:05Z","type":"session_meta","payload":{"id":"child-a","cwd":"/tmp/demo","source":{"subagent":{"thread_spawn":{"parent_thread_id":"codex-parent"}}}}}), + json!({"timestamp":"2026-01-01T00:00:06Z","type":"turn_context","payload":{"model":"gpt-test-mini"}}), + json!({"timestamp":"2026-01-01T00:00:07Z","type":"event_msg","payload":{"type":"task_started","turn_id":"child-turn"}}), + json!({"timestamp":"2026-01-01T00:00:08Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"reviewed"}]}}), + json!({"timestamp":"2026-01-01T00:00:09Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"reviewed"}}), + ], + ); + + let output = tmp.path().join("spans"); + import_transcript( + &transcript, + ImportSource::Codex, + options(&output), + None, + false, + ) + .await + .unwrap(); + + let rows = rows(&output.join("codex-parent.ndjson")); + assert_eq!( + inserted(&rows, "task"), + 4, + "session, two turns, and subagent" + ); + assert_eq!(inserted(&rows, "tool"), 1, "spawn_agent call"); + let inserts = rows + .iter() + .filter_map(|op| op.get("Insert")) + .collect::>(); + let root = inserts + .iter() + .find(|row| row.get("name").and_then(Value::as_str) == Some("codex: demo")) + .unwrap(); + let turn = inserts + .iter() + .find(|row| row.get("name").and_then(Value::as_str) == Some("turn: main-turn")) + .unwrap(); + let child = inserts + .iter() + .find(|row| row.get("name").and_then(Value::as_str) == Some("subagent: child-a")) + .unwrap(); + assert_eq!(child["root_span_id"], root["span_id"]); + assert_eq!(child["parent_span_ids"][0], turn["span_id"]); +} + #[tokio::test] async fn imports_native_codex_rollout_through_codex_translator() { let tmp = tempfile::tempdir().unwrap();