diff --git a/bt-daemon/README.md b/bt-daemon/README.md index aee0f40..0eee934 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -45,6 +45,8 @@ hook processes or JavaScript plugins. Each coding agent reads an independent non-credential `braintrust.json` file: - Codex: `~/.codex/braintrust.json` +- Muse Code: `$XDG_CONFIG_HOME/muse/braintrust.json`, falling back to + `~/.config/muse/braintrust.json` - Claude Code: `~/.claude/braintrust.json` - OpenCode: `$XDG_CONFIG_HOME/opencode/braintrust.json`, falling back to `~/.config/opencode/braintrust.json` @@ -135,13 +137,21 @@ echo '{"session_id":"s1","hook_event_name":"Stop"}' | ./target/debug/bt- The first `hook` spawns the daemon detached; it idles out after 5 minutes. -`import ` has a different purpose from restart +`import ` has a different purpose from restart recovery. It locates the native transcript in the selected agent's standard session store, synthesizes the lifecycle triggers that can be recovered from that transcript, and sends them through the normal translator and sink to create a trace for the past session. Hook-only facts absent from a native transcript are not invented. +Muse imports invoke its documented local `muse export --session` interface and +read export schema version 1. `import muse --all` enumerates durable session +IDs through Muse's read-only MSP `session/list` interface, then exports each +completed session. `--attach` is not available yet because it needs a durable +MSP subscription and cursor-following implementation. Muse also has no safe +invocation-local configuration overlay, so `run muse` is deliberately +unavailable; use persistent `enable muse` setup. + Add `--attach` to keep following an active Codex, Claude, or Antigravity transcript until Ctrl-C. `run [ARGS...]` launches the selected agent with inherited stdio and injects live Braintrust hooks for that invocation, so it diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 25d5516..608c495 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -56,8 +56,10 @@ use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; use clap::{Args, ValueEnum}; use std::ffi::OsString; use std::path::PathBuf; +use std::process::Stdio; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines}; use wire::{ method, Envelope, ManagedRunFlushParams, SessionConfig, SessionRoute, StatusResult, PROTOCOL_VERSION, @@ -65,6 +67,7 @@ use wire::{ const MANAGED_RUN_ID_ENV: &str = "BT_TRACE_MANAGED_RUN_ID"; const MANAGED_RUN_FLUSH_TIMEOUT_MS: u64 = 10_000; +const MUSE_MSP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); /// Arguments for `serve`. #[derive(Debug, Clone, Args)] @@ -228,6 +231,7 @@ pub struct ImportArgs { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum ImportSource { Codex, + Muse, #[value(name = "claude", alias = "claude-code")] Claude, #[value(name = "antigravity", alias = "agy")] @@ -602,6 +606,21 @@ pub async fn run_import( .map(|components| wire::TraceDestination::ParentSpan { components }) .or_else(|| args.destination.clone()); apply_import_destination(&mut config, destination)?; + if args.source == ImportSource::Muse { + if args.attach { + anyhow::bail!( + "Muse import does not support --attach yet; use a completed export-v1 session" + ) + } + let exports = prepare_muse_exports(&args).await?; + return import_muse_exports_with_ledger( + &exports, + opts, + config, + Some(paths::data_dir(None)), + ) + .await; + } let files = transcript_import::resolve_transcripts(&args.session_ids, args.all, args.source)?; let ledger_dir = paths::data_dir(None); if args.attach { @@ -618,6 +637,213 @@ pub async fn run_import( import_transcripts_with_ledger(&files, args.source, opts, config, Some(ledger_dir)).await } +async fn prepare_muse_exports(args: &ImportArgs) -> anyhow::Result { + let session_ids = if args.all { + muse_msp_session_ids().await? + } else { + args.session_ids.clone() + }; + if session_ids.is_empty() { + anyhow::bail!("no durable Muse sessions found to import") + } + export_muse_sessions(&session_ids).await +} + +async fn export_muse_sessions(session_ids: &[String]) -> anyhow::Result { + let directory = tempfile::Builder::new() + .prefix("bt-muse-import-") + .tempdir()?; + let executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| OsString::from("muse")); + for (index, session_id) in session_ids.iter().enumerate() { + let output = tokio::process::Command::new(&executable) + .kill_on_drop(true) + .args([ + OsString::from("export"), + OsString::from("--session"), + OsString::from(session_id), + OsString::from("--out"), + ]) + .arg(directory.path().join(format!("{index}.json"))) + .output() + .await + .with_context(|| { + format!( + "run {} export for Muse session {session_id}", + executable.to_string_lossy() + ) + })?; + if !output.status.success() { + anyhow::bail!( + "Muse export for session {session_id} failed: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + } + } + Ok(directory) +} + +/// List durable Muse sessions through the read-only MSP surface. This never +/// resumes a session, obtains a writer lease, or subscribes to its view; the +/// export command remains the authoritative content reader. +async fn muse_msp_session_ids() -> anyhow::Result> { + const MAX_PAGES: usize = 1_000; + let executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| OsString::from("muse")); + let mut child = tokio::process::Command::new(&executable) + .kill_on_drop(true) + .arg("serve") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| { + format!( + "run {} serve to list Muse sessions", + executable.to_string_lossy() + ) + })?; + let mut stdin = child.stdin.take().expect("piped stdin"); + let stdout = child.stdout.take().expect("piped stdout"); + let mut reader = BufReader::new(stdout).lines(); + let mut request_id = 1u64; + muse_msp_request( + &mut stdin, + &mut reader, + request_id, + "initialize", + serde_json::json!({"clientInfo":{"name":"braintrust_muse_import","version":"1"}}), + ) + .await?; + // Muse follows the JSON-RPC/MSP initialization handshake: requests after + // `initialize` are rejected until this notification has been flushed. + muse_msp_notify(&mut stdin, "initialized").await?; + let mut cursor = None; + let mut sessions = Vec::new(); + for _ in 0..MAX_PAGES { + request_id += 1; + let result = muse_msp_request( + &mut stdin, + &mut reader, + request_id, + "session/list", + serde_json::json!({"cursor": cursor, "limit": 200}), + ) + .await?; + let page = result + .get("sessions") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| anyhow::anyhow!("Muse session/list response has no sessions array"))?; + sessions.extend(page.iter().filter_map(|session| { + session + .get("sessionId") + .or_else(|| session.get("session_id")) + .and_then(serde_json::Value::as_str) + .filter(|id| transcript_import::validate_session_id(id).is_ok()) + .map(str::to_owned) + })); + cursor = result + .get("nextCursor") + .and_then(serde_json::Value::as_str) + .map(str::to_owned); + if cursor.is_none() { + break; + } + } + drop(stdin); + let output = child.wait_with_output().await?; + if !output.status.success() { + anyhow::bail!( + "Muse MSP session listing failed: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + } + sessions.sort(); + sessions.dedup(); + Ok(sessions) +} + +async fn muse_msp_notify(stdin: &mut W, method: &str) -> anyhow::Result<()> +where + W: AsyncWrite + Unpin, +{ + let mut frame = serde_json::to_string(&serde_json::json!({"jsonrpc":"2.0", "method":method}))?; + frame.push('\n'); + stdin.write_all(frame.as_bytes()).await?; + stdin.flush().await?; + Ok(()) +} + +async fn muse_msp_request( + stdin: &mut W, + reader: &mut Lines, + id: u64, + method: &str, + params: serde_json::Value, +) -> anyhow::Result +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin, +{ + let mut request = serde_json::to_string( + &serde_json::json!({"jsonrpc":"2.0", "id":id, "method":method, "params":params}), + )?; + request.push('\n'); + stdin.write_all(request.as_bytes()).await?; + stdin.flush().await?; + tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { + loop { + let line = reader + .next_line() + .await? + .ok_or_else(|| anyhow::anyhow!("Muse MSP closed before responding to {method}"))?; + let frame: serde_json::Value = serde_json::from_str(&line) + .with_context(|| format!("parse Muse MSP frame while waiting for {method}"))?; + if frame.get("id").and_then(serde_json::Value::as_u64) != Some(id) { + continue; + } + if let Some(error) = frame.get("error") { + anyhow::bail!("Muse MSP {method} failed: {error}") + } + return frame + .get("result") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Muse MSP {method} response has no result")); + } + }) + .await + .map_err(|_| anyhow::anyhow!("Muse MSP timed out waiting for {method}"))? +} + +async fn import_muse_exports_with_ledger( + directory: &tempfile::TempDir, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + let mut processor = ImportProcessor::new(opts, config, ledger_dir); + let mut exports = std::fs::read_dir(directory.path()) + .with_context(|| format!("read Muse exports {}", directory.path().display()))? + .flatten() + .map(|entry| entry.path()) + .collect::>(); + exports.sort(); + let mut summaries = Vec::new(); + for export in exports { + let entries = transcript_import::muse::envelopes(&export)?; + let session_ids = entries + .iter() + .map(|entry| entry.session_id.clone()) + .collect::>(); + processor.process(entries).await?; + for session_id in session_ids { + if let Some(summary) = processor.finish_session(&session_id).await? { + summaries.push(summary); + } + } + } + summaries.extend(processor.finish().await?); + Ok(summaries) +} + 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"); @@ -1369,6 +1595,7 @@ mod tests { use super::*; use clap::Parser; use serde_json::json; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; #[derive(Debug, Parser)] struct ImportCli { @@ -1451,6 +1678,46 @@ mod tests { assert!(!should_flush_hook_event("turn_completed", true)); } + #[tokio::test] + async fn muse_msp_handshake_writes_notification_and_skips_notifications() { + let (client, server) = tokio::io::duplex(4_096); + let (client_reader, mut client_writer) = tokio::io::split(client); + let (server_reader, mut server_writer) = tokio::io::split(server); + let server = tokio::spawn(async move { + let mut incoming = BufReader::new(server_reader).lines(); + let initialized = incoming.next_line().await.unwrap().unwrap(); + assert_eq!( + serde_json::from_str::(&initialized).unwrap(), + json!({"jsonrpc": "2.0", "method": "initialized"}) + ); + let request = incoming.next_line().await.unwrap().unwrap(); + assert_eq!( + serde_json::from_str::(&request).unwrap()["id"], + 7 + ); + server_writer + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"session/changed\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n") + .await + .unwrap(); + }); + + let mut reader = BufReader::new(client_reader).lines(); + muse_msp_notify(&mut client_writer, "initialized") + .await + .unwrap(); + let result = muse_msp_request( + &mut client_writer, + &mut reader, + 7, + "session/list", + json!({"cursor": null, "limit": 200}), + ) + .await + .unwrap(); + assert_eq!(result, json!({"sessions": []})); + server.await.unwrap(); + } + #[test] fn additional_metadata_overrides_a_route_only_with_a_json_object() { let mut route = SessionRoute { @@ -1506,6 +1773,11 @@ mod tests { assert!(all.session_ids.is_empty()); assert!(all.all); + let muse = ImportCli::try_parse_from(["test", "muse", "session-a"]) + .unwrap() + .args; + assert_eq!(muse.source, ImportSource::Muse); + assert!(ImportCli::try_parse_from(["test", "codex"]).is_err()); assert!(ImportCli::try_parse_from(["test", "codex", "session-a", "--all"]).is_err()); } diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index d591f5c..fb97341 100644 --- a/bt-daemon/src/main.rs +++ b/bt-daemon/src/main.rs @@ -110,7 +110,7 @@ enum Command { }, /// Print daemon/session status. Status(StatusArgs), - /// Import a past Codex or Claude Code session by its resume id. + /// Import a past coding-agent session by its resume id. Import(ImportArgs), /// Launch a coding agent with live tracing hooks for this invocation. Run { diff --git a/bt-daemon/src/paths.rs b/bt-daemon/src/paths.rs index 1e3048b..d07fcbd 100644 --- a/bt-daemon/src/paths.rs +++ b/bt-daemon/src/paths.rs @@ -95,6 +95,7 @@ pub fn agent_settings_path(source: &str, explicit: Option<&Path>) -> PathBuf { } match source { "codex" => home().join(".codex").join("braintrust.json"), + "muse" => muse_config_dir().join("braintrust.json"), "claude" | "claude-code" => claude_config_dir().join("braintrust.json"), "opencode" => std::env::var_os("XDG_CONFIG_HOME") .filter(|path| !path.is_empty()) @@ -112,6 +113,15 @@ pub fn agent_settings_path(source: &str, explicit: Option<&Path>) -> PathBuf { } } +/// Muse Code's user configuration directory. +pub(crate) fn muse_config_dir() -> PathBuf { + std::env::var_os("XDG_CONFIG_HOME") + .filter(|path| !path.is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| home().join(".config")) + .join("muse") +} + /// Claude Code's shared settings file. Braintrust never mutates this file; /// setup only inspects it for obsolete tracing-specific environment entries. pub(crate) fn claude_settings_path() -> PathBuf { diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index c742283..1c77386 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -24,6 +24,21 @@ const OPENCODE_PACKAGE_MANIFEST: &str = include_str!("../../src/plugins/opencode/content/package.json"); const PI_PACKAGE_MANIFEST: &str = include_str!("../../src/plugins/pi/content/package.json"); const ANTIGRAVITY_PLUGIN: &str = "braintrust-antigravity-tracing"; +const MUSE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreLLMCall", + "PostLLMCall", + "PreCompact", + "PostCompact", + "SubagentStart", + "SubagentStop", + "Stop", + "SessionEnd", +]; const LEGACY_CLAUDE_TRACING_ENV_KEYS: [&str; 2] = ["BRAINTRUST_CC_PROJECT", "BRAINTRUST_CC_DEBUG"]; #[cfg(unix)] const ANTIGRAVITY_PLUGIN_SOURCE: &str = @@ -686,6 +701,85 @@ fn update_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { runner.run("pi", &["update", PI_PACKAGE]) } +/// Install Braintrust's hook file through Muse's supported managed-hook +/// pointer. We never merge into `hooks`: that map belongs to the user and +/// other tools. A pre-existing non-Braintrust managed path is an explicit +/// conflict rather than something setup may overwrite. +fn muse_hook_path(config_dir: &Path) -> PathBuf { + config_dir.join("braintrust-hooks.json") +} + +fn muse_hook_config() -> Map { + let command = "bt trace hook --source muse --flush-on-turn-end || true"; + let hooks = MUSE_HOOK_EVENTS + .iter() + .map(|event| { + ( + (*event).to_string(), + serde_json::json!([{ + "matcher": "*", + "hooks": [{"type": "command", "command": command}], + }]), + ) + }) + .collect::>(); + Map::from_iter([ + ("schema_version".into(), Value::from(1)), + ("hooks".into(), Value::Object(hooks)), + ]) +} + +fn setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { + let settings_path = config_dir.join("settings.json"); + let hook_path = muse_hook_path(config_dir); + let mut settings = load_object(&settings_path)?; + let owned_path = hook_path.to_string_lossy(); + if let Some(existing) = settings.get("managed_hooks_path").and_then(Value::as_str) { + if existing != owned_path { + bail!( + "Muse already uses managed_hooks_path {}; refusing to overwrite it", + existing + ); + } + } + settings.entry("schema_version").or_insert(Value::from(1)); + settings.insert( + "managed_hooks_path".into(), + Value::String(owned_path.into_owned()), + ); + write_object_atomic(&hook_path, muse_hook_config())?; + write_object_atomic(&settings_path, settings) +} + +fn disable_muse_at(config_dir: &Path) -> anyhow::Result<()> { + let settings_path = config_dir.join("settings.json"); + let hook_path = muse_hook_path(config_dir); + let mut settings = load_object(&settings_path)?; + if settings.get("managed_hooks_path").and_then(Value::as_str) + == Some(hook_path.to_string_lossy().as_ref()) + { + settings.remove("managed_hooks_path"); + write_object_atomic(&settings_path, settings)?; + } + match std::fs::remove_file(&hook_path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error).with_context(|| { + format!( + "failed to remove Muse hook configuration: {}", + hook_path.display() + ) + }), + } +} + +fn setup_muse() -> anyhow::Result<()> { + setup_muse_at(&paths::muse_config_dir()) +} +fn disable_muse() -> anyhow::Result<()> { + disable_muse_at(&paths::muse_config_dir()) +} + fn antigravity_home(config_dir: &Path) -> anyhow::Result<&Path> { if config_dir.file_name().and_then(|part| part.to_str()) != Some("config") { bail!( @@ -882,6 +976,7 @@ pub fn run_disable(agent: SetupAgent) -> anyhow::Result { SetupAgent::OpenCode => disable_opencode(), SetupAgent::Pi => disable_pi(&mut runner), SetupAgent::Grok => disable_grok(&mut runner), + SetupAgent::Muse => disable_muse(), SetupAgent::Antigravity => disable_antigravity(&mut runner), }; let settings_path = paths::agent_settings_path(source, None); @@ -904,6 +999,7 @@ pub fn run_update(agent: SetupAgent) -> anyhow::Result { SetupAgent::OpenCode => update_opencode()?, SetupAgent::Pi => update_pi(&mut runner)?, SetupAgent::Grok => update_grok(&mut runner)?, + SetupAgent::Muse => {} SetupAgent::Antigravity => update_antigravity(&mut runner)?, } Ok(TraceCommandOutput::update(source, display_name)) @@ -916,6 +1012,7 @@ fn agent_details(agent: SetupAgent) -> (&'static str, &'static str) { SetupAgent::OpenCode => ("opencode", "OpenCode"), SetupAgent::Pi => ("pi", "Pi"), SetupAgent::Grok => ("grok", "Grok"), + SetupAgent::Muse => ("muse", "Muse Code"), SetupAgent::Antigravity => ("antigravity", "Google Antigravity"), } } @@ -946,6 +1043,10 @@ pub fn run_enable(args: EnableArgs, route: SessionRoute) -> anyhow::Result { + setup_muse()?; + ("muse", "Muse Code") + } SetupAgent::Antigravity => { setup_antigravity(&mut runner)?; ("antigravity", "Google Antigravity") @@ -1629,6 +1730,57 @@ mod tests { ); } + #[test] + fn muse_setup_owns_only_its_managed_hook_file() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + std::fs::create_dir_all(&config_dir).unwrap(); + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"theme":"dark"}"#, + ) + .unwrap(); + + setup_muse_at(&config_dir).unwrap(); + setup_muse_at(&config_dir).unwrap(); + let settings: Value = + serde_json::from_slice(&std::fs::read(config_dir.join("settings.json")).unwrap()) + .unwrap(); + assert_eq!(settings["theme"], "dark"); + assert_eq!( + settings["managed_hooks_path"], + muse_hook_path(&config_dir).to_string_lossy().as_ref() + ); + let hooks: Value = + serde_json::from_slice(&std::fs::read(muse_hook_path(&config_dir)).unwrap()).unwrap(); + assert_eq!(hooks["schema_version"], 1); + assert!(hooks["hooks"]["PreLLMCall"].is_array()); + + disable_muse_at(&config_dir).unwrap(); + let settings: Value = + serde_json::from_slice(&std::fs::read(config_dir.join("settings.json")).unwrap()) + .unwrap(); + assert_eq!(settings["theme"], "dark"); + assert!(settings.get("managed_hooks_path").is_none()); + assert!(!muse_hook_path(&config_dir).exists()); + } + + #[test] + fn muse_setup_refuses_another_managed_hook_owner() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + std::fs::create_dir_all(&config_dir).unwrap(); + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"managed_hooks_path":"/other/hooks.json"}"#, + ) + .unwrap(); + assert!(setup_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + } + #[test] fn disabling_removes_only_the_braintrust_settings_file() { let temp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs index fd9bb19..801a089 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -63,6 +63,7 @@ pub struct DoctorArgs { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum DoctorAgent { Codex, + Muse, #[value(name = "claude", alias = "claude-code")] Claude, #[value(name = "opencode", alias = "open-code")] @@ -76,6 +77,7 @@ impl DoctorAgent { pub(crate) fn source(self) -> &'static str { match self { Self::Codex => "codex", + Self::Muse => "muse", Self::Claude => "claude", Self::OpenCode => "opencode", Self::Pi => "pi", @@ -87,6 +89,7 @@ impl DoctorAgent { pub(crate) fn display_name(self) -> &'static str { match self { Self::Codex => "Codex", + Self::Muse => "Muse Code", Self::Claude => "Claude Code", Self::OpenCode => "OpenCode", Self::Pi => "Pi", @@ -132,6 +135,8 @@ pub struct UpdateArgs { pub enum SetupAgent { /// Install the published Codex tracing plugin. Codex, + /// Configure Muse Code tracing hooks. + Muse, /// Install the published Claude Code tracing plugin. Claude, /// Configure the published OpenCode tracing plugin. @@ -270,6 +275,7 @@ mod tests { fn doctor_accepts_every_supported_agent_alias() { for (agent, expected, source, display_name) in [ ("codex", DoctorAgent::Codex, "codex", "Codex"), + ("muse", DoctorAgent::Muse, "muse", "Muse Code"), ("claude-code", DoctorAgent::Claude, "claude", "Claude Code"), ("open-code", DoctorAgent::OpenCode, "opencode", "OpenCode"), ("pi", DoctorAgent::Pi, "pi", "Pi"), @@ -315,6 +321,27 @@ mod tests { assert!(Cli::try_parse_from(["bt", "setup", "antigravity", "--disable"]).is_err()); } + #[test] + fn muse_uses_shared_enable_and_disable_commands() { + for command in ["enable", "setup"] { + let parsed = Cli::try_parse_from(["bt", command, "muse"]).unwrap(); + assert!(matches!( + parsed.trace.command, + TraceCommand::Setup(SetupArgs { + agent: SetupAgent::Muse, + .. + }) + )); + } + let parsed = Cli::try_parse_from(["bt", "disable", "muse"]).unwrap(); + assert!(matches!( + parsed.trace.command, + TraceCommand::Disable(DisableArgs { + agent: SetupAgent::Muse + }) + )); + } + #[test] fn grok_uses_shared_enable_disable_and_doctor_commands() { for command in ["enable", "setup"] { @@ -350,7 +377,15 @@ mod tests { #[test] fn update_accepts_every_setup_agent() { - for agent in ["codex", "claude", "opencode", "pi", "grok", "antigravity"] { + for agent in [ + "codex", + "claude", + "opencode", + "pi", + "grok", + "muse", + "antigravity", + ] { assert!(matches!( Cli::try_parse_from(["bt", "update", agent]) .unwrap() diff --git a/bt-daemon/src/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs index 5d045b6..fc7aba5 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -9,6 +9,7 @@ use std::path::{Path, PathBuf}; mod antigravity; mod claude; mod codex; +pub(crate) mod muse; pub(crate) fn resolve_transcripts( session_ids: &[String], @@ -41,6 +42,7 @@ fn transcript_roots(source: ImportSource) -> Vec { .unwrap_or_else(|| PathBuf::from(".")); match source { ImportSource::Codex => codex::roots(&home), + ImportSource::Muse => Vec::new(), ImportSource::Claude => claude::roots(&home), ImportSource::Antigravity => antigravity::roots(&home), } @@ -124,6 +126,7 @@ fn find_jsonl_files(directory: &Path, matches: &mut Vec) { fn transcript_session_id(path: &Path, source: ImportSource) -> Option { match source { ImportSource::Codex => codex::transcript_session_id(path), + ImportSource::Muse => None, ImportSource::Claude => claude::transcript_session_id(path), ImportSource::Antigravity => antigravity::transcript_session_id(path), } @@ -141,6 +144,7 @@ fn resolve_transcript_in( find_jsonl_files(root, &mut candidates); matches.extend(candidates.into_iter().filter(|path| match source { ImportSource::Codex => codex::filename_matches(path, session_id), + ImportSource::Muse => false, ImportSource::Claude => claude::filename_matches(path, session_id), ImportSource::Antigravity => antigravity::filename_matches(path, session_id), })); @@ -174,7 +178,7 @@ fn resolve_transcript_in( } } -fn validate_session_id(session_id: &str) -> anyhow::Result<()> { +pub(crate) fn validate_session_id(session_id: &str) -> anyhow::Result<()> { if session_id.is_empty() || !session_id .bytes() @@ -188,6 +192,7 @@ fn validate_session_id(session_id: &str) -> anyhow::Result<()> { fn source_name(source: ImportSource) -> &'static str { match source { ImportSource::Codex => "Codex", + ImportSource::Muse => "Muse Code", ImportSource::Claude => "Claude Code", ImportSource::Antigravity => "Google Antigravity", } @@ -210,6 +215,7 @@ fn envelopes_from_records( ) -> anyhow::Result> { match source { ImportSource::Codex => codex::envelopes(path, &records.values), + ImportSource::Muse => muse::envelopes(path), ImportSource::Claude => claude::envelopes( path, &records.values, @@ -353,6 +359,7 @@ pub(crate) struct TranscriptTail { enum TailState { Codex(codex::Tail), + Muse, Claude(claude::Tail), Antigravity(antigravity::Tail), } @@ -372,6 +379,7 @@ impl TranscriptTail { fn new_state(source: ImportSource) -> TailState { match source { ImportSource::Codex => TailState::Codex(codex::Tail::default()), + ImportSource::Muse => TailState::Muse, ImportSource::Claude => TailState::Claude(claude::Tail::default()), ImportSource::Antigravity => TailState::Antigravity(antigravity::Tail::default()), } @@ -404,6 +412,7 @@ impl TranscriptTail { .len(); match &mut self.state { TailState::Codex(state) => state.poll(events, len, finalize), + TailState::Muse => Ok(events), TailState::Claude(state) => state.poll(events, len, finalize), TailState::Antigravity(state) => state.poll(events, len, finalize), } diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs new file mode 100644 index 0000000..0d4dd18 --- /dev/null +++ b/bt-daemon/src/transcript_import/muse.rs @@ -0,0 +1,254 @@ +//! Versioned Muse Code export-v1 reader. +//! +//! Muse's durable local export is the authoritative import input. Hook +//! payloads are intentionally not replayed here because their LLM output is a +//! preview and their subagent attribution is incomplete. + +use super::{envelope, validate_session_id}; +use crate::wire::Envelope; +use anyhow::{bail, Context}; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::path::Path; + +pub(crate) fn envelopes(path: &Path) -> anyhow::Result> { + let export: Value = serde_json::from_slice( + &std::fs::read(path).with_context(|| format!("read Muse export {}", path.display()))?, + ) + .with_context(|| format!("parse Muse export {}", path.display()))?; + if export.get("export_schema_version").and_then(Value::as_i64) != Some(1) { + bail!( + "unsupported Muse export schema in {}; expected version 1", + path.display() + ); + } + let session = export + .get("sessions") + .and_then(Value::as_array) + .and_then(|sessions| sessions.first()) + .ok_or_else(|| anyhow::anyhow!("Muse export {} has no sessions", path.display()))?; + let session_id = session + .get("session_id") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("Muse export {} has no session_id", path.display()))?; + validate_session_id(session_id)?; + let version = export + .get("exporter_version") + .and_then(|version| { + version + .as_str() + .or_else(|| version.get("semver").and_then(Value::as_str)) + }) + .map(str::to_owned); + let events = export + .get("events") + .and_then(Value::as_array) + .ok_or_else(|| anyhow::anyhow!("Muse export {} has no events", path.display()))?; + let transcript_path = path.to_string_lossy().into_owned(); + let first_ts = events.iter().filter_map(recorded_ms).next().unwrap_or(0); + let mut out = vec![event( + version.clone(), + session_id, + "SessionStart", + first_ts, + json!({ + "session_id": session_id, + "hook_event_name": "SessionStart", + "source": "muse_export_v1", + "trajectory_id": session.get("trajectory_id"), + "root_session_id": session.get("root_session_id"), + "transcript_path": transcript_path, + }), + )]; + let mut requests = HashMap::::new(); + let mut assistant = HashMap::::new(); + let mut completed = HashMap::::new(); + let mut last_ts = first_ts; + for item in events { + let Some(envelope) = item.get("envelope") else { + continue; + }; + let Some(payload) = envelope.get("payload") else { + continue; + }; + let run_id = payload.get("run_id").and_then(Value::as_str); + let native = payload.get("event").and_then(Value::as_object); + let Some(native) = native else { continue }; + let Some(kind) = native.get("kind").and_then(Value::as_str) else { + continue; + }; + let ts = recorded_ms(item).unwrap_or(last_ts); + last_ts = last_ts.max(ts); + match kind { + "started" if payload.get("kind").and_then(Value::as_str) == Some("run") => { + let Some(run_id) = run_id else { continue }; + out.push(event( + version.clone(), + session_id, + "UserPromptSubmit", + ts, + json!({ + "session_id": session_id, + "hook_event_name": "UserPromptSubmit", + "turn_id": run_id, + "prompt": native.get("prompt").cloned().unwrap_or(Value::Null), + "source": "muse_export_v1", + }), + )); + } + "model_input_trace_recorded" => { + let Some(run_id) = run_id else { continue }; + let Some(request_id) = native.get("request_record_id").and_then(Value::as_str) + else { + continue; + }; + requests.insert(run_id.to_owned(), request_id.to_owned()); + out.push(event( + version.clone(), + session_id, + "PreLLMCall", + ts, + json!({ + "session_id": session_id, + "hook_event_name": "PreLLMCall", + "turn_id": run_id, + "request_id": request_id, + "messages": native.get("bounded").cloned().unwrap_or(Value::Null), + "source": "muse_export_v1", + }), + )); + } + "model_completed" => { + let Some(run_id) = run_id else { continue }; + let Some(request_id) = requests.get(run_id) else { + continue; + }; + completed.insert( + run_id.to_owned(), + ( + request_id.to_owned(), + native.get("usage").cloned().unwrap_or(Value::Null), + ), + ); + } + "assistant_message_committed" => { + if let Some(run_id) = run_id { + let text = native.get("text").cloned().unwrap_or(Value::Null); + assistant.insert(run_id.to_owned(), text.clone()); + if let Some((request_id, usage)) = completed.remove(run_id) { + out.push(event( + version.clone(), + session_id, + "PostLLMCall", + ts, + json!({ + "session_id": session_id, + "hook_event_name": "PostLLMCall", + "turn_id": run_id, + "request_id": request_id, + "usage": usage, + "output_text_preview": text, + "source": "muse_export_v1", + }), + )); + } + } + } + "terminal" => { + let Some(run_id) = run_id else { continue }; + if let Some((request_id, usage)) = completed.remove(run_id) { + out.push(event(version.clone(), session_id, "PostLLMCall", ts, json!({ + "session_id": session_id, + "hook_event_name": "PostLLMCall", + "turn_id": run_id, + "request_id": request_id, + "usage": usage, + "output_text_preview": assistant.get(run_id).cloned().unwrap_or(Value::Null), + "source": "muse_export_v1", + }))); + } + out.push(event( + version.clone(), + session_id, + "Stop", + ts, + json!({ + "session_id": session_id, + "hook_event_name": "Stop", + "turn_id": run_id, + "last_assistant_message": assistant.remove(run_id).unwrap_or(Value::Null), + "source": "muse_export_v1", + }), + )); + } + _ => {} + } + } + out.push(event( + version, + session_id, + "SessionEnd", + last_ts, + json!({ + "session_id": session_id, + "hook_event_name": "SessionEnd", + "reason": session.pointer("/session_end/exit_reason"), + "source": "muse_export_v1", + }), + )); + Ok(out) +} + +fn recorded_ms(item: &Value) -> Option { + item.get("recorded_at")? + .as_i64() + .map(|micros| micros / 1_000) +} + +fn event( + version: Option, + session_id: &str, + name: &str, + ts: i64, + payload: Value, +) -> Envelope { + envelope("muse", version, session_id, name, ts, payload) +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn export_v1_becomes_shared_muse_lifecycle_events() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("export.json"); + std::fs::write(&path, serde_json::to_vec(&json!({ + "export_schema_version": 1, "exporter_version": "1.1.1", + "sessions": [{"session_id":"session-1", "trajectory_id":"t", "root_session_id":"session-1", "session_end":{"exit_reason":"clean"}}], + "events": [ + {"recorded_at": 1_000_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}, + {"recorded_at": 1_001_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","request_record_id":"request-1","bounded":{"message":"hello"}}}}}, + {"recorded_at": 1_002_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_completed","usage":{"input_tokens":1}}}}}, + {"recorded_at": 1_003_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"world"}}}}, + {"recorded_at": 1_004_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal"}}}} + ] + })).unwrap()).unwrap(); + let events = envelopes(&path).unwrap(); + assert_eq!( + events + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + [ + "SessionStart", + "UserPromptSubmit", + "PreLLMCall", + "PostLLMCall", + "Stop", + "SessionEnd" + ] + ); + assert_eq!(events[2].ts_ms, 1001); + assert_eq!(events[4].payload["last_assistant_message"], "world"); + } +} diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index d30f28c..18d91a4 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -13,6 +13,7 @@ mod codex; mod debug; mod git; mod grok; +mod muse; mod opencode; mod pi; mod recent; @@ -23,6 +24,7 @@ pub use claude::ClaudeTranslatorFactory; pub use codex::CodexTranslatorFactory; pub use debug::DebugTranslatorFactory; pub use grok::GrokTranslatorFactory; +pub use muse::MuseTranslatorFactory; pub use opencode::OpenCodeTranslatorFactory; pub use pi::PiTranslatorFactory; @@ -171,6 +173,7 @@ impl Registry { r.register(Box::new(ClaudeTranslatorFactory::new(git.clone()))); r.register(Box::new(CodexTranslatorFactory::new(git.clone()))); r.register(Box::new(GrokTranslatorFactory::new(git.clone()))); + r.register(Box::new(MuseTranslatorFactory::new(git.clone()))); r.register(Box::new(OpenCodeTranslatorFactory::new(git.clone()))); r.register(Box::new(PiTranslatorFactory::new(git))); r @@ -195,6 +198,7 @@ impl Registry { "antigravity" => "antigravity", "codex" => "codex", "grok" => "grok", + "muse" | "muse-code" => "muse", "pi" => "pi", "debug" => "debug", _ => return None, diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs new file mode 100644 index 0000000..4e69f1e --- /dev/null +++ b/bt-daemon/src/translate/muse.rs @@ -0,0 +1,401 @@ +//! Muse Code hook translator. +//! +//! Muse hooks deliberately remain a thin, fail-open transport. This state +//! machine only relies on fields observed in Muse 1.1's documented hook +//! payloads. An MSP/export reader can feed the same event names later without +//! changing trace construction. + +use super::git::GitMetadataCache; +use super::{ + local_username, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, +}; +use crate::ids; +use crate::wire::Envelope; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::Arc; + +pub struct MuseTranslatorFactory { + git: Arc, +} +impl MuseTranslatorFactory { + pub(super) fn new(git: Arc) -> Self { + Self { git } + } +} +impl TranslatorFactory for MuseTranslatorFactory { + fn source(&self) -> &str { + "muse" + } + fn create(&self, session_id: &str) -> Box { + Box::new(MuseTranslator::new(session_id, self.git.clone())) + } +} + +#[derive(Clone)] +struct OpenOp { + id: String, + parents: Vec, + name: String, + start_ms: i64, + input: Value, + kind: SpanType, +} + +struct MuseTranslator { + session_id: String, + root_id: String, + effective_root_id: String, + external_parent: Option, + opened: bool, + turns: HashMap, + llms: HashMap, + tools: HashMap, + special: HashMap, + last_ts: i64, + git: Arc, +} + +impl MuseTranslator { + fn new(session_id: &str, git: Arc) -> Self { + Self { + session_id: session_id.into(), + root_id: ids::span_id(session_id, "root"), + effective_root_id: String::new(), + external_parent: None, + opened: false, + turns: HashMap::new(), + llms: HashMap::new(), + tools: HashMap::new(), + special: HashMap::new(), + last_ts: 0, + git, + } + } + fn ensure_root(&mut self, e: &Envelope, ctx: &SessionCtx) -> Vec { + if self.opened { + return Vec::new(); + } + self.opened = true; + let attached = ctx + .config + .as_ref() + .map(|c| c.attached_span_ids()) + .unwrap_or_default(); + self.external_parent = attached.0; + self.effective_root_id = attached + .1 + .or_else(|| self.external_parent.clone()) + .unwrap_or_else(|| self.root_id.clone()); + let mut metadata = ctx + .config + .as_ref() + .and_then(|c| c.additional_metadata.as_ref()) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + metadata.insert("source".into(), json!("muse")); + metadata.insert("session_id".into(), json!(ctx.session_id)); + metadata.insert("username".into(), json!(local_username())); + metadata.insert("muse_version".into(), json!(e.source_version)); + for key in [ + "cwd", + "model", + "permission_mode", + "trajectory_id", + "root_session_id", + ] { + if let Some(value) = e.payload.get(key) { + metadata.insert(key.into(), value.clone()); + } + } + vec![SpanOp::Insert(SpanRow { + span_id: self.root_id.clone(), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: self.external_parent.clone().into_iter().collect(), + name: "Muse Code".into(), + span_type: SpanType::Task, + start_ms: Some(e.ts_ms), + metadata: Some(Value::Object(metadata)), + ..Default::default() + })] + } + fn close( + root_span_id: &str, + map: &mut HashMap, + key: &str, + ts: i64, + output: Option, + metrics: Option, + error: Option, + ) -> Vec { + let Some(open) = map.remove(key) else { + return Vec::new(); + }; + vec![SpanOp::Insert(SpanRow { + span_id: open.id, + root_span_id: root_span_id.into(), + parent_span_ids: open.parents, + name: open.name, + span_type: open.kind, + start_ms: Some(open.start_ms), + end_ms: Some(ts), + input: Some(open.input), + output, + metrics, + error, + ..Default::default() + })] + } + fn close_all(&mut self, ts: i64, error: Option<&str>) -> Vec { + let mut out = Vec::new(); + for map in [ + &mut self.llms, + &mut self.tools, + &mut self.special, + &mut self.turns, + ] { + for (_, open) in std::mem::take(map) { + out.push(SpanOp::Insert(SpanRow { + span_id: open.id, + root_span_id: self.effective_root_id.clone(), + parent_span_ids: open.parents, + name: open.name, + span_type: open.kind, + start_ms: Some(open.start_ms), + end_ms: Some(ts), + input: Some(open.input), + error: error.map(str::to_owned), + ..Default::default() + })); + } + } + out + } +} + +impl AgentTranslator for MuseTranslator { + fn handle(&mut self, e: &Envelope, ctx: &SessionCtx) -> anyhow::Result> { + self.last_ts = self.last_ts.max(e.ts_ms); + let mut out = self.ensure_root(e, ctx); + let p = &e.payload; + let turn_id = p.get("turn_id").and_then(Value::as_str); + match e.event.as_str() { + "UserPromptSubmit" => { + let key = turn_id.unwrap_or("unknown"); + if !self.turns.contains_key(key) { + self.turns.insert( + key.into(), + OpenOp { + id: ids::span_id(&self.session_id, &format!("turn:{key}")), + parents: vec![self.root_id.clone()], + name: "Turn".into(), + start_ms: e.ts_ms, + input: p.get("prompt").cloned().unwrap_or(Value::Null), + kind: SpanType::Task, + }, + ); + } + } + "PreLLMCall" => { + let key = p + .get("request_id") + .and_then(Value::as_str) + .unwrap_or("unknown"); + self.llms.entry(key.into()).or_insert_with(|| OpenOp { id: ids::span_id(&self.session_id, &format!("llm:{key}")), parents: vec![turn_id.and_then(|id| self.turns.get(id)).map(|op| op.id.clone()).unwrap_or_else(|| self.root_id.clone())], name: "llm".into(), start_ms: e.ts_ms, input: json!({"messages": p.get("messages"), "tools": p.get("tools"), "provider": p.get("provider")}), kind: SpanType::Llm }); + } + "PostLLMCall" => { + let key = p + .get("request_id") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let metrics = p.get("usage").cloned(); + let output = p.get("output_text_preview").cloned(); + let error = p.get("error").and_then(Value::as_str).map(str::to_owned); + out.extend(Self::close( + &self.effective_root_id, + &mut self.llms, + key, + e.ts_ms, + output, + metrics, + error, + )); + } + "PreToolUse" => { + let key = p + .get("tool_use_id") + .or_else(|| p.get("tool_call_id")) + .and_then(Value::as_str) + .unwrap_or("unknown"); + self.tools.entry(key.into()).or_insert_with(|| OpenOp { + id: ids::span_id(&self.session_id, &format!("tool:{key}")), + parents: vec![turn_id + .and_then(|id| self.turns.get(id)) + .map(|op| op.id.clone()) + .unwrap_or_else(|| self.root_id.clone())], + name: p + .get("tool_name") + .and_then(Value::as_str) + .unwrap_or("tool") + .into(), + start_ms: e.ts_ms, + input: p + .get("tool_input") + .or_else(|| p.get("input")) + .cloned() + .unwrap_or(Value::Null), + kind: SpanType::Tool, + }); + } + "PostToolUse" => { + let key = p + .get("tool_use_id") + .or_else(|| p.get("tool_call_id")) + .and_then(Value::as_str) + .unwrap_or("unknown"); + out.extend(Self::close( + &self.effective_root_id, + &mut self.tools, + key, + e.ts_ms, + p.get("tool_response").or_else(|| p.get("output")).cloned(), + None, + p.get("error").and_then(Value::as_str).map(str::to_owned), + )); + } + "PermissionRequest" | "PreCompact" | "SubagentStart" => { + let key = format!( + "{}:{}", + e.event, + p.get("request_id") + .or_else(|| p.get("subagent_id")) + .and_then(Value::as_str) + .unwrap_or("unknown") + ); + self.special.entry(key.clone()).or_insert_with(|| OpenOp { + id: ids::span_id(&self.session_id, &format!("special:{key}")), + parents: vec![turn_id + .and_then(|id| self.turns.get(id)) + .map(|op| op.id.clone()) + .unwrap_or_else(|| self.root_id.clone())], + name: e.event.clone(), + start_ms: e.ts_ms, + input: p.clone(), + kind: SpanType::Task, + }); + } + "PostCompact" | "SubagentStop" => { + let start = if e.event == "PostCompact" { + "PreCompact" + } else { + "SubagentStart" + }; + let key = format!( + "{}:{}", + start, + p.get("request_id") + .or_else(|| p.get("subagent_id")) + .and_then(Value::as_str) + .unwrap_or("unknown") + ); + out.extend(Self::close( + &self.effective_root_id, + &mut self.special, + &key, + e.ts_ms, + Some(p.clone()), + None, + None, + )); + } + "Stop" => { + if let Some(id) = turn_id { + out.extend(Self::close( + &self.effective_root_id, + &mut self.turns, + id, + e.ts_ms, + p.get("last_assistant_message").cloned(), + None, + None, + )); + } + } + "SessionEnd" => { + out.extend(self.close_all(e.ts_ms, None)); + if self.opened { + out.push(SpanOp::Insert(SpanRow { + span_id: self.root_id.clone(), + root_span_id: self.effective_root_id.clone(), + end_ms: Some(e.ts_ms), + ..Default::default() + })); + self.opened = false; + } + } + _ => {} + } + self.git + .enrich_rows(p.get("cwd").and_then(Value::as_str), &mut out); + Ok(out) + } + fn finalize(&mut self, _ctx: &SessionCtx) -> anyhow::Result> { + let mut out = self.close_all(self.last_ts, Some("Interrupted before completion")); + if self.opened { + out.push(SpanOp::Insert(SpanRow { + span_id: self.root_id.clone(), + root_span_id: self.effective_root_id.clone(), + end_ms: Some(self.last_ts), + error: Some("Interrupted before completion".into()), + ..Default::default() + })); + self.opened = false; + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::translate::Registry; + use crate::wire::Envelope; + fn event(name: &str, payload: Value) -> Envelope { + Envelope { + source: "muse".into(), + source_version: Some("1.1.1".into()), + plugin_version: None, + session_id: "session-1".into(), + event: name.into(), + ts_ms: 1000, + managed_run_id: None, + capture: None, + payload, + route: None, + config: None, + } + } + #[test] + fn registry_and_llm_pairing_are_stable() { + let registry = Registry::default_agents(); + assert_eq!(registry.canonical_source("muse-code"), Some("muse")); + let mut translator = registry.create("muse", "session-1"); + let ctx = SessionCtx { + session_id: "session-1".into(), + config: None, + }; + translator + .handle( + &event( + "UserPromptSubmit", + json!({"session_id":"session-1","turn_id":"turn-1","prompt":"hello"}), + ), + &ctx, + ) + .unwrap(); + translator.handle(&event("PreLLMCall", json!({"session_id":"session-1","turn_id":"turn-1","request_id":"req-1","messages":[]})), &ctx).unwrap(); + let rows = translator.handle(&event("PostLLMCall", json!({"session_id":"session-1","turn_id":"turn-1","request_id":"req-1","output_text_preview":"hi","usage":{"input_tokens":1}})), &ctx).unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.name == "llm" && row.output == Some(json!("hi"))))); + } +} diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 86156e5..c1ff6fd 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -14,6 +14,43 @@ fn write_jsonl(path: &std::path::Path, records: &[Value]) { } } +#[tokio::test] +async fn importing_muse_export_uses_the_production_translator() { + let tmp = tempfile::tempdir().unwrap(); + let export = tmp.path().join("muse-export.json"); + std::fs::write(&export, serde_json::to_vec(&json!({ + "export_schema_version": 1, + "exporter_version": "1.1.1", + "sessions": [{ + "session_id": "muse-import", + "trajectory_id": "trajectory-muse-import", + "root_session_id": "muse-import", + "session_end": {"exit_reason": "clean"} + }], + "events": [ + {"recorded_at": 1_000_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"started", "prompt":"trace this"}}}}, + {"recorded_at": 1_001_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"model_input_trace_recorded", "request_record_id":"request-1", "bounded":{"request_digest":"sha256:test"}}}}}, + {"recorded_at": 1_002_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"model_completed", "usage":{"input_tokens":3, "output_tokens":2}}}}}, + {"recorded_at": 1_003_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"assistant_message_committed", "text":"traced"}}}}, + {"recorded_at": 1_004_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"terminal"}}}} + ] + })).unwrap()).unwrap(); + let output = tmp.path().join("spans"); + import_transcript(&export, ImportSource::Muse, options(&output), None, false) + .await + .unwrap(); + let output_rows = rows(&output.join("muse-import.ndjson")); + assert_eq!(inserted(&output_rows, "llm"), 1); + assert!(output_rows + .iter() + .any(|row| { row.pointer("/Insert/output").and_then(Value::as_str) == Some("traced") })); + assert!(output_rows.iter().any(|row| { + row.pointer("/Insert/metrics/input_tokens") + .and_then(Value::as_i64) + == Some(3) + })); +} + #[tokio::test] async fn importing_large_codex_rollout_drains_translator_continuations() { const CALLS: usize = 32;