From 0c57308751cfa6a6acf9d5a5e4b1b51c885a5b0b Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 11 Sep 2026 03:52:05 +0800 Subject: [PATCH 01/17] Add Muse Code hook tracing --- bt-daemon/src/paths.rs | 10 + bt-daemon/src/setup.rs | 1856 +++++++++++++++++++++++++++++++ bt-daemon/src/trace_command.rs | 120 +- bt-daemon/src/translate/mod.rs | 141 +-- bt-daemon/src/translate/muse.rs | 401 +++++++ 5 files changed, 2366 insertions(+), 162 deletions(-) create mode 100644 bt-daemon/src/setup.rs create mode 100644 bt-daemon/src/translate/muse.rs diff --git a/bt-daemon/src/paths.rs b/bt-daemon/src/paths.rs index c7799c85..adf3962c 100644 --- a/bt-daemon/src/paths.rs +++ b/bt-daemon/src/paths.rs @@ -103,6 +103,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()) @@ -121,6 +122,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 new file mode 100644 index 00000000..6293e837 --- /dev/null +++ b/bt-daemon/src/setup.rs @@ -0,0 +1,1856 @@ +//! Persistent installation and configuration for coding-agent tracing plugins. + +use crate::paths; +use crate::trace_command::{EnableArgs, SetupAgent}; +use crate::wire::SessionRoute; +use crate::TraceCommandOutput; +use anyhow::{bail, Context}; +use serde_json::{Map, Value}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::process::Command as ProcessCommand; + +const CODEX_MARKETPLACE: &str = "braintrust-codex-plugins"; +const CODEX_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-codex-plugin"; +const CODEX_PLUGIN: &str = "trace-codex@braintrust-codex-plugins"; +const CLAUDE_MARKETPLACE: &str = "braintrust-claude-plugin"; +const CLAUDE_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-claude-plugin"; +const CLAUDE_PLUGIN: &str = "trace-claude-code@braintrust-claude-plugin"; +const GROK_PLUGIN: &str = "trace-grok"; +const GROK_PLUGIN_SOURCE: &str = "braintrustdata/braintrust-grok-plugin"; +const OPENCODE_PACKAGE: &str = "@braintrust/trace-opencode"; +const PI_PACKAGE: &str = "@braintrust/pi-extension"; +const OPENCODE_PACKAGE_MANIFEST: &str = + include_str!("../../src/plugins/opencode/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"]; +const ANTIGRAVITY_PLUGIN_SOURCE: &str = + "https://github.com/braintrustdata/braintrust-antigravity-plugin"; + +fn package_version(manifest: &str) -> anyhow::Result { + let manifest = serde_json::from_str::(manifest)?; + manifest + .get("version") + .and_then(Value::as_str) + .map(str::to_owned) + .ok_or_else(|| anyhow::anyhow!("package manifest has no string version")) +} + +fn npm_major_spec(package: &str, manifest: &str) -> anyhow::Result { + let version = package_version(manifest)?; + let major = version + .split_once('.') + .map(|(major, _)| major) + .ok_or_else(|| anyhow::anyhow!("package version is not semver: {version}"))? + .parse::() + .with_context(|| format!("package version is not semver: {version}"))?; + Ok(format!("{package}@^{major}")) +} + +fn opencode_plugin_spec() -> anyhow::Result { + npm_major_spec(OPENCODE_PACKAGE, OPENCODE_PACKAGE_MANIFEST) +} + +pub(crate) fn pi_plugin_spec() -> String { + format!("npm:{PI_PACKAGE}") +} + +fn version_is_older(installed: &str, expected: &str) -> bool { + let parse = |version: &str| { + let mut parts = version.split(['.', '-', '+']); + Some(( + parts.next()?.parse::().ok()?, + parts.next().unwrap_or("0").parse::().ok()?, + parts.next().unwrap_or("0").parse::().ok()?, + )) + }; + parse(installed) + .zip(parse(expected)) + .is_some_and(|(installed, expected)| installed < expected) +} + +pub(crate) fn update_warning(source: &str) -> Option { + let stale = match source { + "codex" => installed_json_version("codex", &["plugin", "list", "--json"], |value| { + codex_plugin(value).and_then(|plugin| plugin.get("version")).and_then(Value::as_str) + }, &package_version(include_str!("../../src/plugins/codex/content/plugins/trace-codex/.codex-plugin/plugin.json")).ok()?), + "claude" => installed_json_version("claude", &["plugin", "list", "--json"], |value| { + claude_plugin(value).and_then(|plugin| plugin.get("version")).and_then(Value::as_str) + }, &package_version(include_str!("../../src/plugins/claude/content/plugins/trace-claude-code/.claude-plugin/plugin.json")).ok()?), + "opencode" => opencode_update_required(), + "pi" => pi_update_required(), + _ => false, + }; + stale.then(|| format!("tracing plugin is out of date; run `bt trace update {source}`")) +} + +fn installed_json_version( + program: &str, + args: &[&str], + find_version: impl Fn(&Value) -> Option<&str>, + expected: &str, +) -> bool { + let output = ProcessCommand::new(program).args(args).output(); + let Ok(output) = output else { return false }; + if !output.status.success() { + return false; + } + serde_json::from_slice::(&output.stdout) + .ok() + .and_then(|value| find_version(&value).map(str::to_owned)) + .is_some_and(|installed| version_is_older(&installed, expected)) +} + +fn opencode_update_required() -> bool { + let settings_path = paths::agent_settings_path("opencode", None); + let path = settings_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("opencode.json"); + let Ok(raw) = std::fs::read(path) else { + return false; + }; + let Ok(config) = serde_json::from_slice::(&raw) else { + return false; + }; + let expected = opencode_plugin_spec().ok(); + config + .get("plugin") + .and_then(Value::as_array) + .and_then(|plugins| { + plugins.iter().find_map(|plugin| { + plugin + .as_str() + .filter(|plugin| plugin.starts_with(OPENCODE_PACKAGE)) + }) + }) + .is_some_and(|plugin| Some(plugin) != expected.as_deref()) +} + +fn pi_update_required() -> bool { + let output = ProcessCommand::new("pi").arg("list").output(); + let Ok(output) = output else { return false }; + if !output.status.success() { + return false; + } + let installed = String::from_utf8_lossy(&output.stdout); + let expected = pi_plugin_spec(); + installed.lines().any(|line| { + let plugin = line.trim(); + plugin.starts_with(&expected) && plugin != expected + }) +} + +trait CommandRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result; + fn json_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result; + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()>; + fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()>; +} + +struct SystemCommandRunner; + +impl CommandRunner for SystemCommandRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { + let output = ProcessCommand::new(program) + .args(args) + .output() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); + } + serde_json::from_slice(&output.stdout) + .with_context(|| format!("`{program} {}` returned invalid JSON", args.join(" "))) + } + + fn json_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result { + let output = ProcessCommand::new(program) + .args(args) + .env("HOME", home) + .output() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); + } + serde_json::from_slice(&output.stdout) + .with_context(|| format!("`{program} {}` returned invalid JSON", args.join(" "))) + } + + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { + let status = ProcessCommand::new(program) + .args(args) + .status() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !status.success() { + bail!("`{program} {}` failed with {status}", args.join(" ")); + } + Ok(()) + } + + fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()> { + let output = ProcessCommand::new(program) + .args(args) + .env("HOME", home) + .output() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); + } + Ok(()) + } +} + +fn github_repo_matches(source: &str, expected: &str) -> bool { + let source = source.trim().trim_end_matches('/'); + let source = source.strip_suffix(".git").unwrap_or(source); + let source = source + .strip_prefix("https://github.com/") + .or_else(|| source.strip_prefix("git@github.com:")) + .unwrap_or(source); + source == expected +} + +fn codex_marketplace(value: &Value) -> Option<&Value> { + value + .get("marketplaces") + .and_then(Value::as_array)? + .iter() + .find(|item| item.get("name").and_then(Value::as_str) == Some(CODEX_MARKETPLACE)) +} + +fn codex_marketplace_is_published(item: &Value) -> bool { + item.get("marketplaceSource") + .and_then(|source| source.get("source")) + .and_then(Value::as_str) + .is_some_and(|source| github_repo_matches(source, CODEX_MARKETPLACE_SOURCE)) +} + +fn setup_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let marketplaces = runner.json("codex", &["plugin", "marketplace", "list", "--json"])?; + match codex_marketplace(&marketplaces) { + Some(marketplace) if codex_marketplace_is_published(marketplace) => runner.run( + "codex", + &["plugin", "marketplace", "upgrade", CODEX_MARKETPLACE], + )?, + Some(_) => { + runner.run( + "codex", + &["plugin", "marketplace", "remove", CODEX_MARKETPLACE], + )?; + runner.run( + "codex", + &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], + )?; + } + None => runner.run( + "codex", + &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], + )?, + } + + // Adding is idempotent and reconciles the installed cache to the refreshed + // marketplace snapshot. + runner.run("codex", &["plugin", "add", CODEX_PLUGIN]) +} + +fn codex_plugin(value: &Value) -> Option<&Value> { + value + .get("installed") + .and_then(Value::as_array)? + .iter() + .find(|item| item.get("pluginId").and_then(Value::as_str) == Some(CODEX_PLUGIN)) +} + +fn disable_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("codex", &["plugin", "list", "--json"])?; + if codex_plugin(&plugins).is_some() { + runner.run("codex", &["plugin", "remove", CODEX_PLUGIN, "--json"])?; + } + Ok(()) +} + +fn update_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("codex", &["plugin", "list", "--json"])?; + if codex_plugin(&plugins).is_none() { + bail!("Codex tracing plugin is not installed; run `bt trace enable codex`"); + } + let marketplaces = runner.json("codex", &["plugin", "marketplace", "list", "--json"])?; + let marketplace = codex_marketplace(&marketplaces).ok_or_else(|| { + anyhow::anyhow!("Codex tracing marketplace is not installed; run `bt trace enable codex`") + })?; + if !codex_marketplace_is_published(marketplace) { + bail!("Codex tracing marketplace is not the published Braintrust marketplace; run `bt trace enable codex`"); + } + runner.run( + "codex", + &["plugin", "marketplace", "upgrade", CODEX_MARKETPLACE], + )?; + runner.run("codex", &["plugin", "add", CODEX_PLUGIN]) +} + +fn claude_marketplace(value: &Value) -> Option<&Value> { + value + .as_array()? + .iter() + .find(|item| item.get("name").and_then(Value::as_str) == Some(CLAUDE_MARKETPLACE)) +} + +fn claude_marketplace_is_published(item: &Value) -> bool { + item.get("source").and_then(Value::as_str) == Some("github") + && item + .get("repo") + .and_then(Value::as_str) + .is_some_and(|repo| github_repo_matches(repo, CLAUDE_MARKETPLACE_SOURCE)) +} + +fn claude_plugin(value: &Value) -> Option<&Value> { + value + .as_array()? + .iter() + .find(|item| item.get("id").and_then(Value::as_str) == Some(CLAUDE_PLUGIN)) +} + +fn setup_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let marketplaces = runner.json("claude", &["plugin", "marketplace", "list", "--json"])?; + let marketplace_replaced = match claude_marketplace(&marketplaces) { + Some(marketplace) if claude_marketplace_is_published(marketplace) => { + runner.run( + "claude", + &["plugin", "marketplace", "update", CLAUDE_MARKETPLACE], + )?; + false + } + Some(_) => { + runner.run( + "claude", + &["plugin", "marketplace", "remove", CLAUDE_MARKETPLACE], + )?; + runner.run( + "claude", + &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], + )?; + true + } + None => { + runner.run( + "claude", + &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], + )?; + false + } + }; + + // Claude removes a marketplace's installed plugins when that marketplace + // is removed, so replacing a stale source requires a fresh installation. + if marketplace_replaced { + return runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]); + } + + let plugins = runner.json("claude", &["plugin", "list", "--json"])?; + match claude_plugin(&plugins) { + None => runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]), + Some(plugin) => { + runner.run("claude", &["plugin", "update", CLAUDE_PLUGIN])?; + if plugin.get("enabled").and_then(Value::as_bool) == Some(false) { + runner.run("claude", &["plugin", "enable", CLAUDE_PLUGIN])?; + } + Ok(()) + } + } +} + +fn legacy_claude_tracing_env_keys(path: &Path) -> anyhow::Result> { + let raw = match std::fs::read(path) { + Ok(raw) => raw, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => { + return Err(error) + .with_context(|| format!("failed to inspect Claude settings: {}", path.display())) + } + }; + let value: Value = serde_json::from_slice(&raw) + .with_context(|| format!("Claude settings are not valid JSON: {}", path.display()))?; + let env = value.get("env").and_then(Value::as_object); + Ok(LEGACY_CLAUDE_TRACING_ENV_KEYS + .into_iter() + .filter(|key| env.is_some_and(|env| env.contains_key(*key))) + .collect()) +} + +fn warn_legacy_claude_tracing_env() { + let path = paths::claude_settings_path(); + let keys = match legacy_claude_tracing_env_keys(&path) { + Ok(keys) => keys, + Err(error) => { + eprintln!("warning: could not inspect legacy Claude tracing settings: {error}"); + return; + } + }; + if !keys.is_empty() { + eprintln!( + "warning: obsolete Braintrust tracing settings in {}: {}; remove these keys from env. BRAINTRUST_API_KEY is left unchanged because the Braintrust MCP plugin may use it.", + path.display(), + keys.join(", ") + ); + } +} + +fn disable_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("claude", &["plugin", "list", "--json"])?; + if claude_plugin(&plugins).is_some() { + runner.run("claude", &["plugin", "uninstall", CLAUDE_PLUGIN])?; + } + Ok(()) +} + +fn grok_plugin(value: &Value) -> Option<&Value> { + value + .as_array()? + .iter() + .find(|item| item.get("name").and_then(Value::as_str) == Some(GROK_PLUGIN)) +} + +fn grok_plugin_is_published(item: &Value) -> bool { + item.get("source") + .and_then(Value::as_str) + .is_some_and(|source| github_repo_matches(source, GROK_PLUGIN_SOURCE)) +} + +fn setup_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("grok", &["plugin", "list", "--json"])?; + // `bt trace enable grok` is the user's trust boundary. Grok's `--trust` + // applies to this plugin installation and does not change folder trust. + match grok_plugin(&plugins) { + Some(plugin) if grok_plugin_is_published(plugin) => { + runner.run("grok", &["plugin", "update", GROK_PLUGIN])?; + } + Some(_) => { + runner.run("grok", &["plugin", "uninstall", GROK_PLUGIN, "--confirm"])?; + runner.run( + "grok", + &["plugin", "install", GROK_PLUGIN_SOURCE, "--trust"], + )?; + } + None => runner.run( + "grok", + &["plugin", "install", GROK_PLUGIN_SOURCE, "--trust"], + )?, + } + runner.run("grok", &["plugin", "enable", GROK_PLUGIN]) +} + +fn disable_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("grok", &["plugin", "list", "--json"])?; + if grok_plugin(&plugins).is_some_and(grok_plugin_is_published) { + runner.run("grok", &["plugin", "uninstall", GROK_PLUGIN, "--confirm"])?; + } + Ok(()) +} + +fn update_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("grok", &["plugin", "list", "--json"])?; + let plugin = grok_plugin(&plugins).ok_or_else(|| { + anyhow::anyhow!("Grok tracing plugin is not installed; run `bt trace enable grok`") + })?; + if !grok_plugin_is_published(plugin) { + bail!("Grok tracing plugin is not the published Braintrust plugin; run `bt trace enable grok`"); + } + runner.run("grok", &["plugin", "update", GROK_PLUGIN]) +} + +fn update_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugins = runner.json("claude", &["plugin", "list", "--json"])?; + if claude_plugin(&plugins).is_none() { + bail!("Claude Code tracing plugin is not installed; run `bt trace enable claude`"); + } + let marketplaces = runner.json("claude", &["plugin", "marketplace", "list", "--json"])?; + let marketplace = claude_marketplace(&marketplaces).ok_or_else(|| { + anyhow::anyhow!( + "Claude Code tracing marketplace is not installed; run `bt trace enable claude`" + ) + })?; + if !claude_marketplace_is_published(marketplace) { + bail!("Claude Code tracing marketplace is not the published Braintrust marketplace; run `bt trace enable claude`"); + } + runner.run( + "claude", + &["plugin", "marketplace", "update", CLAUDE_MARKETPLACE], + )?; + runner.run("claude", &["plugin", "update", CLAUDE_PLUGIN]) +} + +fn load_object(path: &Path) -> anyhow::Result> { + match std::fs::read(path) { + Ok(raw) => { + let value: Value = serde_json::from_slice(&raw) + .with_context(|| format!("invalid JSON configuration: {}", path.display()))?; + value.as_object().cloned().ok_or_else(|| { + anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) + }) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Map::new()), + Err(error) => { + Err(error).with_context(|| format!("failed to read configuration: {}", path.display())) + } + } +} + +fn write_object_atomic(path: &Path, object: Map) -> anyhow::Result<()> { + crate::settings::with_settings_lock(path, || write_object_atomic_unlocked(path, object)) +} + +fn write_object_atomic_unlocked(path: &Path, object: Map) -> anyhow::Result<()> { + let parent = path + .parent() + .ok_or_else(|| anyhow::anyhow!("configuration path has no parent: {}", path.display()))?; + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create configuration directory: {}", + parent.display() + ) + })?; + let mut encoded = serde_json::to_string_pretty(&Value::Object(object))?; + encoded.push('\n'); + let mut temporary = tempfile::NamedTempFile::new_in(parent) + .with_context(|| format!("failed to create temporary file in {}", parent.display()))?; + temporary.write_all(encoded.as_bytes()).with_context(|| { + format!( + "failed to write temporary configuration for {}", + path.display() + ) + })?; + temporary + .persist(path) + .map_err(|error| error.error) + .with_context(|| format!("failed to replace configuration: {}", path.display()))?; + Ok(()) +} + +fn setup_opencode_at(path: &Path) -> anyhow::Result<()> { + let mut config = load_object(path)?; + let plugins = config + .entry("plugin") + .or_insert_with(|| Value::Array(Vec::new())) + .as_array_mut() + .ok_or_else(|| { + anyhow::anyhow!( + "OpenCode `plugin` config must be an array: {}", + path.display() + ) + })?; + plugins.retain(|plugin| { + plugin.as_str().is_none_or(|plugin| { + plugin != "@braintrust/trace-opencode" + && !plugin.starts_with("@braintrust/trace-opencode@") + }) + }); + plugins.push(Value::String(opencode_plugin_spec()?)); + write_object_atomic(path, config) +} + +fn setup_opencode() -> anyhow::Result<()> { + let settings_path = paths::agent_settings_path("opencode", None); + let path = settings_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("opencode.json"); + setup_opencode_at(&path) +} + +fn update_opencode_at(path: &Path) -> anyhow::Result<()> { + let mut config = match std::fs::read(path) { + Ok(raw) => serde_json::from_slice::(&raw) + .with_context(|| format!("invalid JSON configuration: {}", path.display()))? + .as_object() + .cloned() + .ok_or_else(|| { + anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) + })?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + bail!("OpenCode tracing plugin is not installed; run `bt trace enable opencode`") + } + Err(error) => { + return Err(error) + .with_context(|| format!("failed to read configuration: {}", path.display())) + } + }; + let plugins = config + .get_mut("plugin") + .and_then(Value::as_array_mut) + .ok_or_else(|| { + anyhow::anyhow!( + "OpenCode tracing plugin is not installed; run `bt trace enable opencode`" + ) + })?; + let found = plugins.iter().any(|plugin| { + plugin.as_str().is_some_and(|plugin| { + plugin == OPENCODE_PACKAGE || plugin.starts_with(&format!("{OPENCODE_PACKAGE}@")) + }) + }); + if !found { + bail!("OpenCode tracing plugin is not installed; run `bt trace enable opencode`"); + } + plugins.retain(|plugin| { + plugin.as_str().is_none_or(|plugin| { + plugin != OPENCODE_PACKAGE && !plugin.starts_with(&format!("{OPENCODE_PACKAGE}@")) + }) + }); + plugins.push(Value::String(opencode_plugin_spec()?)); + write_object_atomic(path, config) +} + +fn update_opencode() -> anyhow::Result<()> { + let settings_path = paths::agent_settings_path("opencode", None); + let path = settings_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("opencode.json"); + update_opencode_at(&path) +} + +fn remove_opencode_plugin_at(path: &Path) -> anyhow::Result<()> { + let mut config = match std::fs::read(path) { + Ok(raw) => serde_json::from_slice::(&raw) + .with_context(|| format!("invalid JSON configuration: {}", path.display()))? + .as_object() + .cloned() + .ok_or_else(|| { + anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) + })?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => { + return Err(error) + .with_context(|| format!("failed to read configuration: {}", path.display())) + } + }; + let Some(plugins) = config.get_mut("plugin") else { + return Ok(()); + }; + let plugins = plugins.as_array_mut().ok_or_else(|| { + anyhow::anyhow!( + "OpenCode `plugin` config must be an array: {}", + path.display() + ) + })?; + let original_len = plugins.len(); + plugins.retain(|plugin| { + plugin.as_str().is_none_or(|plugin| { + plugin != "@braintrust/trace-opencode" + && !plugin.starts_with("@braintrust/trace-opencode@") + }) + }); + if plugins.len() != original_len { + write_object_atomic(path, config)?; + } + Ok(()) +} + +fn disable_opencode() -> anyhow::Result<()> { + let settings_path = paths::agent_settings_path("opencode", None); + let path = settings_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("opencode.json"); + remove_opencode_plugin_at(&path) +} + +fn setup_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugin = pi_plugin_spec(); + runner.run("pi", &["install", &plugin]) +} + +fn disable_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugin = pi_plugin_spec(); + runner.run("pi", &["uninstall", &plugin]) +} + +fn update_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let plugin = pi_plugin_spec(); + runner.run("pi", &["update", &plugin]) +} + +/// 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!( + "Antigravity configuration directory must end in `.gemini/config`: {}", + config_dir.display() + ); + } + let gemini_dir = config_dir.parent().ok_or_else(|| { + anyhow::anyhow!( + "Antigravity configuration directory has no parent: {}", + config_dir.display() + ) + })?; + if gemini_dir.file_name().and_then(|part| part.to_str()) != Some(".gemini") { + bail!( + "Antigravity configuration directory must end in `.gemini/config`: {}", + config_dir.display() + ); + } + gemini_dir.parent().ok_or_else(|| { + anyhow::anyhow!( + "Antigravity configuration directory has no home parent: {}", + config_dir.display() + ) + }) +} + +fn remove_legacy_antigravity_registration(config_dir: &Path) -> anyhow::Result<()> { + let hooks_path = config_dir.join("hooks.json"); + if !hooks_path.exists() { + return Ok(()); + } + let mut hooks = load_object(&hooks_path)?; + if hooks.remove(ANTIGRAVITY_PLUGIN).is_some() { + write_object_atomic(&hooks_path, hooks)?; + } + Ok(()) +} + +fn setup_antigravity_at(runner: &mut impl CommandRunner, config_dir: &Path) -> anyhow::Result<()> { + runner.run_in_home( + "agy", + &["plugin", "install", ANTIGRAVITY_PLUGIN_SOURCE], + antigravity_home(config_dir)?, + )?; + runner.run_in_home( + "agy", + &["plugin", "enable", ANTIGRAVITY_PLUGIN], + antigravity_home(config_dir)?, + )?; + + remove_legacy_antigravity_registration(config_dir) +} + +fn setup_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + setup_antigravity_at(runner, &paths::antigravity_config_dir()) +} + +fn disable_antigravity_at( + runner: &mut impl CommandRunner, + config_dir: &Path, +) -> anyhow::Result<()> { + // Antigravity's uninstall command is idempotent when the plugin is absent. + runner.run_in_home( + "agy", + &["plugin", "uninstall", ANTIGRAVITY_PLUGIN], + antigravity_home(config_dir)?, + )?; + remove_legacy_antigravity_registration(config_dir) +} + +fn disable_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + disable_antigravity_at(runner, &paths::antigravity_config_dir()) +} + +fn update_antigravity_at(runner: &mut impl CommandRunner, config_dir: &Path) -> anyhow::Result<()> { + let home = antigravity_home(config_dir)?; + let plugins = runner.json_in_home("agy", &["plugin", "list"], home)?; + let installed = plugins + .get("imports") + .and_then(Value::as_array) + .is_some_and(|imports| { + imports.iter().any(|plugin| { + plugin.get("name").and_then(Value::as_str) == Some(ANTIGRAVITY_PLUGIN) + }) + }); + if !installed { + bail!( + "Google Antigravity tracing plugin is not installed; run `bt trace enable antigravity`" + ); + } + runner.run_in_home( + "agy", + &["plugin", "install", ANTIGRAVITY_PLUGIN_SOURCE], + home, + ) +} + +fn update_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + update_antigravity_at(runner, &paths::antigravity_config_dir()) +} + +fn enable_tracing_at(path: &Path, mut route: SessionRoute) -> anyhow::Result<()> { + let mut settings = load_object(path)?; + if route.additional_metadata.is_none() { + route.additional_metadata = settings + .get("route") + .and_then(|route| route.get("additional_metadata")) + .or_else(|| settings.get("additional_metadata")) + .filter(|metadata| metadata.is_object()) + .cloned(); + } + if route.tags.is_empty() { + route.tags = settings + .get("route") + .and_then(|route| route.get("tags")) + .or_else(|| settings.get("tags")) + .and_then(Value::as_array) + .map(|tags| { + tags.iter() + .filter_map(Value::as_str) + .map(ToOwned::to_owned) + .collect() + }) + .unwrap_or_default(); + } + if route.span_plugins.is_empty() { + route.span_plugins = settings + .get("route") + .and_then(|route| route.get("span_plugins")) + .and_then(|plugins| serde_json::from_value(plugins.clone()).ok()) + .unwrap_or_default(); + } + settings.insert("trace_to_braintrust".into(), Value::Bool(true)); + settings.insert("route".into(), serde_json::to_value(route)?); + for key in [ + "traceToBraintrust", + "profile", + "org_name", + "api_key", + "api_url", + "app_url", + "project", + "destination", + "additional_metadata", + "tags", + ] { + settings.remove(key); + } + write_object_atomic(path, settings)?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)) + .with_context(|| format!("failed to protect agent settings: {}", path.display()))?; + } + Ok(()) +} + +fn enable_tracing(source: &str, route: SessionRoute) -> anyhow::Result { + let path = paths::agent_settings_path(source, None); + enable_tracing_at(&path, route)?; + Ok(path) +} + +fn remove_tracing_settings(path: &Path) -> anyhow::Result<()> { + match std::fs::remove_file(path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error) + .with_context(|| format!("failed to remove tracing settings: {}", path.display())), + } +} + +fn finish_disable(adapter_result: anyhow::Result<()>, settings_path: &Path) -> anyhow::Result<()> { + let settings_result = remove_tracing_settings(settings_path); + adapter_result.and(settings_result) +} + +/// Uninstall an agent's tracing adapter and remove its Braintrust-owned settings. +pub fn run_disable(agent: SetupAgent) -> anyhow::Result { + let mut runner = SystemCommandRunner; + let (source, display_name) = agent_details(agent); + let adapter_result = match agent { + SetupAgent::Codex => disable_codex(&mut runner), + SetupAgent::Claude => disable_claude(&mut runner), + 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); + finish_disable(adapter_result, &settings_path)?; + Ok(TraceCommandOutput::disable( + source, + display_name, + settings_path, + )) +} + +/// Update an already installed tracing adapter without writing tracing settings, +/// enabling a disabled plugin, or creating an agent configuration file. +pub fn run_update(agent: SetupAgent) -> anyhow::Result { + let mut runner = SystemCommandRunner; + let (source, display_name) = agent_details(agent); + match agent { + SetupAgent::Codex => update_codex(&mut runner)?, + SetupAgent::Claude => update_claude(&mut runner)?, + 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)) +} + +fn agent_details(agent: SetupAgent) -> (&'static str, &'static str) { + match agent { + SetupAgent::Codex => ("codex", "Codex"), + SetupAgent::Claude => ("claude", "Claude Code"), + SetupAgent::OpenCode => ("opencode", "OpenCode"), + SetupAgent::Pi => ("pi", "Pi"), + SetupAgent::Grok => ("grok", "Grok"), + SetupAgent::Muse => ("muse", "Muse Code"), + SetupAgent::Antigravity => ("antigravity", "Google Antigravity"), + } +} + +/// Install or refresh one agent's published tracing adapter and persist its +/// non-secret route selection. +pub fn run_enable(args: EnableArgs, route: SessionRoute) -> anyhow::Result { + let mut runner = SystemCommandRunner; + let (source, display_name) = match args.agent { + SetupAgent::Codex => { + setup_codex(&mut runner)?; + ("codex", "Codex") + } + SetupAgent::Claude => { + setup_claude(&mut runner)?; + warn_legacy_claude_tracing_env(); + ("claude", "Claude Code") + } + SetupAgent::OpenCode => { + setup_opencode()?; + ("opencode", "OpenCode") + } + SetupAgent::Pi => { + setup_pi(&mut runner)?; + ("pi", "Pi") + } + SetupAgent::Grok => { + setup_grok(&mut runner)?; + ("grok", "Grok") + } + SetupAgent::Muse => { + setup_muse()?; + ("muse", "Muse Code") + } + SetupAgent::Antigravity => { + setup_antigravity(&mut runner)?; + ("antigravity", "Google Antigravity") + } + }; + let settings_path = enable_tracing(source, route)?; + Ok(TraceCommandOutput::setup( + source, + display_name, + settings_path, + )) +} + +/// Backwards-compatible library entry point for hosts that used the former setup name. +pub fn run_setup(args: EnableArgs, route: SessionRoute) -> anyhow::Result { + run_enable(args, route) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::wire::{AuthSelection, TraceDestination}; + use std::collections::VecDeque; + + struct FakeRunner { + responses: VecDeque, + calls: Vec, + home_calls: Vec<(String, PathBuf)>, + } + + impl FakeRunner { + fn new(responses: impl IntoIterator) -> Self { + Self { + responses: responses.into_iter().collect(), + calls: Vec::new(), + home_calls: Vec::new(), + } + } + + fn called(&self, command: &str) -> bool { + self.calls.iter().any(|call| call == command) + } + } + + impl CommandRunner for FakeRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { + self.calls.push(format!("{program} {}", args.join(" "))); + self.responses + .pop_front() + .ok_or_else(|| anyhow::anyhow!("missing fake JSON response")) + } + + fn json_in_home( + &mut self, + program: &str, + args: &[&str], + home: &Path, + ) -> anyhow::Result { + self.calls.push(format!("{program} {}", args.join(" "))); + self.home_calls + .push((format!("{program} {}", args.join(" ")), home.to_path_buf())); + self.responses + .pop_front() + .ok_or_else(|| anyhow::anyhow!("missing fake JSON response")) + } + + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { + self.calls.push(format!("{program} {}", args.join(" "))); + Ok(()) + } + + fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()> { + self.calls.push(format!("{program} {}", args.join(" "))); + self.home_calls + .push((format!("{program} {}", args.join(" ")), home.to_path_buf())); + Ok(()) + } + } + + struct MissingAgyRunner; + + impl CommandRunner for MissingAgyRunner { + fn json(&mut self, _: &str, _: &[&str]) -> anyhow::Result { + unreachable!() + } + + fn json_in_home(&mut self, _: &str, _: &[&str], _: &Path) -> anyhow::Result { + unreachable!() + } + + fn run(&mut self, _: &str, _: &[&str]) -> anyhow::Result<()> { + unreachable!() + } + + fn run_in_home(&mut self, program: &str, _: &[&str], _: &Path) -> anyhow::Result<()> { + anyhow::bail!("failed to run `{program}`; install {program} and ensure it is on PATH") + } + } + + #[test] + fn codex_installs_from_the_published_marketplace_when_missing() { + let mut runner = FakeRunner::new([serde_json::json!({"marketplaces": []})]); + + setup_codex(&mut runner).unwrap(); + + assert!( + runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") + ); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + } + + #[test] + fn codex_refreshes_the_published_marketplace_and_plugin() { + let mut runner = FakeRunner::new([serde_json::json!({ + "marketplaces": [{ + "name": CODEX_MARKETPLACE, + "marketplaceSource": { + "sourceType": "github", + "source": CODEX_MARKETPLACE_SOURCE + } + }] + })]); + + setup_codex(&mut runner).unwrap(); + + assert!(runner.called("codex plugin marketplace upgrade braintrust-codex-plugins")); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + assert!(!runner.called("codex plugin marketplace remove braintrust-codex-plugins")); + } + + #[test] + fn codex_replaces_a_same_name_local_marketplace() { + let mut runner = FakeRunner::new([serde_json::json!({ + "marketplaces": [{ + "name": CODEX_MARKETPLACE, + "marketplaceSource": {"sourceType": "local", "source": "/tmp/stale"} + }] + })]); + + setup_codex(&mut runner).unwrap(); + + assert!(runner.called("codex plugin marketplace remove braintrust-codex-plugins")); + assert!( + runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") + ); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + } + + #[test] + fn claude_installs_from_the_published_marketplace_when_missing() { + let mut runner = FakeRunner::new([serde_json::json!([]), serde_json::json!([])]); + + setup_claude(&mut runner).unwrap(); + + assert!( + runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") + ); + assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); + } + + #[test] + fn claude_refreshes_the_published_marketplace_and_plugin() { + let mut runner = FakeRunner::new([ + serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "github", + "repo": CLAUDE_MARKETPLACE_SOURCE + }]), + serde_json::json!([{ + "id": CLAUDE_PLUGIN, + "version": "1.4.4", + "enabled": true + }]), + ]); + + setup_claude(&mut runner).unwrap(); + + assert!(runner.called("claude plugin marketplace update braintrust-claude-plugin")); + assert!(runner.called("claude plugin update trace-claude-code@braintrust-claude-plugin")); + assert!(!runner.called("claude plugin marketplace remove braintrust-claude-plugin")); + } + + #[test] + fn claude_replaces_a_same_name_local_marketplace() { + let mut runner = FakeRunner::new([serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "directory", + "path": "/tmp/stale" + }])]); + + setup_claude(&mut runner).unwrap(); + + assert!(runner.called("claude plugin marketplace remove braintrust-claude-plugin")); + assert!( + runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") + ); + assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); + } + + #[test] + fn claude_updates_then_enables_a_disabled_plugin() { + let mut runner = FakeRunner::new([ + serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "github", + "repo": CLAUDE_MARKETPLACE_SOURCE + }]), + serde_json::json!([{"id": CLAUDE_PLUGIN, "enabled": false}]), + ]); + + setup_claude(&mut runner).unwrap(); + + let update = runner + .calls + .iter() + .position(|call| { + call == "claude plugin update trace-claude-code@braintrust-claude-plugin" + }) + .unwrap(); + let enable = runner + .calls + .iter() + .position(|call| { + call == "claude plugin enable trace-claude-code@braintrust-claude-plugin" + }) + .unwrap(); + assert!(update < enable); + } + + #[test] + fn claude_legacy_env_diagnostic_is_read_only_and_preserves_api_key() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("settings.json"); + let contents = r#"{ + "env": { + "BRAINTRUST_CC_PROJECT": "old-project", + "BRAINTRUST_CC_DEBUG": "1", + "BRAINTRUST_API_KEY": "still-used-by-mcp", + "OTHER": "value" + } + }"#; + std::fs::write(&path, contents).unwrap(); + + assert_eq!( + legacy_claude_tracing_env_keys(&path).unwrap(), + ["BRAINTRUST_CC_PROJECT", "BRAINTRUST_CC_DEBUG"] + ); + assert_eq!(std::fs::read_to_string(path).unwrap(), contents); + } + + #[test] + fn grok_installs_published_plugin_then_enables_it_when_missing() { + let mut runner = FakeRunner::new([serde_json::json!([ + {"name": "other-plugin", "source": "somebody/other-plugin"} + ])]); + + setup_grok(&mut runner).unwrap(); + + assert_eq!( + runner.calls, + [ + "grok plugin list --json", + "grok plugin install braintrustdata/braintrust-grok-plugin --trust", + "grok plugin enable trace-grok", + ] + ); + } + + #[test] + fn grok_repeated_enable_updates_and_enables_the_published_plugin() { + let mut runner = FakeRunner::new([serde_json::json!([ + { + "name": GROK_PLUGIN, + "source": "https://github.com/braintrustdata/braintrust-grok-plugin.git", + "status": "installed" + }, + {"name": "other-plugin", "source": "somebody/other-plugin"} + ])]); + + setup_grok(&mut runner).unwrap(); + + assert_eq!( + runner.calls, + [ + "grok plugin list --json", + "grok plugin update trace-grok", + "grok plugin enable trace-grok", + ] + ); + } + + #[test] + fn grok_update_requires_the_published_plugin_without_enabling_it() { + let mut runner = FakeRunner::new([serde_json::json!([{ + "name": GROK_PLUGIN, + "source": GROK_PLUGIN_SOURCE, + "status": "installed" + }])]); + + update_grok(&mut runner).unwrap(); + + assert_eq!( + runner.calls, + ["grok plugin list --json", "grok plugin update trace-grok"] + ); + } + + #[test] + fn grok_enable_reconciles_only_a_conflicting_same_name_plugin() { + let mut runner = FakeRunner::new([serde_json::json!([ + { + "name": GROK_PLUGIN, + "source": "/tmp/local-trace-grok", + "status": "installed" + }, + {"name": "other-plugin", "source": "somebody/other-plugin"} + ])]); + + setup_grok(&mut runner).unwrap(); + + assert_eq!( + runner.calls, + [ + "grok plugin list --json", + "grok plugin uninstall trace-grok --confirm", + "grok plugin install braintrustdata/braintrust-grok-plugin --trust", + "grok plugin enable trace-grok", + ] + ); + } + + #[test] + fn grok_disable_removes_only_the_published_braintrust_plugin() { + let mut runner = FakeRunner::new([serde_json::json!([ + { + "name": GROK_PLUGIN, + "source": GROK_PLUGIN_SOURCE, + "status": "installed" + }, + {"name": "other-plugin", "source": "somebody/other-plugin"} + ])]); + + disable_grok(&mut runner).unwrap(); + + assert_eq!( + runner.calls, + [ + "grok plugin list --json", + "grok plugin uninstall trace-grok --confirm", + ] + ); + + let mut local = FakeRunner::new([serde_json::json!([{ + "name": GROK_PLUGIN, + "source": "/tmp/local-trace-grok", + "status": "installed" + }])]); + disable_grok(&mut local).unwrap(); + assert_eq!(local.calls, ["grok plugin list --json"]); + } + + #[test] + fn opencode_reconciles_the_published_plugin_and_preserves_config() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("opencode.json"); + std::fs::write( + &path, + r#"{"plugin":["other","@braintrust/trace-opencode@0.9.0"],"model":"test/model"}"#, + ) + .unwrap(); + + setup_opencode_at(&path).unwrap(); + + let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert_eq!(config["model"], "test/model"); + assert_eq!( + config["plugin"], + serde_json::json!(["other", opencode_plugin_spec().unwrap()]) + ); + } + + #[test] + fn pi_installs_the_latest_published_extension() { + let mut runner = FakeRunner::new([]); + + setup_pi(&mut runner).unwrap(); + + assert!(runner.called("pi install npm:@braintrust/pi-extension")); + } + + #[test] + fn pi_updates_the_npm_extension_without_installing_it() { + let mut runner = FakeRunner::new([]); + + update_pi(&mut runner).unwrap(); + + assert!(runner.called("pi update npm:@braintrust/pi-extension")); + assert!(!runner + .calls + .iter() + .any(|call| call.starts_with("pi install "))); + } + + #[test] + fn opencode_update_requires_an_existing_plugin_and_preserves_other_config() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("opencode.json"); + std::fs::write( + &path, + r#"{"plugin":["other","@braintrust/trace-opencode@^1"],"model":"test/model"}"#, + ) + .unwrap(); + + update_opencode_at(&path).unwrap(); + + let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert_eq!(config["model"], "test/model"); + assert_eq!( + config["plugin"], + serde_json::json!(["other", opencode_plugin_spec().unwrap()]) + ); + + let missing = temp.path().join("missing.json"); + assert!(update_opencode_at(&missing).is_err()); + assert!(!missing.exists()); + } + + #[test] + fn codex_update_refuses_to_install_a_missing_plugin() { + let mut runner = FakeRunner::new([serde_json::json!({"installed": []})]); + + assert!(update_codex(&mut runner).is_err()); + assert!(!runner + .calls + .iter() + .any(|call| call.starts_with("codex plugin add "))); + } + + #[test] + fn antigravity_installs_published_plugin_and_removes_legacy_registration() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + std::fs::create_dir_all(&config_dir).unwrap(); + let hooks_path = config_dir.join("hooks.json"); + std::fs::write( + &hooks_path, + serde_json::to_vec(&serde_json::json!({ + ANTIGRAVITY_PLUGIN: {"Stop": []}, + "other-plugin": {"Stop": [{"type": "command", "command": "other"}]} + })) + .unwrap(), + ) + .unwrap(); + let mut runner = FakeRunner::new([]); + + setup_antigravity_at(&mut runner, &config_dir).unwrap(); + + assert!(runner.called(&format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"))); + assert!(runner.called(&format!("agy plugin enable {ANTIGRAVITY_PLUGIN}"))); + let hooks: Value = serde_json::from_slice(&std::fs::read(&hooks_path).unwrap()).unwrap(); + assert_eq!(hooks["other-plugin"]["Stop"][0]["command"], "other"); + assert!(hooks.get(ANTIGRAVITY_PLUGIN).is_none()); + } + + #[test] + fn antigravity_update_checks_and_updates_the_same_overridden_home() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + std::fs::create_dir_all(&config_dir).unwrap(); + let expected_home = temp.path().to_path_buf(); + let mut runner = FakeRunner::new([serde_json::json!({ + "imports": [{"name": ANTIGRAVITY_PLUGIN}], + })]); + + update_antigravity_at(&mut runner, &config_dir).unwrap(); + + assert_eq!( + runner.home_calls, + vec![ + ("agy plugin list".into(), expected_home.clone()), + ( + format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"), + expected_home, + ), + ] + ); + } + + #[test] + fn antigravity_setup_is_idempotent() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + std::fs::create_dir_all(&config_dir).unwrap(); + std::fs::write( + config_dir.join("hooks.json"), + serde_json::to_vec(&serde_json::json!({ + ANTIGRAVITY_PLUGIN: {"Stop": []}, + "other-plugin": {"Stop": []} + })) + .unwrap(), + ) + .unwrap(); + let mut runner = FakeRunner::new([]); + + setup_antigravity_at(&mut runner, &config_dir).unwrap(); + let first = std::fs::read(config_dir.join("hooks.json")).unwrap(); + setup_antigravity_at(&mut runner, &config_dir).unwrap(); + let second = std::fs::read(config_dir.join("hooks.json")).unwrap(); + + assert_eq!(first, second); + assert_eq!( + runner + .calls + .iter() + .filter(|call| { + call.as_str() == format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}") + }) + .count(), + 2 + ); + } + + #[test] + fn antigravity_setup_relies_on_native_plugin_hooks() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + let mut runner = FakeRunner::new([]); + + setup_antigravity_at(&mut runner, &config_dir).unwrap(); + + assert!(!config_dir.join("hooks.json").exists()); + assert!(runner.called(&format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"))); + assert!(runner.called(&format!("agy plugin enable {ANTIGRAVITY_PLUGIN}"))); + } + + #[test] + fn antigravity_setup_reports_a_missing_cli_without_changing_hooks() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + std::fs::create_dir_all(&config_dir).unwrap(); + let hooks_path = config_dir.join("hooks.json"); + let original = br#"{"other-plugin":{"Stop":[]}}"#; + std::fs::write(&hooks_path, original).unwrap(); + + let error = setup_antigravity_at(&mut MissingAgyRunner, &config_dir).unwrap_err(); + + assert_eq!( + error.to_string(), + "failed to run `agy`; install agy and ensure it is on PATH" + ); + assert_eq!(std::fs::read(hooks_path).unwrap(), original); + } + + #[test] + fn antigravity_disable_removes_only_the_managed_registration() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join(".gemini/config"); + std::fs::create_dir_all(&config_dir).unwrap(); + let hooks_path = config_dir.join("hooks.json"); + std::fs::write( + &hooks_path, + serde_json::to_vec(&serde_json::json!({ + ANTIGRAVITY_PLUGIN: {"Stop": []}, + "other-plugin": {"Stop": [{"type": "command", "command": "other"}]} + })) + .unwrap(), + ) + .unwrap(); + let mut runner = FakeRunner::new([]); + + disable_antigravity_at(&mut runner, &config_dir).unwrap(); + + assert!(runner.called("agy plugin uninstall braintrust-antigravity-tracing")); + let hooks: Value = serde_json::from_slice(&std::fs::read(&hooks_path).unwrap()).unwrap(); + assert!(hooks.get(ANTIGRAVITY_PLUGIN).is_none()); + assert_eq!(hooks["other-plugin"]["Stop"][0]["command"], "other"); + } + + #[test] + fn tracing_settings_preserve_unrelated_fields_and_remove_legacy_keys() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write( + &path, + r#"{ + "traceToBraintrust": false, + "profile": "stale-profile", + "org_name": "stale-org", + "api_key": "stale-secret", + "api_url": "https://stale-api.example", + "app_url": "https://stale-app.example", + "project": "old", + "destination": "old-destination", + "additional_metadata": {"migrated": true}, + "auth": {"type": "legacy"} + }"#, + ) + .unwrap(); + let route = SessionRoute { + auth: AuthSelection { + source: crate::wire::AuthSource::SavedProfile, + profile_id: None, + profile: Some("work".into()), + org_name: Some("Braintrust SDKs".into()), + }, + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("coding-agents".into()), + }), + ..SessionRoute::default() + }; + + enable_tracing_at(&path, route).unwrap(); + + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!(settings["trace_to_braintrust"], true); + assert_eq!(settings["route"]["auth"]["profile"], "work"); + assert_eq!( + settings["route"]["destination"]["project_name"], + "coding-agents" + ); + assert_eq!(settings["auth"]["type"], "legacy"); + assert_eq!(settings["route"]["additional_metadata"]["migrated"], true); + for key in [ + "traceToBraintrust", + "profile", + "org_name", + "api_key", + "api_url", + "app_url", + "project", + "destination", + "additional_metadata", + ] { + assert!(settings.get(key).is_none(), "legacy key {key} remained"); + } + } + + #[test] + fn tracing_settings_preserve_metadata_until_setup_explicitly_replaces_it() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write( + &path, + r#"{"route":{"additional_metadata":{"ci":true},"tags":["saved"]}}"#, + ) + .unwrap(); + + let route = SessionRoute::default(); + enable_tracing_at(&path, route).unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["additional_metadata"], + serde_json::json!({"ci": true}) + ); + assert_eq!(settings["route"]["tags"], serde_json::json!(["saved"])); + + let route = SessionRoute { + additional_metadata: Some(serde_json::json!({"run_id": "new"})), + tags: vec!["replacement".to_string()], + ..SessionRoute::default() + }; + enable_tracing_at(&path, route).unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["additional_metadata"], + serde_json::json!({"run_id": "new"}) + ); + assert_eq!( + settings["route"]["tags"], + serde_json::json!(["replacement"]) + ); + } + + #[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(); + let path = temp.path().join("braintrust.json"); + std::fs::write( + &path, + r#"{"traceToBraintrust":true,"route":{"destination":{"project_name":"coding-agents"}},"other":true}"#, + ) + .unwrap(); + + remove_tracing_settings(&path).unwrap(); + assert!(!path.exists()); + } + + #[test] + fn disabling_removes_settings_even_when_adapter_cleanup_fails() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write(&path, "{}").unwrap(); + + let error = finish_disable(Err(anyhow::anyhow!("adapter unavailable")), &path) + .unwrap_err() + .to_string(); + + assert_eq!(error, "adapter unavailable"); + assert!(!path.exists()); + } + + #[test] + fn disabling_installed_plugins_uses_each_agents_uninstall_command() { + let mut codex = FakeRunner::new([serde_json::json!({ + "installed": [{"pluginId": CODEX_PLUGIN}] + })]); + disable_codex(&mut codex).unwrap(); + assert!(codex.called("codex plugin remove trace-codex@braintrust-codex-plugins --json")); + + let mut claude = FakeRunner::new([serde_json::json!([{"id": CLAUDE_PLUGIN}])]); + disable_claude(&mut claude).unwrap(); + assert!(claude.called("claude plugin uninstall trace-claude-code@braintrust-claude-plugin")); + + let mut pi = FakeRunner::new([]); + disable_pi(&mut pi).unwrap(); + assert!(pi.called("pi uninstall npm:@braintrust/pi-extension")); + } + + #[test] + fn disabling_opencode_removes_only_the_managed_plugin() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("opencode.json"); + std::fs::write( + &path, + serde_json::to_vec(&serde_json::json!({ + "plugin": ["other", opencode_plugin_spec().unwrap()], + "model": "test/model", + })) + .unwrap(), + ) + .unwrap(); + + remove_opencode_plugin_at(&path).unwrap(); + + let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert_eq!(config["plugin"], serde_json::json!(["other"])); + assert_eq!(config["model"], "test/model"); + } + + #[test] + fn tracing_settings_preserve_plugins_until_setup_explicitly_replaces_them() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write(&path, r#"{"route":{"span_plugins":["old.mjs"]}}"#).unwrap(); + + enable_tracing_at(&path, SessionRoute::default()).unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["span_plugins"], + serde_json::json!(["old.mjs"]) + ); + + enable_tracing_at( + &path, + SessionRoute { + span_plugins: vec![PathBuf::from("first.mjs"), PathBuf::from("second.mjs")], + ..SessionRoute::default() + }, + ) + .unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["span_plugins"], + serde_json::json!(["first.mjs", "second.mjs"]) + ); + } +} diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs index e2d80014..eb308473 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -4,7 +4,6 @@ //! integrations they control. Hosts such as `bt` provide global auth flags and //! dispatch these commands without duplicating agent-specific CLI knowledge. -use crate::agents::{self, Agent, Identity}; use crate::{HookArgs, ImportArgs, RunArgs, ServeArgs, StatusArgs}; use clap::{Args, Subcommand, ValueEnum}; use std::path::PathBuf; @@ -30,7 +29,7 @@ pub enum TraceCommand { /// Run the tracing daemon (foreground). #[command(hide = true)] Daemon(ServeArgs), - /// Capture one coding-agent hook event from stdin; exits nonzero if capture fails. + /// Forward one coding-agent hook event (read from stdin) to the daemon. #[command(hide = true)] Hook(HookArgs), /// Print daemon/session status. @@ -64,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")] @@ -71,19 +71,30 @@ pub enum DoctorAgent { Pi, Antigravity, Grok, - Cursor, } impl DoctorAgent { - pub(crate) fn identity(self) -> &'static Identity { + pub(crate) fn source(self) -> &'static str { match self { - Self::Codex => agents::Codex.identity(), - Self::Claude => agents::Claude.identity(), - Self::OpenCode => agents::OpenCode.identity(), - Self::Pi => agents::Pi.identity(), - Self::Antigravity => agents::Antigravity.identity(), - Self::Grok => agents::Grok.identity(), - Self::Cursor => agents::Cursor.identity(), + Self::Codex => "codex", + Self::Muse => "muse", + Self::Claude => "claude", + Self::OpenCode => "opencode", + Self::Pi => "pi", + Self::Antigravity => "antigravity", + Self::Grok => "grok", + } + } + + 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", + Self::Antigravity => "Google Antigravity", + Self::Grok => "Grok", } } } @@ -128,6 +139,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. @@ -137,26 +150,10 @@ pub enum SetupAgent { Pi, /// Install the published Grok tracing plugin. Grok, - /// Install the published Cursor tracing plugin. - Cursor, /// Install the Google Antigravity tracing hooks. Antigravity, } -impl SetupAgent { - pub(crate) fn identity(self) -> &'static Identity { - match self { - Self::Codex => agents::Codex.identity(), - Self::Claude => agents::Claude.identity(), - Self::OpenCode => agents::OpenCode.identity(), - Self::Pi => agents::Pi.identity(), - Self::Grok => agents::Grok.identity(), - Self::Cursor => agents::Cursor.identity(), - Self::Antigravity => agents::Antigravity.identity(), - } - } -} - #[cfg(test)] mod tests { use super::*; @@ -245,16 +242,6 @@ mod tests { .. }) if value == r#"{"import":true}"# )); - - let cursor_import = Cli::try_parse_from(["bt", "import", "cursor", "session-id"]).unwrap(); - assert!(matches!( - cursor_import.trace.command, - TraceCommand::Import(ImportArgs { - source: crate::ImportSource::Cursor, - ref session_ids, - .. - }) if session_ids == &["session-id"] - )); } #[test] @@ -288,27 +275,11 @@ mod tests { )); } - #[test] - fn run_accepts_cursor_and_forwards_agent_arguments() { - let parsed = Cli::try_parse_from([ - "bt", "run", "cursor", "--tag", "ci", "--", "--print", "hello", - ]) - .unwrap(); - assert!(matches!( - parsed.trace.command, - TraceCommand::Run(RunArgs { - source: crate::RunSource::Cursor, - ref agent_args, - ref tags, - .. - }) if agent_args == &["--print", "hello"] && tags == &["ci"] - )); - } - #[test] 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"), @@ -319,15 +290,14 @@ mod tests { "Google Antigravity", ), ("grok", DoctorAgent::Grok, "grok", "Grok"), - ("cursor", DoctorAgent::Cursor, "cursor", "Cursor"), ] { let cli = Cli::try_parse_from(["bt", "doctor", agent]).unwrap(); assert!(matches!( cli.trace.command, TraceCommand::Doctor(DoctorArgs { agent }) if agent == expected )); - assert_eq!(expected.identity().id, source); - assert_eq!(expected.identity().display_name, display_name); + assert_eq!(expected.source(), source); + assert_eq!(expected.display_name(), display_name); } } @@ -356,7 +326,28 @@ mod tests { } #[test] - fn grok_uses_shared_enable_and_disable_commands() { + 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"] { let parsed = Cli::try_parse_from(["bt", command, "grok"]).unwrap(); assert!(matches!( @@ -375,6 +366,17 @@ mod tests { agent: SetupAgent::Grok }) )); + + let parsed = Cli::try_parse_from(["bt", "doctor", "grok"]).unwrap(); + assert!(matches!( + parsed.trace.command, + TraceCommand::Doctor(DoctorArgs { + agent: DoctorAgent::Grok + }) + )); + + assert_eq!(DoctorAgent::Grok.source(), "grok"); + assert_eq!(DoctorAgent::Grok.display_name(), "Grok"); } #[test] @@ -385,7 +387,7 @@ mod tests { "opencode", "pi", "grok", - "cursor", + "muse", "antigravity", ] { assert!(matches!( diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index 52327dea..754a91b2 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -10,22 +10,24 @@ mod antigravity; mod claude; mod codex; -pub(crate) mod cursor; mod debug; -mod decode; mod git; mod grok; +mod muse; mod opencode; mod pi; -pub(crate) use pi::request_config as pi_request_config; mod recent; mod tool; -mod turn_lineage; -pub use turn_lineage::TURN_SPAN_ID_KEY; -use debug::DebugTranslatorFactory; +pub use antigravity::AntigravityTranslatorFactory; +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; -use crate::agents::{registrar, Agent}; use crate::wire::{Envelope, SessionConfig}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -52,14 +54,6 @@ pub enum SpanType { Wait, } -/// Immutable versions captured when a coding-agent span is first created. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct OriginSnapshot { - pub plugin_version: Option, - pub bt_version: String, - pub source_version: Option, -} - /// A resolved span row, ready for a sink to insert or merge. Field set is the /// subset every current plugin uses; extend as translators need more. #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -90,21 +84,9 @@ pub struct SpanRow { #[serde(skip)] #[doc(hidden)] pub late_merge_key: Option, - /// Span id of the user turn that owns this row, recorded by the translator - /// from its native turn tracking. [`turn_lineage`] publishes it as - /// `metadata.turn_span_id`. - #[serde(skip)] - #[doc(hidden)] - pub turn_span_id: Option, /// Labels for filtering in Braintrust (e.g. `compaction`, `permission-request`). #[serde(default, skip_serializing_if = "Option::is_none")] pub tags: Option>, - /// Translator-owned marker for turn creation, not other task spans. - #[serde(skip)] - pub is_turn: bool, - /// Delivery-owned creation snapshot, present on each coding-agent span. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub origin: Option, } /// A span operation. `Insert` creates (or replaces) a row; `Merge` updates an @@ -178,27 +160,15 @@ pub trait AgentTranslator: Send { } } -/// Builds translator instances for an agent's hook source. Each agent type -/// implements this in its translator module and registers itself in -/// [`crate::agents`]; the source is the agent's [`Identity::source`]. -/// -/// [`Identity::source`]: crate::agents::Identity::source -pub(crate) trait TranslatorFactory: Agent { +/// Builds translator instances for a given `source`. +pub trait TranslatorFactory: Send + Sync { + fn source(&self) -> &str; fn create(&self, session_id: &str) -> Box; } /// Maps canonical and supported alias source strings to translator factories. pub struct Registry { - factories: HashMap>, -} - -/// Map an agent name or alias to the source its hooks report. Unknown -/// sources pass through. -pub(crate) fn canonical_source_name(source: &str) -> &str { - registrar() - .translate - .get(source) - .map_or(source, |agent| agent.identity().source) + factories: HashMap>, } impl Registry { @@ -207,16 +177,20 @@ impl Registry { let mut r = Registry { factories: HashMap::new(), }; - r.insert(Arc::new(DebugTranslatorFactory)); - for agent in registrar().translate.shared() { - r.insert(agent.clone()); - } + r.register(Box::new(DebugTranslatorFactory)); + let git = Arc::new(git::GitMetadataCache::default()); + r.register(Box::new(AntigravityTranslatorFactory::new(git.clone()))); + 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 } - fn insert(&mut self, factory: Arc) { - self.factories - .insert(factory.identity().source.to_string(), factory); + pub fn register(&mut self, factory: Box) { + self.factories.insert(factory.source().to_string(), factory); } /// Known sources, for the `initialize` capabilities list. @@ -228,11 +202,16 @@ impl Registry { /// Resolve daemon source aliases to one stable identity. pub fn canonical_source<'a>(&'a self, source: &'a str) -> Option<&'a str> { - // Only agent sources and the debug translator are accepted. - let canonical = match registrar().translate.get(source) { - Some(agent) => agent.identity().source, - None if source == "debug" => "debug", - None => return None, + let canonical = match source { + "claude" | "claude-code" => "claude-code", + "open-code" | "opencode" => "opencode", + "antigravity" => "antigravity", + "codex" => "codex", + "grok" => "grok", + "muse" | "muse-code" => "muse", + "pi" => "pi", + "debug" => "debug", + _ => return None, }; self.factories.contains_key(canonical).then_some(canonical) } @@ -252,21 +231,17 @@ impl Registry { )) } - pub(crate) fn create_checked_with_session_namespace( + pub(crate) fn create_checked_with_session_key( &self, source: &str, - session_namespace: &str, + session_key: &str, ) -> anyhow::Result> { let canonical = self .canonical_source(source) .ok_or_else(|| anyhow::anyhow!("unsupported coding-agent source {source:?}"))?; let factory = self.factories.get(canonical); - let factory = factory.ok_or_else(|| { - anyhow::anyhow!("no translator factory for coding-agent source {canonical:?}") - })?; - Ok(Box::new(turn_lineage::TurnLineage::new( - factory.create(session_namespace), - ))) + let factory = factory.expect("canonical source must have a factory"); + Ok(factory.create(session_key)) } /// Create a known translator. Production ingress uses @@ -277,43 +252,3 @@ impl Registry { .unwrap_or_else(|error| panic!("{error}")) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn every_agent_source_has_a_translator() { - assert_eq!( - Registry::default_agents().sources(), - [ - "antigravity", - "claude-code", - "codex", - "cursor", - "debug", - "grok", - "opencode", - "pi" - ] - ); - } - - #[test] - fn source_aliases_resolve_to_hook_sources() { - let registry = Registry::default_agents(); - for (alias, source) in [ - ("claude", "claude-code"), - ("claude-code", "claude-code"), - ("open-code", "opencode"), - ("opencode", "opencode"), - ("codex", "codex"), - ("debug", "debug"), - ] { - assert_eq!(registry.canonical_source(alias), Some(source)); - } - assert_eq!(registry.canonical_source("agy"), None); - assert_eq!(canonical_source_name("claude"), "claude-code"); - assert_eq!(canonical_source_name("unknown"), "unknown"); - } -} diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs new file mode 100644 index 00000000..4e69f1ef --- /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"))))); + } +} From fb17b58619821aaadd02937a84f386b62ec85d8d Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 11 Sep 2026 04:00:16 +0800 Subject: [PATCH 02/17] Import Muse export sessions --- bt-daemon/src/lib.rs | 2392 ++++++++++++++++++++++- bt-daemon/src/main.rs | 2 +- bt-daemon/src/transcript_import/mod.rs | 750 +++---- bt-daemon/src/transcript_import/muse.rs | 254 +++ bt-daemon/tests/replay.rs | 37 + 5 files changed, 2880 insertions(+), 555 deletions(-) create mode 100644 bt-daemon/src/transcript_import/muse.rs diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 61a3f58f..1386e87c 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -12,58 +12,39 @@ pub mod paths; -mod agents; -mod args; mod client; mod command_output; mod correlation; mod delivery_ledger; mod dispatch; -mod env; -mod hook; mod ids; mod journal; mod managed_run; mod pipeline; mod plugin_diagnostics; pub(crate) mod process; -mod route; mod server; mod settings; mod setup; mod sink; mod span_processor; -mod subprocess; mod trace_command; mod trace_runtime; mod transcript_import; mod transcript_mirror; mod translate; mod transport; -#[cfg(windows)] -mod win_acl; pub mod wire; -pub use args::{ - HookArgs, ImportArgs, ImportSource, ParentObjectType, RunArgs, RunHookCommand, RunSource, - ServeArgs, StatusArgs, -}; -pub use client::{ - flush_managed_run, flush_session, forward_envelope, run_status, shutdown_daemon, HostInfo, -}; +pub use client::HostInfo; pub use command_output::{ - AuthDiagnostic, DaemonDiagnostic, DaemonStatus, DoctorCommandOutput, ImportSummary, - OutputFormat, SetupCommandOutput, StatusCommandOutput, StopCommandOutput, TraceCommandOutput, + AuthDiagnostic, DoctorCommandOutput, ImportSummary, OutputFormat, SetupCommandOutput, + StatusCommandOutput, StopCommandOutput, TraceCommandOutput, }; -pub use hook::run_hook; #[doc(hidden)] pub use journal::source_journal_path; -pub use managed_run::run_traced; pub use plugin_diagnostics::PluginDiagnostic; -pub use server::{ - braintrust_serve_options, debug_serve_options, run_serve, AuthLease, AuthProvider, - AuthResolveReason, ServeOptions, -}; +pub use server::{AuthLease, AuthProvider, AuthResolveReason, ServeOptions}; pub use setup::{run_disable, run_enable, run_setup, run_update}; pub use sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, Sink, SinkFactory}; pub use trace_command::{ @@ -73,8 +54,2367 @@ pub use trace_command::{ pub use trace_runtime::{ run_trace, HostRouteSelection, RouteRequirements, TraceHostContext, TraceHostServices, }; -pub use transcript_import::{import_transcript, import_transcripts, run_import}; pub use translate::{ - AgentTranslator, OriginSnapshot, Registry, SessionCtx, SpanOp, SpanRow, SpanType, - TURN_SPAN_ID_KEY, + AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, +}; + +use anyhow::Context; +use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; +use clap::{Args, ValueEnum}; +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; +use wire::{ + method, Envelope, ManagedRunFlushParams, SessionConfig, SessionRoute, StatusResult, + PROTOCOL_VERSION, }; + +const MANAGED_RUN_ID_ENV: &str = "BT_TRACE_MANAGED_RUN_ID"; +const MANAGED_RUN_FLUSH_TIMEOUT_MS: u64 = 10_000; + +/// Arguments for `serve`. +#[derive(Debug, Clone, Args)] +pub struct ServeArgs { + /// Socket path override (default: see docs/protocol.md). + #[arg(long)] + pub socket: Option, + /// Data/journal directory override. + #[arg(long)] + pub data_dir: Option, + /// Exit after this many seconds idle (no activity, empty queues). 0 + /// disables the watchdog. + #[arg(long, default_value_t = 300)] + pub idle_timeout_secs: u64, + /// Retire a session's in-memory state after this many seconds without + /// traffic. A later event rebuilds it from the journal. 0 disables + /// retirement. + #[arg(long, default_value_t = 30)] + pub session_idle_timeout_secs: u64, +} + +/// Arguments for `hook`. +#[derive(Debug, Clone, Args)] +pub struct HookArgs { + /// Which translator should interpret this event's payload. + #[arg(long)] + pub source: String, + /// Optional agent version, forwarded for payload-drift handling. + #[arg(long)] + pub source_version: Option, + /// Optional instrumentation package version, forwarded independently from + /// the coding agent version. + #[arg(long)] + pub plugin_version: Option, + /// Socket path override. + #[arg(long)] + pub socket: Option, + /// JSON field in the payload holding the session id. + #[arg(long, default_value = "session_id")] + pub session_id_field: String, + /// JSON field in the payload holding the event name. + #[arg(long, default_value = "hook_event_name")] + pub event_field: String, + /// Explicit event name (overrides `--event-field` lookup). + #[arg(long)] + pub event: Option, + /// JSON field holding a transcript path. When present, capture the file + /// length observed by this hook so deterministic journal replay cannot + /// read transcript records written by later lifecycle events. + #[arg(long)] + pub transcript_path_field: Option, + /// Fail instead of spawning a daemon if none is running. + #[arg(long)] + pub no_spawn: bool, + /// Ask the daemon to flush the session after a turn-ending event. The + /// flush is scheduled out-of-band; hook capture still returns immediately + /// after the durable journal write. + #[arg(long)] + pub flush_on_turn_end: bool, + /// Deprecated compatibility option. Hook capture never waits for daemon + /// translation, reporting, or flushing. + #[arg(long, default_value_t = 10_000, hide = true)] + pub flush_timeout_ms: u64, + /// JSON object merged into root-span metadata. Deliberately not read from + /// the environment: a hook fires automatically on every event, so its + /// configuration must come only from the persisted agent route (or, for a + /// managed child, the invocation settings `run` injected). + #[arg(long)] + pub additional_metadata: Option, + /// Marks the hook definition injected by `run`; inherited plugin hooks do + /// not carry this flag and are suppressed for the managed child. + #[arg(long, hide = true)] + pub managed_run_hook: bool, +} + +/// Arguments for `status`. +#[derive(Debug, Clone, Args)] +pub struct StatusArgs { + #[arg(long)] + pub socket: Option, + /// Limit to one session. + #[arg(long)] + pub session_id: Option, +} + +/// Arguments for importing a past coding-agent session. +#[derive(Debug, Clone, Args)] +pub struct ImportArgs { + /// Agent that produced the session. + #[arg(value_enum)] + pub source: ImportSource, + /// 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( + long, + value_name = "DESTINATION", + conflicts_with_all = ["parent", "parent_span_id", "parent_root_span_id", "parent_object_type", "parent_object_id", "parent_project"] + )] + pub destination: Option, + /// Attach below the opaque value returned by a Braintrust SDK span.export(). + #[arg( + long, + value_name = "SPAN_EXPORT", + conflicts_with_all = ["destination", "parent_span_id", "parent_root_span_id", "parent_object_type", "parent_object_id", "parent_project"] + )] + pub parent: Option, + /// Attach below this span ID. Also requires --parent-root-span-id and + /// --parent-object-type, plus --parent-object-id or --parent-project. + #[arg( + long, + value_name = "SPAN_ID", + requires_all = ["parent_root_span_id", "parent_object_type"], + conflicts_with_all = ["destination", "parent"] + )] + pub parent_span_id: Option, + /// Root span ID for --parent-span-id. + #[arg(long, value_name = "ROOT_SPAN_ID", requires = "parent_span_id")] + pub parent_root_span_id: Option, + /// Braintrust object type for --parent-span-id. + #[arg(long, value_enum, requires = "parent_span_id")] + pub parent_object_type: Option, + /// Braintrust object ID for --parent-span-id. + #[arg( + long, + value_name = "OBJECT_ID", + requires = "parent_span_id", + conflicts_with = "parent_project" + )] + pub parent_object_id: Option, + /// Braintrust project name for a project-logs parent. + #[arg( + long, + value_name = "PROJECT", + requires = "parent_span_id", + conflicts_with = "parent_object_id" + )] + pub parent_project: Option, + /// Keep following the transcript until Ctrl-C, importing new turns as the + /// coding-agent session grows. + #[arg(long, conflicts_with = "all")] + pub attach: bool, + /// JSON object merged into every imported root span's metadata. + #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] + pub additional_metadata: Option, + /// Tag applied to each imported root span. May be repeated or comma-separated. + #[arg(long = "tag", env = "BRAINTRUST_TAGS", value_delimiter = ',')] + pub tags: Vec, + /// JavaScript span transform for this import. Repeat to compose an isolated + /// transform chain; persistent setup plugins are not included. + #[arg(long, value_name = "PATH")] + pub plugin: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum ImportSource { + Codex, + Muse, + #[value(name = "claude", alias = "claude-code")] + Claude, + #[value(name = "antigravity", alias = "agy")] + Antigravity, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum ParentObjectType { + Experiment, + #[value(name = "project_logs", alias = "project-logs")] + ProjectLogs, + #[value(name = "playground_logs", alias = "playground-logs")] + PlaygroundLogs, +} + +impl From for SpanObjectType { + fn from(value: ParentObjectType) -> Self { + match value { + ParentObjectType::Experiment => Self::Experiment, + ParentObjectType::ProjectLogs => Self::ProjectLogs, + ParentObjectType::PlaygroundLogs => Self::PlaygroundLogs, + } + } +} + +impl ImportArgs { + pub fn has_destination_override(&self) -> bool { + self.destination.is_some() || self.parent.is_some() || self.parent_span_id.is_some() + } +} + +/// Arguments for launching a coding agent with invocation-local live hooks. +#[derive(Debug, Clone, Args)] +#[command(trailing_var_arg = true)] +pub struct RunArgs { + /// Coding agent to launch. + #[arg(value_enum)] + pub source: RunSource, + /// JSON object merged into root-span metadata for this invocation. + #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] + pub additional_metadata: Option, + /// Tag applied to each root span for this invocation. May be repeated or comma-separated. + #[arg(long = "tag", env = "BRAINTRUST_TAGS", value_delimiter = ',')] + pub tags: Vec, + /// JavaScript span transform for this invocation. Repeat to compose an + /// isolated transform chain; persistent setup plugins are not included. + #[arg(long, value_name = "PATH")] + pub plugin: Vec, + /// Arguments forwarded verbatim to the coding agent. + #[arg(allow_hyphen_values = true)] + pub agent_args: Vec, +} + +/// Front-end command used by a managed agent run to forward one hook payload. +/// +/// The standalone binary uses `[bt-daemon, hook]`; the embedded `bt` front-end +/// uses its own equivalent prefix. `run_traced` appends `--source `. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RunHookCommand { + pub program: OsString, + pub args: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum RunSource { + Codex, + #[value(name = "claude", alias = "claude-code")] + Claude, + #[value(name = "opencode", alias = "open-code")] + OpenCode, + Pi, +} + +/// Run the daemon until shutdown. +pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { + server::run(args, opts).await +} + +fn build_hook_envelope( + args: &HookArgs, + route: SessionRoute, + payload: serde_json::Value, + session_id: String, + event: String, +) -> Envelope { + Envelope { + source: args.source.clone(), + source_version: args.source_version.clone(), + plugin_version: args.plugin_version.clone(), + session_id, + event, + ts_ms: now_ms(), + managed_run_id: std::env::var(MANAGED_RUN_ID_ENV) + .ok() + .filter(|value| !value.is_empty()), + capture: None, + payload, + route: Some(route), + config: None, + } +} + +pub(crate) fn should_flush_hook_event(event: &str, flush_on_turn_end: bool) -> bool { + matches!(event, "SessionEnd" | "session_end") + || (flush_on_turn_end + && matches!( + event, + "Stop" + | "stop" + | "StopFailure" + | "stop_failure" + | "StopCancelled" + | "stop_cancelled" + | "SubagentStop" + | "subagent_stop" + )) +} + +pub(crate) fn should_flush_ingress_event(env: &wire::Envelope) -> bool { + let native = env.payload.get("event").unwrap_or(&env.payload); + let flush_on_turn_end = matches!( + env.route.as_ref().map(|route| route.flush_mode), + Some(wire::FlushMode::FlushOnTurnEnd) + ); + should_flush_hook_event(&env.event, flush_on_turn_end) + || match env.source.as_str() { + "pi" => match env.event.as_str() { + // Preserve Pi's explicit lifecycle flushes even in batched mode. + "session_shutdown" | "session_compact" | "session_tree" => true, + // OMP can end an attempt while keeping the same turn open for retry. + "agent_end" => { + flush_on_turn_end + && native.get("willRetry").and_then(serde_json::Value::as_bool) + != Some(true) + } + _ => false, + }, + "opencode" => matches!( + env.event.as_str(), + "session.idle" | "session.deleted" | "session.error" | "server.instance.disposed" + ), + _ => false, + } +} + +/// Capture one hook event from `stdin` and forward it to the daemon. +/// +/// `route` contains only non-secret profile and destination selection. +/// Returns `Ok` once the daemon has durably journaled the event. Callers that +/// must never fail the agent's turn should treat any `Err` as non-fatal and +/// exit 0. +pub async fn run_hook( + mut args: HookArgs, + mut route: SessionRoute, + host: HostInfo, +) -> anyhow::Result<()> { + // A managed run injects its own hook definitions. Suppress an inherited + // Braintrust plugin hook for the same child, but allow the injected hook + // process, which carries the second marker. + if std::env::var_os("_BT_TRACE_MANAGED_RUN").is_some() && !args.managed_run_hook { + return Ok(()); + } + let settings = settings::AgentSettings::load(&args.source); + if !settings.tracing_enabled() { + return Ok(()); + } + let mut payload = read_stdin_json()?; + + resolve_dynamic_hook_versions(&mut args, &payload); + + if let Some(field) = &args.transcript_path_field { + add_transcript_observation(&mut payload, field); + } + + let session_id = json_str_field(&payload, &args.session_id_field) + .ok_or_else(|| anyhow::anyhow!("no `{}` field in hook payload", args.session_id_field))?; + let event = args + .event + .clone() + .or_else(|| json_str_field(&payload, &args.event_field)) + .unwrap_or_default(); + + if let Some(configured_route) = settings.route { + route = configured_route; + } + if args.flush_on_turn_end { + route.flush_mode = wire::FlushMode::FlushOnTurnEnd; + } + apply_additional_metadata(&mut route, args.additional_metadata.as_deref())?; + let env = build_hook_envelope(&args, route, payload, session_id, event); + + let socket = paths::socket_path(args.socket.as_deref()); + forward_envelope(&env, &socket, &host, args.no_spawn).await?; + + Ok(()) +} + +fn resolve_dynamic_hook_versions(args: &mut HookArgs, payload: &serde_json::Value) { + if args.source_version.is_none() { + args.source_version = source_version_from_env(&args.source) + .or_else(|| source_version_from_payload(&args.source, payload)); + } + if args.plugin_version.is_none() { + args.plugin_version = plugin_version_from_plugin_root(&args.source); + } +} + +fn source_version_from_env(source: &str) -> Option { + let key = match source { + "claude-code" => "CLAUDE_CODE_VERSION", + "codex" => "CODEX_VERSION", + "grok" => "GROK_VERSION", + "antigravity" => "AGY_VERSION", + _ => return None, + }; + std::env::var(key) + .ok() + .filter(|value| !value.trim().is_empty()) +} + +fn source_version_from_payload(source: &str, payload: &serde_json::Value) -> Option { + let fields: &[&str] = match source { + "codex" => &["cli_version", "version"], + "claude-code" | "grok" | "antigravity" => &["version", "cli_version"], + _ => &[], + }; + fields + .iter() + .find_map(|field| json_str_field(payload, field)) +} + +fn plugin_version_from_plugin_root(source: &str) -> Option { + let root = match source { + "claude-code" => std::env::var_os("CLAUDE_PLUGIN_ROOT"), + "codex" => std::env::var_os("PLUGIN_ROOT"), + "grok" => std::env::var_os("GROK_PLUGIN_ROOT"), + // Antigravity runs plugin hooks with the plugin root as cwd. + "antigravity" => std::env::current_dir() + .ok() + .map(|path| path.into_os_string()), + _ => None, + }?; + let manifest = match source { + "claude-code" => ".claude-plugin/plugin.json", + "codex" => ".codex-plugin/plugin.json", + "grok" => ".grok-plugin/plugin.json", + "antigravity" => "plugin.json", + _ => return None, + }; + plugin_version_from_manifest(&PathBuf::from(root).join(manifest)) +} + +fn plugin_version_from_manifest(path: &std::path::Path) -> Option { + serde_json::from_slice::(&std::fs::read(path).ok()?) + .ok()? + .get("version")? + .as_str() + .map(str::to_owned) +} + +/// Apply one invocation-local JSON metadata override to a non-secret route. +/// +/// The route is then carried unchanged through live hooks, managed runs, and +/// transcript import. Keeping validation here gives every public command the +/// same contract and prevents individual agent shims from parsing JSON. +pub(crate) fn apply_additional_metadata( + route: &mut SessionRoute, + additional_metadata: Option<&str>, +) -> anyhow::Result<()> { + let Some(metadata) = additional_metadata else { + return Ok(()); + }; + let value: serde_json::Value = serde_json::from_str(metadata) + .map_err(|e| anyhow::anyhow!("invalid --additional-metadata JSON: {e}"))?; + if !value.is_object() { + anyhow::bail!("--additional-metadata must be a JSON object"); + } + route.additional_metadata = Some(value); + Ok(()) +} + +/// Apply invocation-local root-span tags to a route. Tags are normalized once +/// at the CLI boundary so hook shims and translators only receive valid values. +pub(crate) fn apply_tags(route: &mut SessionRoute, tags: &[String]) -> anyhow::Result<()> { + if tags.is_empty() { + return Ok(()); + } + let mut normalized = Vec::with_capacity(tags.len()); + for tag in tags { + let tag = tag.trim(); + if tag.is_empty() { + anyhow::bail!("--tag must not be empty"); + } + if !normalized.iter().any(|existing| existing == tag) { + normalized.push(tag.to_string()); + } + } + route.tags = normalized; + Ok(()) +} + +fn initialize_params(env: &Envelope, daemon_version: &str) -> serde_json::Value { + serde_json::json!({ + "protocol_version": PROTOCOL_VERSION, + "client": { + "source": env.source, + "daemon_version": daemon_version, + "plugin_version": env.plugin_version, + "pid": std::process::id() + } + }) +} + +/// Ensure a daemon is up and forward one already-built [`Envelope`] to it +/// (`initialize` handshake + `event.log`). Also the seam in-process clients and +/// tests use to send events without going through stdin. +pub async fn forward_envelope( + env: &Envelope, + socket: &std::path::Path, + host: &HostInfo, + no_spawn: bool, +) -> anyhow::Result<()> { + const MAX_HANDOVER_ATTEMPTS: usize = 3; + let mut handover_attempts = 0; + let mut capture_retry_deadline = None; + + loop { + if capture_retry_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) { + anyhow::bail!( + "replacement daemon did not accept the event at {}", + socket.display() + ); + } + + let stream = match client::connect(socket).await { + Ok(stream) => stream, + Err(_) if capture_retry_deadline.is_some() && no_spawn => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; + } + Err(_) => client::ensure_daemon(socket, host, no_spawn).await?, + }; + let mut conn = client::Conn::new(stream); + let initialized = match conn + .request(method::INITIALIZE, initialize_params(env, &host.version)) + .await + { + Ok(initialized) => initialized, + Err(_) if capture_retry_deadline.is_some() => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; + } + Err(error) => return Err(error), + }; + let initialized: wire::InitializeResult = serde_json::from_value(initialized)?; + if daemon_needs_upgrade(&initialized.daemon_version, &host.version) { + if no_spawn { + anyhow::bail!( + "daemon version {} is older than client {} and --no-spawn is set", + initialized.daemon_version, + host.version + ); + } + if handover_attempts == MAX_HANDOVER_ATTEMPTS { + anyhow::bail!( + "daemon version {} is still older than client {} after {MAX_HANDOVER_ATTEMPTS} handover attempts", + initialized.daemon_version, + host.version + ); + } + conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) + .await?; + handover_attempts += 1; + drop(conn); + wait_for_daemon_exit(socket).await; + continue; + } + let result = match conn.request(method::EVENT_LOG, env).await { + Ok(result) => result, + Err(_) if capture_retry_deadline.is_some() => { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + continue; + } + Err(error) => return Err(error), + }; + let result: wire::EventLogResult = serde_json::from_value(result)?; + if result.accepted { + return Ok(()); + } + capture_retry_deadline + .get_or_insert_with(|| tokio::time::Instant::now() + std::time::Duration::from_secs(5)); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } +} + +async fn wait_for_daemon_exit(socket: &std::path::Path) { + for _ in 0..100 { + if client::connect(socket).await.is_err() { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } +} + +/// A shared daemon only moves forward. Released versions use semver; opaque +/// development versions retain the old exact-match handover behavior. +fn daemon_needs_upgrade(daemon_version: &str, client_version: &str) -> bool { + if daemon_version == client_version { + return false; + } + compare_daemon_versions(daemon_version, client_version).is_none_or(|ordering| ordering.is_lt()) +} + +pub(crate) fn compare_daemon_versions(left: &str, right: &str) -> Option { + Some( + semver::Version::parse(left) + .ok()? + .cmp(&semver::Version::parse(right).ok()?), + ) +} + +/// Ask the daemon to flush a session, bounded by `timeout_ms`. A reliable +/// barrier: on return, every event enqueued before the call has been processed +/// and its spans emitted to the sink. +pub async fn flush_session( + session_id: &str, + socket: &std::path::Path, + timeout_ms: u64, +) -> anyhow::Result { + let stream = client::connect(socket).await?; + let mut conn = client::Conn::new(stream); + conn.request( + method::INITIALIZE, + serde_json::json!({ "protocol_version": PROTOCOL_VERSION, "client": { "source": "flush" } }), + ) + .await?; + let params = wire::FlushParams { + session_id: session_id.to_string(), + timeout_ms, + }; + let value = conn.request(method::SESSION_FLUSH, params).await?; + Ok(serde_json::from_value(value)?) +} + +/// Flush every daemon session accepted from one managed child process tree. +/// A missing daemon is not itself a failure: the daemon may have idle-exited +/// after draining every session, so acceptance is checked against the +/// persisted managed-run record instead. +pub async fn flush_managed_run( + managed_run_id: &str, + socket: &std::path::Path, + timeout_ms: u64, +) -> anyhow::Result { + flush_managed_run_in(managed_run_id, socket, timeout_ms, &paths::data_dir(None)).await +} + +async fn flush_managed_run_in( + managed_run_id: &str, + socket: &std::path::Path, + timeout_ms: u64, + data_dir: &std::path::Path, +) -> anyhow::Result { + let stream = match client::connect(socket).await { + Ok(stream) => stream, + Err(_) => { + let accepted_sessions = journal::read_managed_run_keys(data_dir, managed_run_id) + .await + .len() as u64; + return Ok(wire::FlushResult { + flushed: true, + pending: 0, + accepted_sessions, + }); + } + }; + let mut conn = client::Conn::new(stream); + conn.request( + method::INITIALIZE, + serde_json::json!({ + "protocol_version": PROTOCOL_VERSION, + "client": { "source": "managed-run-flush" } + }), + ) + .await?; + let params = ManagedRunFlushParams { + managed_run_id: managed_run_id.to_string(), + timeout_ms, + }; + let value = conn.request(method::MANAGED_RUN_FLUSH, params).await?; + Ok(serde_json::from_value(value)?) +} + +/// Query daemon status. `Ok(None)` means no daemon is running. +pub async fn run_status(args: StatusArgs) -> anyhow::Result> { + let socket = paths::socket_path(args.socket.as_deref()); + let stream = match client::connect(&socket).await { + Ok(s) => s, + Err(_) => return Ok(None), + }; + let mut conn = client::Conn::new(stream); + conn.request( + method::INITIALIZE, + serde_json::json!({ + "protocol_version": PROTOCOL_VERSION, + "client": { "source": "status" } + }), + ) + .await?; + let params = wire::StatusParams { + session_id: args.session_id.clone(), + }; + let value = conn.request(method::STATUS_GET, params).await?; + Ok(Some(serde_json::from_value(value)?)) +} + +/// Request a graceful daemon shutdown. Primarily useful for lifecycle +/// management and transport integration tests. +pub async fn shutdown_daemon(socket: &std::path::Path) -> anyhow::Result<()> { + let stream = client::connect(socket).await?; + let mut conn = client::Conn::new(stream); + conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) + .await?; + Ok(()) +} + +/// Import a native coding-agent transcript through the normal translators and +/// sink. This is separate from daemon journal recovery: import creates traces +/// for a past session, while recovery rebuilds live correlation state. +pub async fn run_import( + args: ImportArgs, + opts: ServeOptions, + mut config: Option, +) -> anyhow::Result> { + validate_import_selection(&args)?; + apply_import_span_plugins(&mut config, &args.plugin)?; + let destination = import_parent_components(&args)? + .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" + ) + } + if args.all { + anyhow::bail!("Muse import does not support --all yet; pass one or more session ids") + } + let exports = export_muse_sessions(&args.session_ids)?; + 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 { + return import_transcript_with_ledger( + &files[0], + args.source, + opts, + config, + true, + Some(ledger_dir), + ) + .await; + } + import_transcripts_with_ledger(&files, args.source, opts, config, Some(ledger_dir)).await +} + +fn apply_import_span_plugins( + config: &mut Option, + plugins: &[PathBuf], +) -> anyhow::Result<()> { + let plugins = resolve_span_plugin_paths(plugins)?; + if let Some(config) = config.as_mut() { + config.span_plugins = plugins; + } else if !plugins.is_empty() { + *config = Some(SessionConfig { + auth: wire::BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + destination: None, + flush_mode: wire::FlushMode::FireAndForget, + additional_metadata: None, + tags: Vec::new(), + span_plugins: plugins, + }); + } + Ok(()) +} + +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 = std::process::Command::new(&executable) + .args([ + OsString::from("export"), + OsString::from("--session"), + OsString::from(session_id), + OsString::from("--out"), + ]) + .arg(directory.path().join(format!("{index}.json"))) + .output() + .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) +} + +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"); + } + if args.attach && args.session_ids.len() != 1 { + anyhow::bail!("--attach requires exactly one session id"); + } + import_parent_components(args)?; + Ok(()) +} + +fn import_parent_components(args: &ImportArgs) -> anyhow::Result> { + if let Some(parent) = &args.parent { + parent + .to_parent_span_info() + .map_err(|error| anyhow::anyhow!("invalid --parent value: {error}"))?; + return Ok(Some(parent.clone())); + } + + let Some(span_id) = args.parent_span_id.as_deref() else { + return Ok(None); + }; + let span_id = span_id.trim(); + if span_id.is_empty() { + anyhow::bail!("--parent-span-id must not be empty"); + } + let root_span_id = args + .parent_root_span_id + .as_deref() + .ok_or_else(|| anyhow::anyhow!("--parent-span-id requires --parent-root-span-id"))?; + let root_span_id = root_span_id.trim(); + if root_span_id.is_empty() { + anyhow::bail!("--parent-root-span-id must not be empty"); + } + let object_type = args + .parent_object_type + .ok_or_else(|| anyhow::anyhow!("--parent-span-id requires --parent-object-type"))?; + + let object_id = args + .parent_object_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + let parent_project = args + .parent_project + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let compute_object_metadata_args = match object_type { + ParentObjectType::ProjectLogs if object_id.is_none() => { + let project = parent_project.ok_or_else(|| { + anyhow::anyhow!( + "project_logs parent requires --parent-object-id or --parent-project" + ) + })?; + Some(serde_json::Map::from_iter([( + "project_name".to_string(), + serde_json::Value::String(project.to_string()), + )])) + } + ParentObjectType::ProjectLogs => None, + _ if parent_project.is_some() => { + anyhow::bail!("--parent-project is only valid with --parent-object-type project_logs") + } + _ if object_id.is_none() => { + anyhow::bail!("this parent object type requires --parent-object-id") + } + _ => None, + }; + + let components = SpanComponents { + object_type: object_type.into(), + object_id, + compute_object_metadata_args, + row_id: None, + span_id: Some(span_id.to_string()), + root_span_id: Some(root_span_id.to_string()), + span_parents: None, + propagated_event: None, + }; + components + .to_parent_span_info() + .map_err(|error| anyhow::anyhow!("invalid parent identifiers: {error}"))?; + Ok(Some(components)) +} + +/// Launch a coding agent with inherited stdio and inject Braintrust hooks for +/// this invocation, without requiring the tracing plugin to be installed or +/// enabled globally. +pub async fn run_traced( + args: RunArgs, + hook_command: RunHookCommand, + mut route: SessionRoute, +) -> anyhow::Result { + apply_run_span_plugins(&mut route, &args.plugin)?; + if route.destination.is_none() { + anyhow::bail!( + "managed run requires a trace destination; select a project, object destination, or parent span" + ); + } + if args.source == RunSource::Codex + && args.agent_args.iter().any(|arg| { + let arg = arg.to_string_lossy(); + arg == "--dangerously-bypass-hook-trust" + || arg.starts_with("--dangerously-bypass-hook-trust=") + }) + { + anyhow::bail!( + "bt trace run codex cannot be combined with --dangerously-bypass-hook-trust; managed tracing preserves Codex hook trust. Remove the flag, or run Codex directly" + ); + } + let (executable_env, default_executable) = match args.source { + RunSource::Codex => ("CODEX_BIN", "codex"), + RunSource::Claude => ("CLAUDE_BIN", "claude"), + RunSource::OpenCode => ("OPENCODE_BIN", "opencode"), + RunSource::Pi => ("PI_BIN", "pi"), + }; + let executable = + std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); + let injected_args = managed_run_args(args.source, &hook_command)?; + let managed_run_id = uuid::Uuid::new_v4().to_string(); + // A caller-supplied socket is already an explicit daemon boundary (for + // example an integration harness or a deliberately isolated runtime). + // Only create our own boundary when environment auth would otherwise use + // the ambient shared daemon. + let isolate_daemon = route.auth.effective_source() == wire::AuthSource::Environment + && std::env::var_os(paths::SOCKET_ENV).is_none(); + let isolated_runtime = isolate_daemon + .then(|| ManagedRunRuntime::new(&managed_run_id)) + .transpose()?; + let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?; + let mut command = tokio::process::Command::new(&executable); + command + .args(injected_args) + .args(args.agent_args) + .env("_BT_TRACE_MANAGED_RUN", "1") + .env(MANAGED_RUN_ID_ENV, &managed_run_id) + .env(settings::INVOCATION_SETTINGS_ENV, invocation_settings); + if let Some(runtime) = &isolated_runtime { + command + .env(paths::SOCKET_ENV, &runtime.socket) + .env(paths::DATA_DIR_ENV, runtime.temp_dir.path()); + } + if args.source == RunSource::OpenCode { + command.env( + "OPENCODE_CONFIG_CONTENT", + opencode_managed_config(std::env::var("OPENCODE_CONFIG_CONTENT").ok().as_deref())?, + ); + } + let mut child = command.spawn().map_err(|error| { + anyhow::anyhow!("failed to launch {}: {error}", executable.to_string_lossy()) + })?; + let interrupt = tokio::signal::ctrl_c(); + tokio::pin!(interrupt); + + let status = tokio::select! { + status = child.wait() => status.map_err(anyhow::Error::from), + result = &mut interrupt => { + match result { + Ok(()) => { + let kill_result = child.start_kill(); + let wait_result = child.wait().await; + kill_result + .map_err(anyhow::Error::from) + .and_then(|()| wait_result.map_err(anyhow::Error::from)) + } + Err(error) => Err(error.into()), + } + } + }; + let socket = isolated_runtime + .as_ref() + .map(|runtime| runtime.socket.clone()) + .unwrap_or_else(|| paths::socket_path(None)); + let data_dir = isolated_runtime + .as_ref() + .map(|runtime| runtime.temp_dir.path().to_path_buf()) + .unwrap_or_else(|| paths::data_dir(None)); + match flush_managed_run_in( + &managed_run_id, + &socket, + MANAGED_RUN_FLUSH_TIMEOUT_MS, + &data_dir, + ) + .await + { + Ok(result) if result.accepted_sessions == 0 => tracing::warn!( + managed_run_id, + "managed run produced no accepted trace events" + ), + Ok(result) if result.flushed => {} + Ok(result) => tracing::warn!( + managed_run_id, + pending = result.pending, + "managed run trace flush timed out" + ), + Err(error) => tracing::warn!(managed_run_id, %error, "managed run trace flush failed"), + } + if let Some(runtime) = &isolated_runtime { + let _ = shutdown_daemon(&socket).await; + if let Err(error) = + crate::plugin_diagnostics::merge(runtime.temp_dir.path(), &paths::data_dir(None)) + { + tracing::warn!(managed_run_id, %error, "failed to preserve managed-run plugin diagnostics"); + } + } + status +} + +struct ManagedRunRuntime { + temp_dir: tempfile::TempDir, + socket: std::path::PathBuf, +} + +impl ManagedRunRuntime { + fn new(_managed_run_id: &str) -> anyhow::Result { + let temp_dir = tempfile::Builder::new().prefix("bt-trace-run-").tempdir()?; + #[cfg(unix)] + let socket = temp_dir.path().join("daemon.sock"); + #[cfg(windows)] + let socket = std::path::PathBuf::from(format!( + r"\\.\pipe\braintrust-bt-daemon-managed-{_managed_run_id}" + )); + Ok(Self { temp_dir, socket }) + } +} + +fn apply_run_span_plugins(route: &mut SessionRoute, plugins: &[PathBuf]) -> anyhow::Result<()> { + route.span_plugins = resolve_span_plugin_paths(plugins)?; + Ok(()) +} + +pub(crate) fn resolve_span_plugin_paths(paths: &[PathBuf]) -> anyhow::Result> { + let paths: Vec<_> = paths + .iter() + .map(|path| { + path.canonicalize().map_err(|error| { + anyhow::anyhow!("could not resolve span plugin {}: {error}", path.display()) + }) + }) + .collect::>()?; + crate::span_processor::validate(&paths)?; + Ok(paths) +} + +fn managed_run_args( + source: RunSource, + hook_command: &RunHookCommand, +) -> anyhow::Result> { + let source_name = match source { + RunSource::Codex => "codex", + RunSource::Claude => "claude", + RunSource::OpenCode => "opencode", + RunSource::Pi => "pi", + }; + match source { + RunSource::Codex | RunSource::Claude => { + let unix_command = managed_hook_shell_command(hook_command, source_name, false)?; + let windows_command = managed_hook_shell_command(hook_command, source_name, true)?; + match source { + RunSource::Codex => Ok(codex_managed_run_args(&unix_command, &windows_command)), + RunSource::Claude => Ok(claude_managed_run_args(if cfg!(windows) { + &windows_command + } else { + &unix_command + })?), + _ => unreachable!(), + } + } + RunSource::OpenCode => Ok(Vec::new()), + RunSource::Pi => { + let extension = match std::env::var_os("BT_TRACE_PI_PLUGIN_SPEC") { + Some(extension) => extension, + None => OsString::from(crate::setup::pi_plugin_spec()), + }; + Ok(vec![OsString::from("-e"), extension]) + } + } +} + +fn opencode_managed_config(existing: Option<&str>) -> anyhow::Result { + let mut config = match existing { + Some(raw) => serde_json::from_str::(raw) + .map_err(|error| anyhow::anyhow!("invalid OPENCODE_CONFIG_CONTENT: {error}"))?, + None => serde_json::json!({}), + }; + let object = config + .as_object_mut() + .ok_or_else(|| anyhow::anyhow!("OPENCODE_CONFIG_CONTENT must be a JSON object"))?; + let plugins = object + .entry("plugin") + .or_insert_with(|| serde_json::Value::Array(Vec::new())) + .as_array_mut() + .ok_or_else(|| anyhow::anyhow!("OPENCODE_CONFIG_CONTENT.plugin must be an array"))?; + let plugin = std::env::var("BT_TRACE_OPENCODE_PLUGIN_SPEC") + .unwrap_or_else(|_| "@braintrust/trace-opencode/tracing".to_string()); + if !plugins.iter().any(|value| value.as_str() == Some(&plugin)) { + plugins.push(serde_json::Value::String(plugin)); + } + Ok(serde_json::to_string(&config)?) +} + +fn managed_hook_shell_command( + hook_command: &RunHookCommand, + source: &str, + windows: bool, +) -> anyhow::Result { + let mut argv = Vec::with_capacity(hook_command.args.len() + 4); + argv.push(hook_command.program.clone()); + argv.extend(hook_command.args.iter().cloned()); + argv.push(OsString::from("--source")); + argv.push(OsString::from(source)); + argv.push(OsString::from("--managed-run-hook")); + let mut rendered = Vec::with_capacity(argv.len()); + for arg in argv { + let arg = arg + .into_string() + .map_err(|_| anyhow::anyhow!("managed hook command contains non-Unicode argv"))?; + rendered.push(if windows { + quote_windows_command_arg(&arg) + } else { + quote_unix_shell_arg(&arg) + }); + } + Ok(rendered.join(" ")) +} + +fn quote_unix_shell_arg(arg: &str) -> String { + format!("'{}'", arg.replace('\'', "'\"'\"'")) +} + +fn quote_windows_command_arg(arg: &str) -> String { + format!("\"{}\"", arg.replace('\\', "/").replace('"', "\"\"")) +} + +const CODEX_RUN_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreCompact", + "PostCompact", + "SubagentStart", + "SubagentStop", + "Stop", + "SessionEnd", +]; + +const CLAUDE_RUN_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "Setup", + "UserPromptSubmit", + "UserPromptExpansion", + "PreToolUse", + "PermissionRequest", + "PermissionDenied", + "PostToolUse", + "PostToolUseFailure", + "PostToolBatch", + "PreCompact", + "PostCompact", + "Notification", + "MessageDisplay", + "SubagentStart", + "SubagentStop", + "TaskCreated", + "TaskCompleted", + "Stop", + "StopFailure", + "SessionEnd", +]; + +fn codex_managed_run_args(unix_command: &str, windows_command: &str) -> Vec { + let unix_command = serde_json::to_string(unix_command).expect("serialize hook command"); + let windows_command = + serde_json::to_string(windows_command).expect("serialize Windows hook command"); + let mut args = vec![OsString::from("--enable"), OsString::from("hooks")]; + for event in CODEX_RUN_HOOK_EVENTS { + args.push(OsString::from("-c")); + args.push(OsString::from(format!( + "hooks.{event}=[{{hooks=[{{type=\"command\",command={unix_command},commandWindows={windows_command}}}]}}]" + ))); + } + args +} + +fn claude_managed_run_args(command: &str) -> anyhow::Result> { + let matcher_group = serde_json::json!([{ + "hooks": [{ + "type": "command", + "command": command, + "async": false + }] + }]); + let hooks = CLAUDE_RUN_HOOK_EVENTS + .iter() + .map(|event| ((*event).to_string(), matcher_group.clone())) + .collect::>(); + Ok(vec![ + OsString::from("--settings"), + OsString::from(serde_json::to_string( + &serde_json::json!({ "hooks": hooks }), + )?), + ]) +} + +fn apply_import_destination( + config: &mut Option, + destination: Option, +) -> anyhow::Result<()> { + if let Some(destination) = destination { + let config = config.as_mut().ok_or_else(|| { + anyhow::anyhow!( + "import destination requires a resolved Braintrust session configuration; \ + use `bt trace import` instead of the standalone `bt-daemon import` command" + ) + })?; + config.destination = Some(destination); + } + Ok(()) +} + +/// Import a native transcript from a known path. Front-ends should normally +/// expose [`run_import`] so users only need the agent's session id; this lower- +/// level entry point is useful for embedding and isolated tests. +pub async fn import_transcript( + file: &std::path::Path, + source: ImportSource, + opts: ServeOptions, + config: Option, + attach: bool, +) -> anyhow::Result> { + import_transcript_with_ledger(file, source, opts, config, attach, None).await +} + +async fn import_transcript_with_ledger( + file: &std::path::Path, + source: ImportSource, + opts: ServeOptions, + config: Option, + attach: bool, + ledger_dir: Option, +) -> anyhow::Result> { + let mut tail = transcript_import::TranscriptTail::new(file.to_path_buf(), source); + let mut processor = ImportProcessor::new(opts, config, ledger_dir); + let shutdown = tokio::signal::ctrl_c(); + tokio::pin!(shutdown); + let mut finalizing = !attach; + loop { + let entries = tail.poll(finalizing)?; + if tail.take_translator_reset() { + processor.reset_translators(); + } + processor.process(entries).await?; + if finalizing { + break; + } + tokio::select! { + result = &mut shutdown => { + result?; + finalizing = true; + } + _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} + } + } + 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> { + import_transcripts_with_ledger(files, source, opts, config, None).await +} + +async fn import_transcripts_with_ledger( + files: &[PathBuf], + source: ImportSource, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + let mut processor = ImportProcessor::new(opts, config, ledger_dir); + let mut summaries = Vec::new(); + 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()))?; + 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) +} + +struct ImportLive { + source: String, + translator: Box, + sink: Box, + ctx: SessionCtx, + pending_ops: usize, + span_ids: std::collections::HashSet, + root_span_id: Option, + destination: Option, + diagnostics_dir: Option, +} + +struct ImportProcessor { + sessions: std::collections::HashMap, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +} + +impl ImportProcessor { + fn new(opts: ServeOptions, config: Option, ledger_dir: Option) -> Self { + Self { + sessions: std::collections::HashMap::new(), + opts, + config, + ledger_dir, + } + } + + async fn process(&mut self, entries: Vec) -> anyhow::Result<()> { + for mut env in entries { + env.config = self.config.clone(); + let sid = env.session_id.clone(); + let live = match self.sessions.get_mut(&sid) { + Some(live) => live, + None => { + let translator = self.opts.translators.create(&env.source, &sid); + let sink = self.opts.sink_factory.create( + &sid, + &env.source, + env.plugin_version.as_deref(), + )?; + let sink: Box = match &self.ledger_dir { + Some(ledger_dir) => Box::new( + delivery_ledger::LedgerSink::new( + sink, + ledger_dir, + &env.source, + &sid, + env.config.as_ref(), + ) + .await, + ), + None => sink, + }; + self.sessions.insert( + sid.clone(), + ImportLive { + source: env.source.clone(), + translator, + sink, + ctx: SessionCtx { + session_id: sid.clone(), + config: None, + }, + pending_ops: 0, + span_ids: std::collections::HashSet::new(), + root_span_id: self + .config + .as_ref() + .and_then(|config| config.attached_span_ids().1), + destination: self + .config + .as_ref() + .and_then(|config| config.destination.clone()), + diagnostics_dir: self.ledger_dir.clone(), + }, + ); + self.sessions.get_mut(&sid).unwrap() + } + }; + if let Some(cfg) = &env.config { + live.sink.configure(cfg); + live.ctx.config = Some(cfg.clone()); + } + let ops = live.translator.handle(&env, &live.ctx)?; + Self::emit_translator_batches(live, ops).await?; + } + Ok(()) + } + + fn reset_translators(&mut self) { + for (session_id, live) in &mut self.sessions { + live.translator = self.opts.translators.create(&live.source, session_id); + } + } + + async fn emit_translator_batches( + live: &mut ImportLive, + first: Vec, + ) -> anyhow::Result<()> { + let mut next = Some(first); + while let Some(ops) = next { + // Imports can contain tens of thousands of SDK log commands. Bound the + // number queued between drains without serializing one network flush + // for every native turn boundary. + const FLUSH_OPS: usize = 500; + for chunk in ops.chunks(FLUSH_OPS) { + for op in chunk { + let row = match op { + SpanOp::Insert(row) | SpanOp::Merge(row) => row, + }; + live.span_ids.insert(row.span_id.clone()); + if live.root_span_id.is_none() && !row.root_span_id.is_empty() { + live.root_span_id = Some(row.root_span_id.clone()); + } + } + let plugins = live + .ctx + .config + .as_ref() + .map(|config| config.span_plugins.as_slice()) + .unwrap_or_default(); + let mut transformed = Vec::with_capacity(chunk.len()); + for op in chunk { + match crate::span_processor::process( + plugins, + op, + &live.source, + &live.ctx.session_id, + ) { + Ok(result) => { + if let Some(failure) = result.failure { + if failure.newly_seen { + if let Some(data_dir) = &live.diagnostics_dir { + if let Err(error) = crate::plugin_diagnostics::record( + data_dir, + &live.source, + &failure.path, + &failure.message, + ) { + tracing::warn!( + session_id = %live.ctx.session_id, + "failed to persist span plugin diagnostic: {error}" + ); + } + } + tracing::warn!( + session_id = %live.ctx.session_id, + plugin = %failure.path.display(), + error = %failure.message, + "span plugin failed during import; span operations are being discarded" + ); + } + } + if let Some(op) = result.op { + transformed.push(op); + } + } + Err(error) => { + if let (Some(data_dir), Some(plugin)) = + (&live.diagnostics_dir, plugins.first()) + { + if let Err(diagnostic_error) = crate::plugin_diagnostics::record( + data_dir, + &live.source, + plugin, + &error.to_string(), + ) { + tracing::warn!( + session_id = %live.ctx.session_id, + "failed to persist span plugin diagnostic: {diagnostic_error}" + ); + } + } + tracing::warn!( + session_id = %live.ctx.session_id, + %error, + "span plugin processor failed during import; span operation discarded" + ); + } + } + } + if !transformed.is_empty() { + live.sink.emit(&transformed).await?; + } + live.pending_ops += transformed.len(); + if live.pending_ops >= FLUSH_OPS { + live.sink.flush().await?; + live.pending_ops = 0; + } + } + next = live.translator.drain_pending(&live.ctx)?; + } + Ok(()) + } + + async fn finish(self) -> anyhow::Result> { + let mut summaries = Vec::new(); + for (sid, mut live) in self.sessions { + let ops = live.translator.flush(&live.ctx)?; + Self::emit_translator_batches(&mut live, ops).await?; + live.sink.flush().await?; + summaries.push(Self::summary(sid, live)); + } + summaries.sort_by(|left, right| left.session_id.cmp(&right.session_id)); + Ok(summaries) + } + + async fn finish_session(&mut self, session_id: &str) -> anyhow::Result> { + let Some(mut live) = self.sessions.remove(session_id) else { + return Ok(None); + }; + let ops = live.translator.flush(&live.ctx)?; + Self::emit_translator_batches(&mut live, ops).await?; + live.sink.flush().await?; + Ok(Some(Self::summary(session_id.to_string(), live))) + } + + fn summary(session_id: String, live: ImportLive) -> ImportSummary { + ImportSummary { + session_id, + destination: live.destination, + root_span_id: live.root_span_id, + span_count: live.span_ids.len(), + finalized: true, + } + } +} + +/// Build debug [`ServeOptions`]: the production translator registry plus a +/// debug sink writing NDJSON under `/spans/`. +pub fn debug_serve_options(version: impl Into, data_dir: &std::path::Path) -> ServeOptions { + ServeOptions { + version: version.into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(DebugSinkFactory { + dir: data_dir.join("spans"), + }), + auth_provider: None, + } +} + +/// Build [`ServeOptions`] with the Braintrust sink. `translators` lets the +/// caller choose the translator registry (debug-only until Phase 3 adds the +/// Codex/Claude translators). Clients are built lazily per session URL, so this +/// is cheap and infallible. +pub fn braintrust_serve_options( + version: impl Into, + sink_config: BraintrustSinkConfig, + translators: Arc, +) -> ServeOptions { + ServeOptions { + version: version.into(), + translators, + sink_factory: Arc::new(BraintrustSinkFactory::new(sink_config)), + auth_provider: None, + } +} + +fn now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) +} + +fn read_stdin_json() -> anyhow::Result { + use std::io::Read; + let mut buf = String::new(); + std::io::stdin().read_to_string(&mut buf)?; + if buf.trim().is_empty() { + anyhow::bail!("empty stdin (expected a JSON hook payload)"); + } + Ok(serde_json::from_str(&buf)?) +} + +/// Read a string-ish field (`session_id` / event name) from the payload, +/// coercing a JSON number to its string form. +fn json_str_field(payload: &serde_json::Value, field: &str) -> Option { + match payload.get(field) { + Some(serde_json::Value::String(s)) => Some(s.clone()), + Some(serde_json::Value::Number(n)) => Some(n.to_string()), + _ => None, + } +} + +/// Stamp the transcript boundary visible when a blocking hook runs. Agent +/// transcripts are append-only, while daemon journal replay may happen after +/// the session has advanced. Recording byte lengths keeps translation causally +/// aligned with each native hook without copying transcript contents into the +/// journal. +fn add_transcript_observation(payload: &mut serde_json::Value, field: &str) { + let Some(path) = json_str_field(payload, field) else { + return; + }; + let transcript = std::path::Path::new(&path); + let mut observation = serde_json::Map::new(); + observation.insert("path".into(), serde_json::Value::String(path.clone())); + if let Ok(metadata) = std::fs::metadata(transcript) { + observation.insert( + "observed_bytes".into(), + serde_json::Value::Number(metadata.len().into()), + ); + } + + if transcript.file_name().and_then(|name| name.to_str()) == Some("transcript.jsonl") { + let full = transcript.with_file_name("transcript_full.jsonl"); + if let Ok(metadata) = std::fs::metadata(&full) { + observation.insert( + "full_path".into(), + serde_json::Value::String(full.to_string_lossy().into_owned()), + ); + observation.insert( + "full_observed_bytes".into(), + serde_json::Value::Number(metadata.len().into()), + ); + } + } + + if let Some(object) = payload.as_object_mut() { + object.insert( + "_bt_transcript_observation".into(), + serde_json::Value::Object(observation), + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use clap::Parser; + use serde_json::json; + + #[derive(Debug, Parser)] + struct ImportCli { + #[command(flatten)] + args: ImportArgs, + } + + #[derive(Debug, Parser)] + struct ServeCli { + #[command(flatten)] + args: ServeArgs, + } + + #[derive(Debug, Parser)] + struct HookCli { + #[command(flatten)] + args: HookArgs, + } + + #[test] + fn serve_defaults_to_short_journal_backed_session_retirement() { + let args = ServeCli::try_parse_from(["test"]).unwrap().args; + assert_eq!(args.session_idle_timeout_secs, 30); + } + + #[test] + fn hook_plugin_version_reaches_envelope_and_every_initialize_attempt() { + let args = HookCli::try_parse_from([ + "test", + "--source", + "grok", + "--source-version", + "1.0.13", + "--plugin-version", + "0.1.0", + ]) + .unwrap() + .args; + assert_eq!(args.source_version.as_deref(), Some("1.0.13")); + assert_eq!(args.plugin_version.as_deref(), Some("0.1.0")); + + let payload = serde_json::json!({"native": "unchanged"}); + let env = build_hook_envelope( + &args, + SessionRoute::default(), + payload.clone(), + "session-1".into(), + "SessionEnd".into(), + ); + assert_eq!(env.source_version.as_deref(), Some("1.0.13")); + assert_eq!(env.plugin_version.as_deref(), Some("0.1.0")); + assert_eq!(env.payload, payload); + + // Both the initial connection and the post-restart retry use this + // shared parameter builder. + let initialize = initialize_params(&env, "1.0.13"); + assert_eq!(initialize["client"]["source"], "grok"); + assert_eq!(initialize["client"]["daemon_version"], "1.0.13"); + assert_eq!(initialize["client"]["plugin_version"], "0.1.0"); + assert_ne!(initialize["client"]["plugin_version"], "1.0.13"); + } + + #[test] + fn hook_versions_are_discovered_without_hook_configuration_literals() { + let payload = serde_json::json!({"cli_version": "2.3.4"}); + assert_eq!( + source_version_from_payload("codex", &payload).as_deref(), + Some("2.3.4") + ); + + let temp = tempfile::tempdir().unwrap(); + let manifest = temp.path().join("plugin.json"); + std::fs::write(&manifest, br#"{"version":"5.6.7"}"#).unwrap(); + assert_eq!( + plugin_version_from_manifest(&manifest).as_deref(), + Some("5.6.7") + ); + } + + #[test] + fn hook_flush_recognizes_native_and_documented_terminal_events() { + for event in ["session_end", "SessionEnd"] { + assert!(should_flush_hook_event(event, false)); + } + for event in [ + "stop", + "Stop", + "stop_failure", + "StopFailure", + "stop_cancelled", + "StopCancelled", + "subagent_stop", + "SubagentStop", + ] { + assert!(!should_flush_hook_event(event, false)); + assert!(should_flush_hook_event(event, true)); + } + assert!(!should_flush_hook_event("turn_completed", true)); + } + + #[test] + fn pi_agent_end_flushes_only_completed_turns_when_enabled() { + let mut env: wire::Envelope = serde_json::from_value(serde_json::json!({ + "source": "pi", "session_id": "pi-turn", "event": "agent_end", + "ts_ms": 1, "payload": {"event": {}}, + "route": {"flush_mode": "flush_on_turn_end"} + })) + .unwrap(); + for native in [ + serde_json::json!({}), + serde_json::json!({"willRetry": false}), + ] { + env.payload = serde_json::json!({"event": native}); + assert!(should_flush_ingress_event(&env)); + } + for payload in [ + serde_json::json!({"event": {"willRetry": true}}), + serde_json::json!({"willRetry": true}), + ] { + env.payload = payload; + assert!(!should_flush_ingress_event(&env)); + } + env.payload = serde_json::json!({"event": {}}); + env.source = "opencode".into(); + assert!(!should_flush_ingress_event(&env)); + env.source = "pi".into(); + env.event = "message_end".into(); + assert!(!should_flush_ingress_event(&env)); + env.event = "agent_end".into(); + env.route.as_mut().unwrap().flush_mode = wire::FlushMode::FireAndForget; + assert!(!should_flush_ingress_event(&env)); + env.route = None; + assert!(!should_flush_ingress_event(&env)); + for event in ["session_shutdown", "session_compact", "session_tree"] { + env.event = event.into(); + assert!(should_flush_ingress_event(&env)); + env.source = "opencode".into(); + assert!(!should_flush_ingress_event(&env)); + env.source = "pi".into(); + } + } + + #[test] + fn opencode_lifecycle_events_flush_from_ingress() { + let mut env: wire::Envelope = serde_json::from_value(serde_json::json!({ + "source": "opencode", "session_id": "opencode-session", "event": "session.idle", + "ts_ms": 1, "payload": {}, + "route": {"flush_mode": "fire_and_forget"} + })) + .unwrap(); + for event in [ + "session.idle", + "session.deleted", + "session.error", + "server.instance.disposed", + ] { + env.event = event.into(); + assert!(should_flush_ingress_event(&env)); + } + env.event = "session.compacted".into(); + assert!(!should_flush_ingress_event(&env)); + env.event = "session.idle".into(); + env.source = "pi".into(); + assert!(!should_flush_ingress_event(&env)); + } + + #[test] + fn span_plugin_paths_are_canonicalized_before_use() { + let dir = tempfile::Builder::new() + .prefix("span-plugin-path-") + .tempdir_in(".") + .unwrap(); + let plugin = dir.path().join("plugin.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let relative = PathBuf::from(dir.path().file_name().unwrap()).join("plugin.mjs"); + + let resolved = resolve_span_plugin_paths(&[relative]).unwrap(); + + assert_eq!(resolved, [plugin.canonicalize().unwrap()]); + assert!(resolved[0].is_absolute()); + } + + #[test] + fn additional_metadata_overrides_a_route_only_with_a_json_object() { + let mut route = SessionRoute { + additional_metadata: Some(serde_json::json!({"saved": true})), + ..SessionRoute::default() + }; + apply_additional_metadata(&mut route, Some(r#"{"run_id":"123"}"#)).unwrap(); + assert_eq!( + route.additional_metadata, + Some(serde_json::json!({"run_id": "123"})) + ); + + let error = apply_additional_metadata(&mut route, Some("[]")).unwrap_err(); + assert!(error.to_string().contains("must be a JSON object")); + let error = apply_additional_metadata(&mut route, Some("not-json")).unwrap_err(); + assert!(error + .to_string() + .contains("invalid --additional-metadata JSON")); + } + + #[test] + fn tags_override_a_route_and_are_normalized() { + let mut route = SessionRoute { + tags: vec!["saved".into()], + ..SessionRoute::default() + }; + apply_tags(&mut route, &[" ci ".into(), "ci".into(), "docs".into()]).unwrap(); + assert_eq!(route.tags, ["ci", "docs"]); + + let error = apply_tags(&mut route, &[" ".into()]).unwrap_err(); + assert!(error.to_string().contains("must not be empty")); + } + + #[test] + fn managed_run_plugins_replace_inherited_plugins() { + let temp = tempfile::tempdir().unwrap(); + let plugin = temp.path().join("run.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let mut route = SessionRoute { + span_plugins: vec![PathBuf::from("persisted.mjs")], + ..SessionRoute::default() + }; + + apply_run_span_plugins(&mut route, &[]).unwrap(); + assert!(route.span_plugins.is_empty()); + + apply_run_span_plugins(&mut route, std::slice::from_ref(&plugin)).unwrap(); + assert_eq!(route.span_plugins, [plugin.canonicalize().unwrap()]); + } + + #[test] + fn import_plugins_replace_inherited_plugins() { + let temp = tempfile::tempdir().unwrap(); + let plugin = temp.path().join("import.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let mut config = Some(SessionConfig { + auth: wire::BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + destination: None, + flush_mode: wire::FlushMode::FireAndForget, + additional_metadata: None, + tags: Vec::new(), + span_plugins: vec![PathBuf::from("persisted.mjs")], + }); + + apply_import_span_plugins(&mut config, &[]).unwrap(); + assert!(config.as_ref().unwrap().span_plugins.is_empty()); + + apply_import_span_plugins(&mut config, std::slice::from_ref(&plugin)).unwrap(); + assert_eq!( + config.unwrap().span_plugins, + [plugin.canonicalize().unwrap()] + ); + } + + #[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); + + 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()); + } + + #[test] + fn import_parent_accepts_complete_cli_identifiers_without_guessing() { + let args = ImportCli::try_parse_from([ + "test", + "codex", + "session-a", + "--parent-span-id", + "span-1", + "--parent-root-span-id", + "root-1", + "--parent-object-type", + "project_logs", + "--parent-project", + "Agents", + ]) + .unwrap() + .args; + + let components = import_parent_components(&args).unwrap().unwrap(); + assert_eq!(components.object_type, SpanObjectType::ProjectLogs); + assert_eq!(components.span_id.as_deref(), Some("span-1")); + assert_eq!(components.root_span_id.as_deref(), Some("root-1")); + assert_eq!( + components + .compute_object_metadata_args + .as_ref() + .and_then(|value| value.get("project_name")) + .and_then(serde_json::Value::as_str), + Some("Agents") + ); + } + + #[test] + fn import_parent_rejects_incomplete_or_conflicting_identifiers() { + assert!(ImportCli::try_parse_from([ + "test", + "codex", + "session-a", + "--parent-span-id", + "span-1", + ]) + .is_err()); + assert!(ImportCli::try_parse_from([ + "test", + "codex", + "session-a", + "--destination", + "project_logs:project-id", + "--parent-span-id", + "span-1", + "--parent-root-span-id", + "root-1", + "--parent-object-type", + "project_logs", + "--parent-project", + "Agents", + ]) + .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, + parent_span_id: None, + parent_root_span_id: None, + parent_object_type: None, + parent_object_id: None, + parent_project: None, + attach: true, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + }; + assert!(validate_import_selection(&args) + .unwrap_err() + .to_string() + .contains("exactly one")); + } + + #[test] + fn json_string_fields_accept_strings_and_numbers_only() { + let payload = serde_json::json!({ + "string": "session", + "number": 42, + "boolean": true, + "null": null + }); + assert_eq!( + json_str_field(&payload, "string").as_deref(), + Some("session") + ); + assert_eq!(json_str_field(&payload, "number").as_deref(), Some("42")); + assert_eq!(json_str_field(&payload, "boolean"), None); + assert_eq!(json_str_field(&payload, "null"), None); + assert_eq!(json_str_field(&payload, "missing"), None); + } + + #[test] + fn clock_returns_a_positive_epoch_timestamp() { + assert!(now_ms() > 0); + } + + #[test] + fn transcript_observation_captures_compact_and_full_boundaries() { + let dir = tempfile::tempdir().unwrap(); + let compact = dir.path().join("transcript.jsonl"); + let full = dir.path().join("transcript_full.jsonl"); + std::fs::write(&compact, b"compact\n").unwrap(); + std::fs::write(&full, b"complete record\n").unwrap(); + let mut payload = serde_json::json!({ + "transcriptPath": compact.to_string_lossy() + }); + + add_transcript_observation(&mut payload, "transcriptPath"); + + let observed = &payload["_bt_transcript_observation"]; + assert_eq!(observed["path"], compact.to_string_lossy().as_ref()); + assert_eq!(observed["observed_bytes"], 8); + assert_eq!(observed["full_path"], full.to_string_lossy().as_ref()); + assert_eq!(observed["full_observed_bytes"], 16); + } + + #[test] + fn import_destination_without_session_config_fails_fast() { + let mut config = None; + let destination = wire::TraceDestination::ProjectLogs { + project_id: Some("project-id".to_string()), + project_name: None, + }; + + let error = apply_import_destination(&mut config, Some(destination)).unwrap_err(); + + assert!(error + .to_string() + .contains("import destination requires a resolved Braintrust session configuration")); + } + + fn import_test_config(project_id: &str) -> SessionConfig { + SessionConfig { + auth: wire::BackendAuth { + token: "test".into(), + api_url: Some("https://api.example.test".into()), + app_url: None, + org_name: Some("test-org".into()), + org_id: Some("test-org-id".into()), + }, + destination: Some(wire::TraceDestination::ProjectLogs { + project_id: Some(project_id.into()), + project_name: None, + }), + flush_mode: wire::FlushMode::FireAndForget, + additional_metadata: None, + tags: Vec::new(), + span_plugins: Vec::new(), + } + } + + fn imported_terminal_span_ids(path: &std::path::Path) -> Vec { + std::fs::read_to_string(path) + .unwrap() + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .filter_map(|op| { + let row = op.get("Insert").or_else(|| op.get("Merge"))?; + row.get("end_ms")?.as_i64()?; + row.get("span_id")?.as_str().map(str::to_string) + }) + .collect() + } + + #[tokio::test] + async fn imports_share_destination_delivery_ledger_and_replay_to_new_destinations() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp.path().join("session.jsonl"); + std::fs::write( + &transcript, + format!( + "{}\n{}\n{}\n{}\n", + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"ledger-import","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":"turn-1"}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"done"}}), + ), + ) + .unwrap(); + let ledger_dir = temp.path().join("ledger"); + let first_output = temp.path().join("first"); + import_transcript_with_ledger( + &transcript, + ImportSource::Codex, + debug_serve_options("test", &first_output), + Some(import_test_config("project-a")), + false, + Some(ledger_dir.clone()), + ) + .await + .unwrap(); + let spans = first_output.join("spans/ledger-import.ndjson"); + let first_terminals = imported_terminal_span_ids(&spans); + assert!(!first_terminals.is_empty()); + let first_operations = std::fs::read_to_string(&spans).unwrap(); + + import_transcript_with_ledger( + &transcript, + ImportSource::Codex, + debug_serve_options("test", &first_output), + Some(import_test_config("project-a")), + false, + Some(ledger_dir.clone()), + ) + .await + .unwrap(); + assert_eq!(imported_terminal_span_ids(&spans), first_terminals); + assert_eq!(std::fs::read_to_string(&spans).unwrap(), first_operations); + + let second_output = temp.path().join("second"); + import_transcript_with_ledger( + &transcript, + ImportSource::Codex, + debug_serve_options("test", &second_output), + Some(import_test_config("project-b")), + false, + Some(ledger_dir), + ) + .await + .unwrap(); + assert_eq!( + imported_terminal_span_ids(&second_output.join("spans/ledger-import.ndjson")), + first_terminals + ); + } + + fn test_run_hook_command() -> RunHookCommand { + RunHookCommand { + program: OsString::from("/opt/Braintrust CLI/bt"), + args: vec![OsString::from("agents"), OsString::from("hook")], + } + } + + #[tokio::test] + async fn managed_run_rejects_a_missing_destination_before_launch() { + let error = run_traced( + RunArgs { + source: RunSource::Codex, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + agent_args: Vec::new(), + }, + test_run_hook_command(), + SessionRoute::default(), + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("requires a trace destination")); + } + + #[tokio::test] + async fn codex_managed_run_rejects_global_hook_trust_bypass_before_launch() { + let error = run_traced( + RunArgs { + source: RunSource::Codex, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + agent_args: vec![OsString::from("--dangerously-bypass-hook-trust")], + }, + test_run_hook_command(), + SessionRoute { + destination: Some(wire::TraceDestination::ProjectLogs { + project_id: Some("project-id".into()), + project_name: None, + }), + ..SessionRoute::default() + }, + ) + .await + .unwrap_err(); + + assert!(error + .to_string() + .contains("cannot be combined with --dangerously-bypass-hook-trust")); + } + + #[test] + fn codex_managed_run_injects_live_hooks() { + let args = managed_run_args(RunSource::Codex, &test_run_hook_command()).unwrap(); + assert_eq!(args[0], "--enable"); + assert_eq!(args[1], "hooks"); + assert!(!args + .iter() + .any(|arg| arg == "--dangerously-bypass-hook-trust")); + assert_eq!( + args.iter().filter(|arg| *arg == "-c").count(), + CODEX_RUN_HOOK_EVENTS.len() + ); + let config = args + .iter() + .find_map(|arg| { + let arg = arg.to_str()?; + arg.starts_with("hooks.SessionStart=").then_some(arg) + }) + .unwrap(); + assert!(config.contains("--managed-run-hook")); + assert!(config.contains("agents")); + assert!(config.contains("hook")); + assert!(config.contains("--source")); + assert!(config.contains("codex")); + assert!(!config.contains("transcript")); + } + + #[test] + fn claude_managed_run_injects_live_hooks() { + let args = managed_run_args(RunSource::Claude, &test_run_hook_command()).unwrap(); + assert_eq!(args[0], "--settings"); + let settings: serde_json::Value = serde_json::from_str(args[1].to_str().unwrap()).unwrap(); + let hooks = settings["hooks"].as_object().unwrap(); + assert_eq!(hooks.len(), CLAUDE_RUN_HOOK_EVENTS.len()); + for event in ["SessionStart", "PreToolUse"] { + let matcher_groups = hooks[event].as_array().unwrap(); + assert_eq!(matcher_groups.len(), 1); + assert!(matcher_groups[0]["hooks"].is_array()); + assert!(matcher_groups[0]["hooks"][0]["hooks"].is_null()); + } + let command = hooks["SessionStart"][0]["hooks"][0]["command"] + .as_str() + .unwrap(); + assert!(command.contains("--managed-run-hook")); + assert!(command.contains("agents")); + assert!(command.contains("hook")); + assert!(command.contains("--source")); + assert!(command.contains("claude")); + assert!(!command.contains("transcript")); + } + + #[test] + fn opencode_managed_run_preserves_inline_config_and_adds_plugin() { + let config = + opencode_managed_config(Some(r#"{"model":"test/model","plugin":["other"]}"#)).unwrap(); + let config: serde_json::Value = serde_json::from_str(&config).unwrap(); + assert_eq!(config["model"], "test/model"); + assert_eq!( + config["plugin"], + serde_json::json!(["other", "@braintrust/trace-opencode/tracing"]) + ); + assert!( + managed_run_args(RunSource::OpenCode, &test_run_hook_command()) + .unwrap() + .is_empty() + ); + } + + #[test] + fn pi_managed_run_loads_the_npm_extension() { + assert_eq!( + managed_run_args(RunSource::Pi, &test_run_hook_command()).unwrap(), + vec![ + OsString::from("-e"), + OsString::from(crate::setup::pi_plugin_spec()), + ] + ); + } + + #[test] + fn managed_hook_commands_quote_frontend_paths() { + let hook = test_run_hook_command(); + let unix = managed_hook_shell_command(&hook, "codex", false).unwrap(); + assert!(unix.contains("'/opt/Braintrust CLI/bt' 'agents' 'hook' '--source' 'codex'")); + let windows = managed_hook_shell_command(&hook, "claude", true).unwrap(); + assert!(windows + .contains("\"/opt/Braintrust CLI/bt\" \"agents\" \"hook\" \"--source\" \"claude\"")); + } +} diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index 2af7390b..87576031 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/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs index 6d713459..0c89065c 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -1,4 +1,3 @@ -use crate::agents::{registrar, Agent}; use crate::wire::Envelope; use crate::ImportSource; use anyhow::{bail, Context}; @@ -10,171 +9,98 @@ use std::path::{Path, PathBuf}; mod antigravity; mod claude; mod codex; -mod cursor; -mod pipeline; - -pub use pipeline::{import_transcript, import_transcripts, run_import}; - -/// Importing one agent's native transcripts. Implemented in each agent's -/// module here and registered in [`crate::agents`]. -pub(crate) trait TranscriptImport: Agent { - /// Every session transcript the agent has saved under `home`. - fn discover(&self, home: &Path) -> anyhow::Result>; - /// The saved transcript for one session under `home`. - fn find(&self, home: &Path, session_id: &str) -> anyhow::Result; - /// State for following one growing transcript. - fn tail(&self) -> Box; -} - -/// Where an agent saves its transcripts and how they are named. Each agent -/// builds one privately; this holds the search shared by all of them. -pub(super) struct TranscriptLayout { - pub display_name: &'static str, - /// Directories searched recursively for `.jsonl` transcripts. - pub roots: Vec, - /// The session a transcript belongs to, if the path is one. - pub session_id: fn(&Path) -> Option, - /// Whether a path is named for a session id. - pub filename_matches: fn(&Path, &str) -> bool, -} - -/// Follows one growing transcript, turning the records read so far into -/// synthetic hook events and emitting only those not emitted before. -pub(crate) trait TailSession: Send { - /// Synthetic hook events for every record read so far. - fn envelopes( - &mut self, - path: &Path, - records: &IncrementalRecords, - ) -> anyhow::Result>; - - /// The events from `events` to emit now. `finalize` ends the session. - fn poll( - &mut self, - events: Vec, - len: u64, - finalize: bool, - ) -> anyhow::Result>; - - /// Whether a rewritten transcript continues with the same translator - /// instead of starting a new one. If so, the fresh tail that replaces this - /// one receives [`Self::rewritten`]. - fn keeps_translator_on_rewrite(&self) -> bool { - false - } - - /// Called on a fresh tail after a rewrite, with the records from before - /// and after it, when [`Self::keeps_translator_on_rewrite`] is true. - fn rewritten(&mut self, _previous: &[Value], _current: &[Value]) {} - - /// Whether an attached import may stop at a final record that is still - /// being written when it shuts down. - fn tolerates_incomplete_final_record(&self) -> bool { - false - } -} - -fn importer(source: ImportSource) -> &'static dyn TranscriptImport { - let name = source.identity().id; - registrar() - .import - .get(name) - .unwrap_or_else(|| panic!("no transcript import is registered for {name}")) -} +pub(crate) mod muse; pub(crate) fn resolve_transcripts( session_ids: &[String], all: bool, source: ImportSource, ) -> anyhow::Result> { - let home = std::env::var_os("HOME") - .or_else(|| std::env::var_os("USERPROFILE")) - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(".")); - let importer = importer(source); match (all, session_ids.is_empty()) { - (true, true) => importer.discover(&home), + (true, true) => discover_transcripts(source), (false, false) => session_ids .iter() - .map(|session_id| importer.find(&home, session_id)) + .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"), } } -impl TranscriptLayout { - /// One transcript per session found under the roots, in path order. - pub fn discover(&self) -> anyhow::Result> { - let mut sessions = BTreeMap::>::new(); - for path in self.candidates() { - if let Some(session_id) = (self.session_id)(&path) { - sessions.entry(session_id).or_default().push(path); - } - } - if sessions.is_empty() { - bail!( - "no {} transcripts found; searched {}", - self.display_name, - self.locations() - ); - } - sessions - .into_iter() - .map(|(session_id, paths)| self.single(&session_id, paths)) - .collect() +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::Muse => Vec::new(), + ImportSource::Claude => claude::roots(&home), + ImportSource::Antigravity => antigravity::roots(&home), } +} - /// The one transcript named for `session_id`. - pub fn find(&self, session_id: &str) -> anyhow::Result { - validate_session_id(session_id)?; - let matches: Vec<_> = self - .candidates() - .into_iter() - .filter(|path| (self.filename_matches)(path, session_id)) - .collect(); - if matches.is_empty() { - bail!( - "no {} transcript found for session {session_id}; searched {}", - self.display_name, - self.locations() - ); - } - self.single(session_id, matches) +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(); - fn candidates(&self) -> Vec { - let mut candidates = Vec::new(); - for root in &self.roots { - find_jsonl_files(root, &mut candidates); - } - candidates.sort(); - candidates.dedup(); - candidates + 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) + ); } - fn single(&self, session_id: &str, paths: Vec) -> anyhow::Result { - match <[PathBuf; 1]>::try_from(paths) { - Ok([path]) => Ok(path), - Err(paths) => bail!( - "multiple {} transcripts found for session {session_id}: {}", - self.display_name, - paths + 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(", ") - ), + .join(", "); + bail!( + "multiple {} transcripts found for session {session_id}: {locations}", + source_name(source) + ); + } } } - - fn locations(&self) -> String { - self.roots - .iter() - .map(|root| root.display().to_string()) - .collect::>() - .join(", ") - } + Ok(resolved) } fn find_jsonl_files(directory: &Path, matches: &mut Vec) { @@ -197,6 +123,61 @@ 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), + } +} + +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::Muse => false, + ImportSource::Claude => claude::filename_matches(path, session_id), + ImportSource::Antigravity => antigravity::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 @@ -208,20 +189,47 @@ fn validate_session_id(session_id: &str) -> anyhow::Result<()> { Ok(()) } +fn source_name(source: ImportSource) -> &'static str { + match source { + ImportSource::Codex => "Codex", + ImportSource::Muse => "Muse Code", + ImportSource::Claude => "Claude Code", + ImportSource::Antigravity => "Google Antigravity", + } +} + #[cfg(test)] pub(crate) fn transcript_envelopes( path: &Path, source: ImportSource, ) -> anyhow::Result> { let mut records = IncrementalRecords::default(); - records.refresh_with_final_partial(path, true, false)?; - importer(source).tail().envelopes(path, &records) + records.refresh(path, true)?; + envelopes_from_records(path, source, &records) +} + +fn envelopes_from_records( + path: &Path, + source: ImportSource, + records: &IncrementalRecords, +) -> anyhow::Result> { + match source { + ImportSource::Codex => codex::envelopes(path, &records.values), + ImportSource::Muse => muse::envelopes(path), + ImportSource::Claude => claude::envelopes( + path, + &records.values, + &records.end_offsets, + records.read_offset, + ), + ImportSource::Antigravity => { + antigravity::envelopes(path, &records.values, &records.end_offsets) + } + } } -/// The records of a transcript read so far, re-read incrementally as it -/// grows. #[derive(Default)] -pub(crate) struct IncrementalRecords { +struct IncrementalRecords { values: Vec, end_offsets: Vec, read_offset: u64, @@ -238,12 +246,7 @@ enum Refresh { } impl IncrementalRecords { - fn refresh_with_final_partial( - &mut self, - path: &Path, - finalize: bool, - allow_incomplete_final_record: bool, - ) -> anyhow::Result { + fn refresh(&mut self, path: &Path, finalize: bool) -> anyhow::Result { let metadata = std::fs::metadata(path) .with_context(|| format!("read transcript metadata {}", path.display()))?; let len = metadata.len(); @@ -295,13 +298,7 @@ impl IncrementalRecords { values.push(value); end_offsets.push(offset); } - Err(_) - if (allow_incomplete_final_record || !finalize) - && offset == len - && !line.ends_with('\n') => - { - break - } + Err(_) if !finalize && offset == len && !line.ends_with('\n') => break, Err(error) => { return Err(error).with_context(|| { format!("parse transcript {} line {}", path.display(), line_count) @@ -353,44 +350,43 @@ fn read_anchor(path: &Path, through: u64, max_len: usize) -> anyhow::Result, + source: ImportSource, + state: TailState, records: IncrementalRecords, retry_envelopes: bool, translator_reset: bool, - allow_incomplete_final_record: bool, +} + +enum TailState { + Codex(codex::Tail), + Muse, + Claude(claude::Tail), + Antigravity(antigravity::Tail), } impl TranscriptTail { pub(crate) fn new(path: PathBuf, source: ImportSource) -> Self { - let importer = importer(source); Self { path, - importer, - session: importer.tail(), + source, + state: Self::new_state(source), records: IncrementalRecords::default(), retry_envelopes: false, translator_reset: false, - allow_incomplete_final_record: false, } } - pub(crate) fn allow_incomplete_final_record_on_shutdown(&mut self) { - self.allow_incomplete_final_record = true; + 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()), + } } pub(crate) fn poll(&mut self, finalize: bool) -> anyhow::Result> { - let keeps_translator = self.session.keeps_translator_on_rewrite(); - let previous_records = if keeps_translator { - self.records.values.clone() - } else { - Vec::new() - }; - let refresh = match self.records.refresh_with_final_partial( - &self.path, - finalize, - self.session.tolerates_incomplete_final_record() && self.allow_incomplete_final_record, - ) { + let refresh = match self.records.refresh(&self.path, finalize) { Ok(refresh) => refresh, Err(_) if !finalize => return Ok(Vec::new()), Err(error) => return Err(error), @@ -399,15 +395,10 @@ impl TranscriptTail { return Ok(Vec::new()); } if refresh == Refresh::Reset { - self.session = self.importer.tail(); - if keeps_translator { - self.session - .rewritten(&previous_records, &self.records.values); - } else { - self.translator_reset = true; - } + self.state = Self::new_state(self.source); + self.translator_reset = true; } - let events = match self.session.envelopes(&self.path, &self.records) { + let events = match envelopes_from_records(&self.path, self.source, &self.records) { Ok(events) => events, Err(_) if !finalize => { self.retry_envelopes = true; @@ -419,7 +410,12 @@ impl TranscriptTail { let len = std::fs::metadata(&self.path) .with_context(|| format!("read transcript metadata {}", self.path.display()))? .len(); - self.session.poll(events, len, finalize) + 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), + } } pub(crate) fn take_translator_reset(&mut self) -> bool { @@ -530,7 +526,7 @@ mod tests { std::fs::write(&transcript, "{}\n").unwrap(); assert_eq!( - codex::layout(vec![root]).find("session-123").unwrap(), + resolve_transcript_in("session-123", ImportSource::Codex, &[root]).unwrap(), transcript ); } @@ -546,7 +542,7 @@ mod tests { std::fs::write(&transcript, "{}\n").unwrap(); assert_eq!( - claude::layout(vec![root]).find("session-123").unwrap(), + resolve_transcript_in("session-123", ImportSource::Claude, &[root]).unwrap(), transcript ); } @@ -567,252 +563,19 @@ mod tests { std::fs::write(unrelated, "{}\n").unwrap(); assert_eq!( - antigravity::layout(vec![root]) - .find("conversation-123") - .unwrap(), + resolve_transcript_in("conversation-123", ImportSource::Antigravity, &[root]).unwrap(), transcript ); } - #[test] - fn finds_cursor_transcripts_by_nested_session_identity_and_ignores_symlinks() { - let temp = tempfile::tempdir().unwrap(); - let root = temp.path().join("projects"); - let transcript = root - .join("workspace") - .join("agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - std::fs::write(&transcript, r#"{"role":"user","message":{"content":"hi"}}"#).unwrap(); - let subagent = root - .join("workspace") - .join("agent-transcripts/agent-session-123/agent-session-123.jsonl"); - std::fs::create_dir_all(subagent.parent().unwrap()).unwrap(); - std::fs::write(&subagent, "{}\n").unwrap(); - - assert_eq!( - cursor::layout(vec![root.clone()]) - .find("session-123") - .unwrap(), - transcript - ); - assert_eq!( - cursor::layout(vec![root]).discover().unwrap(), - vec![transcript] - ); - } - - #[test] - fn cursor_import_emits_only_supported_transcript_enrichment_and_boundaries() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp - .path() - .join("workspace/agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - let contents = concat!( - "{\"role\":\"user\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"question\"}]}}\n", - "{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"tool_use\",\"name\":\"Read\",\"input\":{\"path\":\"a.txt\"}}]}}\n", - "{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"answer\"}]}}\n", - "{\"type\":\"turn_ended\",\"status\":\"success\"}\n" - ); - std::fs::write(&transcript, contents).unwrap(); - let events = transcript_envelopes(&transcript, ImportSource::Cursor).unwrap(); - - assert_eq!( - events - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["_bt_importStart", "_bt_importCheckpoint", "_bt_importStop"] - ); - assert!(events.iter().all(|event| event.source == "cursor")); - assert_eq!( - events[1].payload["_bt_transcript_mirror"]["through"], - contents.len() as u64 - ); - assert_eq!( - events[1].payload["historical_fidelity"], - "user_and_assistant_text_and_terminal_status" - ); - assert_eq!(events[2].payload["status"], "success"); - assert_eq!(events[1].payload["start_time_estimated"], true); - } - - #[test] - fn cursor_import_preserves_terminal_error_message() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp - .path() - .join("workspace/agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - std::fs::write( - &transcript, - "{\"role\":\"user\",\"message\":{\"content\":\"question\"}}\n{\"type\":\"turn_ended\",\"status\":\"error\",\"error\":\"WritableIterable is closed\"}\n", - ) - .unwrap(); - - let events = transcript_envelopes(&transcript, ImportSource::Cursor).unwrap(); - let stop = events - .iter() - .find(|event| event.event == "_bt_importStop") - .unwrap(); - assert_eq!(stop.payload["status"], "error"); - assert_eq!(stop.payload["error_message"], "WritableIterable is closed"); - } - - #[test] - fn cursor_attach_tails_growth_without_repeating_boundaries_or_checkpoints() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp - .path() - .join("workspace/agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - let first_record = b"{\"role\":\"user\",\"message\":{\"content\":\"one\"}}\n"; - std::fs::write(&transcript, first_record).unwrap(); - let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); - let first = tail.poll(false).unwrap(); - assert_eq!( - first - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["_bt_importStart", "_bt_importCheckpoint"] - ); - let snapshot = first[1].payload["_bt_transcript_mirror"]["mirror"] - .as_str() - .unwrap() - .to_owned(); - assert_ne!(PathBuf::from(&snapshot), transcript); - assert_eq!( - first[1].payload["_bt_transcript_mirror"]["through"], - first_record.len() as u64 - ); - assert!(tail.poll(false).unwrap().is_empty()); - - use std::io::Write; - std::fs::OpenOptions::new() - .append(true) - .open(&transcript) - .unwrap() - .write_all(b"{\"role\":\"assistant\",\"message\":{\"content\":\"two\"}}\n") - .unwrap(); - let grown = tail.poll(false).unwrap(); - assert_eq!( - grown - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["_bt_importCheckpoint"] - ); - assert_eq!( - grown[0].payload["_bt_transcript_mirror"]["mirror"], - snapshot - ); - assert_eq!( - std::fs::read_to_string(&snapshot).unwrap().lines().count(), - 2 - ); - assert_eq!( - tail.poll(true) - .unwrap() - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["_bt_importStop"] - ); - std::fs::write(&transcript, b"replacement transcript\n").unwrap(); - let snapshot_contents = std::fs::read_to_string(&snapshot).unwrap(); - assert!(snapshot_contents.contains("\"content\":\"one\"")); - assert!(snapshot_contents.contains("\"content\":\"two\"")); - } - - #[test] - fn cursor_transcript_replacement_rotates_snapshot_without_resetting_translator() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp - .path() - .join("workspace/agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - std::fs::write( - &transcript, - "{\"role\":\"user\",\"message\":{\"content\":\"the original longer prompt\"}}\n", - ) - .unwrap(); - let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); - let first = tail.poll(false).unwrap(); - let first_snapshot = first[1].payload["_bt_transcript_mirror"]["mirror"] - .as_str() - .unwrap() - .to_owned(); - - std::fs::write( - &transcript, - "{\"role\":\"user\",\"message\":{\"content\":\"latest\"}}\n", - ) - .unwrap(); - let replacement = tail.poll(false).unwrap(); - assert_eq!(replacement[0].event, "_bt_importStart"); - assert_eq!(replacement[1].event, "_bt_importCheckpoint"); - let replacement_snapshot = replacement[1].payload["_bt_transcript_mirror"]["mirror"] - .as_str() - .unwrap(); - assert_ne!(first_snapshot, replacement_snapshot); - assert!(!tail.take_translator_reset()); - } - - #[test] - fn cursor_attach_finalizes_after_an_incomplete_trailing_record() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp - .path() - .join("workspace/agent-transcripts/session-123/session-123.jsonl"); - std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); - std::fs::write( - &transcript, - "{\"role\":\"user\",\"message\":{\"content\":\"complete\"}}\n", - ) - .unwrap(); - - let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); - tail.allow_incomplete_final_record_on_shutdown(); - assert_eq!( - tail.poll(false).unwrap().last().unwrap().event, - "_bt_importCheckpoint" - ); - use std::io::Write; - std::fs::OpenOptions::new() - .append(true) - .open(&transcript) - .unwrap() - .write_all(b"{\"role\":\"assistant\",\"message\":") - .unwrap(); - - let final_events = tail.poll(true).unwrap(); - assert_eq!( - final_events - .iter() - .map(|event| event.event.as_str()) - .collect::>(), - vec!["_bt_importCheckpoint", "_bt_importStop"] - ); - } - - #[test] - fn completed_import_still_rejects_an_incomplete_trailing_record() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp.path().join("transcript.jsonl"); - std::fs::write(&transcript, "{\"type\":\"complete\"}\n{\"partial\"").unwrap(); - - let error = IncrementalRecords::default() - .refresh_with_final_partial(&transcript, true, false) - .unwrap_err(); - assert!(error.to_string().contains("parse transcript")); - } - #[test] fn rejects_unsafe_session_ids() { - let error = codex::layout(vec![PathBuf::from("unused")]) - .find("../session") - .unwrap_err(); + let error = resolve_transcript_in( + "../session", + ImportSource::Codex, + &[PathBuf::from("unused")], + ) + .unwrap_err(); assert!(error.to_string().contains("invalid session id")); } @@ -824,16 +587,18 @@ mod tests { std::fs::create_dir_all(&first).unwrap(); std::fs::create_dir_all(&second).unwrap(); - let missing = claude::layout(vec![first.clone(), second.clone()]) - .find("missing") - .unwrap_err(); + 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 = claude::layout(vec![first, second]) - .find("duplicate") - .unwrap_err(); + let ambiguous = + resolve_transcript_in("duplicate", ImportSource::Claude, &[first, second]).unwrap_err(); assert!(ambiguous .to_string() .contains("multiple Claude Code transcripts")); @@ -866,7 +631,7 @@ mod tests { std::fs::write(root.join("not-a-transcript.jsonl"), "{}\n").unwrap(); assert_eq!( - codex::layout(vec![root]).discover().unwrap(), + discover_transcripts_in(ImportSource::Codex, &[root]).unwrap(), vec![first, second] ); } @@ -892,15 +657,13 @@ mod tests { std::fs::write(project.join("notes.jsonl"), "{}\n").unwrap(); assert_eq!( - claude::layout(vec![root]).discover().unwrap(), + discover_transcripts_in(ImportSource::Claude, &[root]).unwrap(), vec![transcript] ); } #[test] - fn claude_import_keeps_delayed_child_and_native_handback_under_launching_turn() { - use crate::translate::{Registry, SessionCtx, SpanOp, SpanType}; - + 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"); @@ -908,12 +671,8 @@ mod tests { 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:03Z","message":{"content":"unrelated question"}}), 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"}]}}), - json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:07Z","origin":{"kind":"peer","handback":true,"senderTaskId":"child-a","from":"child-a","name":"reviewer"},"message":{"content":"Another Claude session sent a message:\nreviewed\nThis agent is working inside the same session."}}), - json!({"type":"assistant","sessionId":"session-a","timestamp":"2026-01-01T00:00:08Z","message":{"id":"handback-response","content":[{"type":"text","text":"accepted review"}]}}), - json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:09Z","message":{"content":"next human question"}}), ]; std::fs::write( &transcript, @@ -931,72 +690,27 @@ mod tests { .unwrap(); let events = transcript_envelopes(&transcript, ImportSource::Claude).unwrap(); - let context = SessionCtx { - session_id: "session-a".into(), - config: None, - }; - let mut translator = Registry::default_agents().create("claude-code", "session-a"); - let mut ops = Vec::new(); - for event in &events { - ops.extend(translator.handle(event, &context).unwrap()); - while let Some(pending) = translator.drain_pending(&context).unwrap() { - ops.extend(pending); - } - } - ops.extend(translator.finalize(&context).unwrap()); - while let Some(pending) = translator.drain_pending(&context).unwrap() { - ops.extend(pending); - } - let inserts = ops - .iter() - .filter_map(|op| match op { - SpanOp::Insert(row) => Some(row), - SpanOp::Merge(_) => None, - }) - .collect::>(); - let turns = inserts + let names = events .iter() - .filter(|row| row.name.starts_with("Turn ")) + .map(|event| event.event.as_str()) .collect::>(); assert_eq!( - turns - .iter() - .map(|row| row.name.as_str()) - .collect::>(), - ["Turn 1", "Turn 2", "Turn 3"] + names, + vec![ + "SessionStart", + "UserPromptSubmit", + "SubagentStart", + "SubagentStop", + "Stop", + "SessionEnd" + ] ); - let launch = turns[0]; - let child = inserts - .iter() - .find(|row| { - row.span_type == SpanType::Task - && row - .metadata - .as_ref() - .and_then(|value| value.get("agent_id")) - == Some(&json!("child-a")) - }) - .unwrap(); - assert_eq!(child.parent_span_ids, std::slice::from_ref(&launch.span_id)); - let continuation = inserts - .iter() - .find(|row| { - row.metadata - .as_ref() - .and_then(|value| value.get("turn_trigger")) - == Some(&json!("agent_message")) - }) - .unwrap(); + assert_eq!(events[2].payload["agent_id"], json!("child-a")); + assert_eq!(events[2].payload["agent_type"], json!("reviewer")); assert_eq!( - continuation.parent_span_ids, - std::slice::from_ref(&launch.span_id) + Path::new(events[3].payload["agent_transcript_path"].as_str().unwrap()), + subagent ); - assert!(ops.iter().any(|op| match op { - SpanOp::Insert(row) | SpanOp::Merge(row) => { - row.span_id == continuation.span_id - && row.output.as_ref() == Some(&json!("accepted review")) - } - })); } #[test] @@ -1280,12 +994,7 @@ mod tests { let path = temp.path().join("session.jsonl"); std::fs::write(&path, "{\"type\":\"session_meta\"").unwrap(); let mut records = IncrementalRecords::default(); - assert_eq!( - records - .refresh_with_final_partial(&path, false, false) - .unwrap(), - Refresh::Unchanged - ); + assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Unchanged); assert!(records.values.is_empty()); std::fs::write( @@ -1293,21 +1002,11 @@ mod tests { "{\"type\":\"session_meta\",\"payload\":{\"id\":\"session\"}}\n", ) .unwrap(); - assert_eq!( - records - .refresh_with_final_partial(&path, false, false) - .unwrap(), - Refresh::Appended - ); + assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Appended); assert_eq!(records.values.len(), 1); std::fs::write(&path, "{\"type\":\"event_msg\"}\n").unwrap(); - assert_eq!( - records - .refresh_with_final_partial(&path, false, false) - .unwrap(), - Refresh::Reset - ); + assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Reset); assert_eq!(records.values, vec![json!({"type":"event_msg"})]); std::fs::write( @@ -1315,12 +1014,7 @@ mod tests { "{\"type\":\"replacement_with_a_longer_prefix\"}\n{\"type\":\"second\"}\n", ) .unwrap(); - assert_eq!( - records - .refresh_with_final_partial(&path, false, false) - .unwrap(), - Refresh::Reset - ); + assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Reset); assert_eq!( records.values, vec![ diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs new file mode 100644 index 00000000..0d4dd189 --- /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/tests/replay.rs b/bt-daemon/tests/replay.rs index 317ff27a..10226f05 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -20,6 +20,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; From 2d5c8adc85c24351b3080b0865aae759466f6fa5 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 11 Sep 2026 04:01:12 +0800 Subject: [PATCH 03/17] Document Muse tracing limitations --- bt-daemon/README.md | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 8f78c933..7d873b09 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -63,6 +63,8 @@ coding agent decides whether to continue its turn. 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` @@ -272,13 +274,19 @@ echo '{"session_id":"s1","hook_event_name":"Stop"}' | ./bt-daemon/target 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. They import completed sessions only: `--attach` +and `--all` require MSP session enumeration and cursor support, which are not +implemented yet. 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 Braintrust hooks or an adapter for that invocation, so it From cf982164ea9b64ab1269650c117a2dc9cb174934 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 11 Sep 2026 04:09:18 +0800 Subject: [PATCH 04/17] Import all Muse sessions through MSP --- bt-daemon/README.md | 10 +- bt-daemon/src/lib.rs | 160 ++++++++++++++++++++++++- bt-daemon/src/transcript_import/mod.rs | 2 +- 3 files changed, 164 insertions(+), 8 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 7d873b09..b1387ff4 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -282,10 +282,12 @@ 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. They import completed sessions only: `--attach` -and `--all` require MSP session enumeration and cursor support, which are not -implemented yet. Muse also has no safe invocation-local configuration overlay, -so `run muse` is deliberately unavailable; use persistent `enable muse` setup. +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 diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 1386e87c..205d2c76 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -62,7 +62,9 @@ use anyhow::Context; use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; use clap::{Args, ValueEnum}; use std::ffi::OsString; +use std::io::{BufRead, BufReader, Write}; use std::path::PathBuf; +use std::process::{Command as ProcessCommand, Stdio}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use wire::{ @@ -785,10 +787,15 @@ pub async fn run_import( "Muse import does not support --attach yet; use a completed export-v1 session" ) } - if args.all { - anyhow::bail!("Muse import does not support --all yet; pass one or more session ids") + let session_ids = if args.all { + muse_msp_session_ids()? + } else { + args.session_ids.clone() + }; + if session_ids.is_empty() { + anyhow::bail!("no durable Muse sessions found to import") } - let exports = export_muse_sessions(&args.session_ids)?; + let exports = export_muse_sessions(&session_ids)?; return import_muse_exports_with_ledger( &exports, opts, @@ -870,6 +877,127 @@ fn export_muse_sessions(session_ids: &[String]) -> anyhow::Result 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 = ProcessCommand::new(&executable) + .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); + 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"}}), + )?; + // 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")?; + 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}), + )?; + 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()?; + if !output.status.success() { + anyhow::bail!( + "Muse MSP session listing failed: {}", + String::from_utf8_lossy(&output.stderr).trim() + ); + } + sessions.sort(); + sessions.dedup(); + Ok(sessions) +} + +fn muse_msp_notify(stdin: &mut impl Write, method: &str) -> anyhow::Result<()> { + serde_json::to_writer( + &mut *stdin, + &serde_json::json!({"jsonrpc":"2.0", "method":method}), + )?; + stdin.write_all(b"\n")?; + stdin.flush()?; + Ok(()) +} + +fn muse_msp_request( + stdin: &mut impl Write, + reader: &mut impl BufRead, + id: u64, + method: &str, + params: serde_json::Value, +) -> anyhow::Result { + serde_json::to_writer( + &mut *stdin, + &serde_json::json!({"jsonrpc":"2.0", "id":id, "method":method, "params":params}), + )?; + stdin.write_all(b"\n")?; + stdin.flush()?; + let mut line = String::new(); + for _ in 0..10_000 { + line.clear(); + if reader.read_line(&mut line)? == 0 { + anyhow::bail!("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")); + } + anyhow::bail!("Muse MSP sent too many frames before responding to {method}") +} + async fn import_muse_exports_with_ledger( directory: &tempfile::TempDir, opts: ServeOptions, @@ -1746,6 +1874,7 @@ mod tests { use super::*; use clap::Parser; use serde_json::json; + use std::io::Cursor; #[derive(Debug, Parser)] struct ImportCli { @@ -1928,6 +2057,31 @@ mod tests { assert!(resolved[0].is_absolute()); } + #[test] + fn muse_msp_handshake_writes_notification_and_skips_notifications() { + let mut stdin = Vec::new(); + muse_msp_notify(&mut stdin, "initialized").unwrap(); + assert_eq!( + serde_json::from_slice::(&stdin).unwrap(), + json!({"jsonrpc": "2.0", "method": "initialized"}) + ); + + let mut reader = BufReader::new(Cursor::new( + b"{\"jsonrpc\":\"2.0\",\"method\":\"session/changed\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n", + )); + let result = muse_msp_request( + &mut stdin, + &mut reader, + 7, + "session/list", + json!({"cursor": null, "limit": 200}), + ) + .unwrap(); + assert_eq!(result, json!({"sessions": []})); + let frames = String::from_utf8(stdin).unwrap(); + assert!(frames.contains("\"method\":\"session/list\"")); + } + #[test] fn additional_metadata_overrides_a_route_only_with_a_json_object() { let mut route = SessionRoute { diff --git a/bt-daemon/src/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs index 0c89065c..08577b67 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -178,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() From 8c920465bca6eaa90d45b14834f8d18254e4586d Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 11 Sep 2026 21:44:10 +0800 Subject: [PATCH 05/17] Stream Muse MSP import requests asynchronously --- bt-daemon/src/lib.rs | 172 +++++++++++++++++++++++++------------------ 1 file changed, 102 insertions(+), 70 deletions(-) diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 205d2c76..f361f6fd 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -62,11 +62,11 @@ use anyhow::Context; use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; use clap::{Args, ValueEnum}; use std::ffi::OsString; -use std::io::{BufRead, BufReader, Write}; use std::path::PathBuf; -use std::process::{Command as ProcessCommand, Stdio}; +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, @@ -74,6 +74,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)] @@ -787,15 +788,7 @@ pub async fn run_import( "Muse import does not support --attach yet; use a completed export-v1 session" ) } - let session_ids = if args.all { - muse_msp_session_ids()? - } else { - args.session_ids.clone() - }; - if session_ids.is_empty() { - anyhow::bail!("no durable Muse sessions found to import") - } - let exports = export_muse_sessions(&session_ids)?; + let exports = prepare_muse_exports(&args).await?; return import_muse_exports_with_ledger( &exports, opts, @@ -846,13 +839,26 @@ fn apply_import_span_plugins( Ok(()) } -fn export_muse_sessions(session_ids: &[String]) -> anyhow::Result { +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 = std::process::Command::new(&executable) + let output = tokio::process::Command::new(&executable) + .kill_on_drop(true) .args([ OsString::from("export"), OsString::from("--session"), @@ -861,6 +867,7 @@ fn export_muse_sessions(session_ids: &[String]) -> anyhow::Result anyhow::Result anyhow::Result> { +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 = ProcessCommand::new(&executable) + let mut child = tokio::process::Command::new(&executable) + .kill_on_drop(true) .arg("serve") .stdin(Stdio::piped()) .stdout(Stdio::piped()) @@ -897,7 +905,7 @@ fn muse_msp_session_ids() -> anyhow::Result> { })?; let mut stdin = child.stdin.take().expect("piped stdin"); let stdout = child.stdout.take().expect("piped stdout"); - let mut reader = BufReader::new(stdout); + let mut reader = BufReader::new(stdout).lines(); let mut request_id = 1u64; muse_msp_request( &mut stdin, @@ -905,10 +913,11 @@ fn muse_msp_session_ids() -> anyhow::Result> { 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")?; + muse_msp_notify(&mut stdin, "initialized").await?; let mut cursor = None; let mut sessions = Vec::new(); for _ in 0..MAX_PAGES { @@ -919,7 +928,8 @@ fn muse_msp_session_ids() -> anyhow::Result> { request_id, "session/list", serde_json::json!({"cursor": cursor, "limit": 200}), - )?; + ) + .await?; let page = result .get("sessions") .and_then(serde_json::Value::as_array) @@ -941,7 +951,7 @@ fn muse_msp_session_ids() -> anyhow::Result> { } } drop(stdin); - let output = child.wait_with_output()?; + let output = child.wait_with_output().await?; if !output.status.success() { anyhow::bail!( "Muse MSP session listing failed: {}", @@ -953,49 +963,56 @@ fn muse_msp_session_ids() -> anyhow::Result> { Ok(sessions) } -fn muse_msp_notify(stdin: &mut impl Write, method: &str) -> anyhow::Result<()> { - serde_json::to_writer( - &mut *stdin, - &serde_json::json!({"jsonrpc":"2.0", "method":method}), - )?; - stdin.write_all(b"\n")?; - stdin.flush()?; +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(()) } -fn muse_msp_request( - stdin: &mut impl Write, - reader: &mut impl BufRead, +async fn muse_msp_request( + stdin: &mut W, + reader: &mut Lines, id: u64, method: &str, params: serde_json::Value, -) -> anyhow::Result { - serde_json::to_writer( - &mut *stdin, +) -> 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}), )?; - stdin.write_all(b"\n")?; - stdin.flush()?; - let mut line = String::new(); - for _ in 0..10_000 { - line.clear(); - if reader.read_line(&mut line)? == 0 { - anyhow::bail!("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}") + 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")); } - return frame - .get("result") - .cloned() - .ok_or_else(|| anyhow::anyhow!("Muse MSP {method} response has no result")); - } - anyhow::bail!("Muse MSP sent too many frames before responding to {method}") + }) + .await + .map_err(|_| anyhow::anyhow!("Muse MSP timed out waiting for {method}"))? } async fn import_muse_exports_with_ledger( @@ -1874,7 +1891,7 @@ mod tests { use super::*; use clap::Parser; use serde_json::json; - use std::io::Cursor; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; #[derive(Debug, Parser)] struct ImportCli { @@ -2057,29 +2074,44 @@ mod tests { assert!(resolved[0].is_absolute()); } - #[test] - fn muse_msp_handshake_writes_notification_and_skips_notifications() { - let mut stdin = Vec::new(); - muse_msp_notify(&mut stdin, "initialized").unwrap(); - assert_eq!( - serde_json::from_slice::(&stdin).unwrap(), - json!({"jsonrpc": "2.0", "method": "initialized"}) - ); + #[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(Cursor::new( - b"{\"jsonrpc\":\"2.0\",\"method\":\"session/changed\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n", - )); + let mut reader = BufReader::new(client_reader).lines(); + muse_msp_notify(&mut client_writer, "initialized") + .await + .unwrap(); let result = muse_msp_request( - &mut stdin, + &mut client_writer, &mut reader, 7, "session/list", json!({"cursor": null, "limit": 200}), ) + .await .unwrap(); assert_eq!(result, json!({"sessions": []})); - let frames = String::from_utf8(stdin).unwrap(); - assert!(frames.contains("\"method\":\"session/list\"")); + server.await.unwrap(); } #[test] From 0acd62741fb429c1a654b6de60145200a9086f02 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 14:58:38 -0400 Subject: [PATCH 06/17] Harden Muse tracing setup and diagnostics --- bt-daemon/src/setup.rs | 293 +++++++++++++++++++++++++++++---- bt-daemon/src/trace_runtime.rs | 3 + 2 files changed, 262 insertions(+), 34 deletions(-) diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index 6293e837..5cfd4861 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -5,6 +5,7 @@ use crate::trace_command::{EnableArgs, SetupAgent}; use crate::wire::SessionRoute; use crate::TraceCommandOutput; use anyhow::{bail, Context}; +use fs2::FileExt; use serde_json::{Map, Value}; use std::io::Write; use std::path::{Path, PathBuf}; @@ -704,7 +705,8 @@ fn muse_hook_path(config_dir: &Path) -> PathBuf { } fn muse_hook_config() -> Map { - let command = "bt trace hook --source muse --flush-on-turn-end || true"; + // Muse treats command-hook failures as non-fatal to the agent turn. + let command = "bt trace hook --source muse --flush-on-turn-end"; let hooks = MUSE_HOOK_EVENTS .iter() .map(|event| { @@ -723,57 +725,220 @@ fn muse_hook_config() -> Map { ]) } -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 - ); +fn legacy_muse_hook_config() -> Map { + let mut config = muse_hook_config(); + if let Some(hooks) = config.get_mut("hooks").and_then(Value::as_object_mut) { + for entries in hooks.values_mut() { + if let Some(command) = entries + .get_mut(0) + .and_then(|entry| entry.get_mut("hooks")) + .and_then(|hooks| hooks.get_mut(0)) + .and_then(|hook| hook.get_mut("command")) + { + *command = + Value::String("bt trace hook --source muse --flush-on-turn-end || true".into()); + } } } - 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) + config +} + +fn muse_version(output: &[u8]) -> anyhow::Result { + let text = std::str::from_utf8(output).context("Muse version is not UTF-8")?; + text.split_whitespace() + .find_map(|word| semver::Version::parse(word).ok()) + .ok_or_else(|| { + anyhow::anyhow!( + "Muse Code returned an unrecognized version: {}", + text.trim() + ) + }) +} + +/// Muse serializes changes to settings.json with this native lock file. Hold +/// it across the read, hook update, and settings update so we cannot discard a +/// concurrent change made by Muse itself. +fn with_muse_settings_lock( + config_dir: &Path, + action: impl FnOnce() -> anyhow::Result, +) -> anyhow::Result { + std::fs::create_dir_all(config_dir) + .with_context(|| format!("create Muse config directory {}", config_dir.display()))?; + let lock_path = config_dir.join(".settings.json.lock"); + let lock = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&lock_path) + .with_context(|| format!("open Muse settings lock {}", lock_path.display()))?; + lock.lock_exclusive() + .with_context(|| format!("lock Muse settings {}", lock_path.display()))?; + let result = action(); + FileExt::unlock(&lock)?; + result +} + +fn muse_hook_is_owned(path: &Path) -> anyhow::Result { + match std::fs::read(path) { + Ok(raw) => { + let value: Value = serde_json::from_slice(&raw) + .with_context(|| format!("invalid Muse hook configuration: {}", path.display()))?; + Ok(value == Value::Object(muse_hook_config()) + || value == Value::Object(legacy_muse_hook_config())) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => { + Err(error).with_context(|| format!("read Muse hook configuration: {}", path.display())) + } + } +} + +fn check_muse_hook_ownership(path: &Path) -> anyhow::Result<()> { + if path.exists() && !muse_hook_is_owned(path)? { + bail!("Muse hook file {} differs from Braintrust's generated hooks; refusing to overwrite or remove it", path.display()); + } + Ok(()) +} + +fn setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + 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") { + if existing.as_str() != Some(owned_path.as_ref()) { + bail!( + "Muse already uses managed_hooks_path {}; refusing to overwrite it", + existing + ); + } + } + check_muse_hook_ownership(&hook_path)?; + settings.entry("schema_version").or_insert(Value::from(1)); + settings.insert( + "managed_hooks_path".into(), + Value::String(owned_path.into_owned()), + ); + write_object_atomic_unlocked(&hook_path, muse_hook_config())?; + write_object_atomic_unlocked(&settings_path, settings) + }) } fn disable_muse_at(config_dir: &Path) -> anyhow::Result<()> { - let settings_path = config_dir.join("settings.json"); + with_muse_settings_lock(config_dir, || { + let settings_path = config_dir.join("settings.json"); + let hook_path = muse_hook_path(config_dir); + let mut settings = load_object(&settings_path)?; + check_muse_hook_ownership(&hook_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_unlocked(&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 update_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + let settings = load_object(&config_dir.join("settings.json"))?; + let hook_path = muse_hook_path(config_dir); + if settings.get("managed_hooks_path").and_then(Value::as_str) + != Some(hook_path.to_string_lossy().as_ref()) + { + bail!("Muse tracing hooks are not installed; run `bt trace enable muse`"); + } + if !hook_path.exists() { + bail!("Muse tracing hook file is missing; run `bt trace enable muse`"); + } + check_muse_hook_ownership(&hook_path)?; + write_object_atomic_unlocked(&hook_path, muse_hook_config()) + }) +} + +pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { + let mut warnings = Vec::new(); + match ProcessCommand::new("muse").arg("--version").output() { + Ok(output) if output.status.success() => match muse_version(&output.stdout) { + Ok(version) if version < semver::Version::new(1, 1, 1) => warnings + .push("Muse Code 1.1.1 or newer is required for verified command hooks".into()), + Ok(_) => {} + Err(error) => warnings.push(error.to_string()), + }, + _ => warnings.push( + "Muse Code executable is unavailable; install Muse Code before enabling tracing".into(), + ), + } + if !ProcessCommand::new("bt") + .arg("--version") + .output() + .is_ok_and(|output| output.status.success()) + { + warnings.push( + "Braintrust `bt` executable is unavailable on PATH; Muse command hooks cannot forward events".into(), + ); + } + let settings = match load_object(&config_dir.join("settings.json")) { + Ok(settings) => settings, + Err(error) => { + warnings.push(format!("Muse settings could not be read: {error}")); + return warnings; + } + }; 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()) + != Some(hook_path.to_string_lossy().as_ref()) { - settings.remove("managed_hooks_path"); - write_object_atomic(&settings_path, settings)?; + warnings.push("Muse managed_hooks_path does not point to Braintrust's hook file; run `bt trace enable muse`".into()); } - 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() - ) - }), + match std::fs::read(&hook_path) { + Ok(raw) => match serde_json::from_slice::(&raw) { + Ok(hooks) if hooks == Value::Object(muse_hook_config()) => {} + Ok(hooks) if hooks == Value::Object(legacy_muse_hook_config()) => warnings.push( + "Muse Braintrust hooks use an older command; run `bt trace update muse`".into(), + ), + Ok(_) => warnings.push("Muse Braintrust hook file differs from expected event definitions; restore it or run `bt trace enable muse`".into()), + Err(error) => warnings.push(format!("Muse Braintrust hook file is invalid JSON: {error}")), + }, + Err(error) => warnings.push(format!("Muse Braintrust hook file could not be read: {error}")), } + warnings } fn setup_muse() -> anyhow::Result<()> { + let output = ProcessCommand::new("muse") + .arg("--version") + .output() + .context("Muse Code is not installed; install Muse Code before enabling tracing")?; + if !output.status.success() { + bail!("Muse Code could not report its version; install Muse Code before enabling tracing"); + } + let installed = muse_version(&output.stdout)?; + if installed < semver::Version::new(1, 1, 1) { + bail!("Muse Code 1.1.1 or newer is required for verified command hooks"); + } setup_muse_at(&paths::muse_config_dir()) } fn disable_muse() -> anyhow::Result<()> { disable_muse_at(&paths::muse_config_dir()) } +fn update_muse() -> anyhow::Result<()> { + update_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!( @@ -987,7 +1152,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::Muse => update_muse()?, SetupAgent::Antigravity => update_antigravity(&mut runner)?, } Ok(TraceCommandOutput::update(source, display_name)) @@ -1734,6 +1899,11 @@ mod tests { 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()); + assert_eq!( + hooks["hooks"]["PreLLMCall"][0]["hooks"][0]["command"], + "bt trace hook --source muse --flush-on-turn-end" + ); + assert!(config_dir.join(".settings.json.lock").exists()); disable_muse_at(&config_dir).unwrap(); let settings: Value = @@ -1760,6 +1930,61 @@ mod tests { .contains("refusing to overwrite")); } + #[test] + fn muse_setup_and_disable_preserve_foreign_hook_file() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + std::fs::create_dir_all(&config_dir).unwrap(); + let hook_path = muse_hook_path(&config_dir); + let custom = r#"{"schema_version":1,"hooks":{"Stop":[]}}"#; + std::fs::write(&hook_path, custom).unwrap(); + + assert!(setup_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + assert!(disable_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + assert_eq!(std::fs::read_to_string(&hook_path).unwrap(), custom); + } + + #[test] + fn muse_update_requires_an_owned_install_and_upgrades_legacy_command() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + assert!(update_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("not installed")); + + setup_muse_at(&config_dir).unwrap(); + let hook_path = muse_hook_path(&config_dir); + write_object_atomic(&hook_path, legacy_muse_hook_config()).unwrap(); + assert!(muse_doctor_warnings(&config_dir) + .iter() + .any(|warning| warning.contains("older command"))); + update_muse_at(&config_dir).unwrap(); + let installed = load_object(&hook_path).unwrap(); + assert_eq!(installed, muse_hook_config()); + + std::fs::write(&hook_path, r#"{"hooks":{}}"#).unwrap(); + assert!(update_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + } + + #[test] + fn muse_version_requires_a_parseable_native_version() { + assert_eq!( + muse_version(b"Muse Code 1.1.1 (1.1.1-R2514.1)\n").unwrap(), + semver::Version::new(1, 1, 1) + ); + assert!(muse_version(b"unexpected output").is_err()); + } + #[test] fn disabling_removes_only_the_braintrust_settings_file() { let temp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 684c97b6..e637b193 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -532,6 +532,9 @@ async fn doctor_output_at( if let Some(warning) = crate::setup::update_warning(source) { warnings.push(warning); } + if args.agent == DoctorAgent::Muse { + warnings.extend(crate::setup::muse_doctor_warnings(&paths::muse_config_dir())); + } let (route, route_source) = match settings.route { Some(route) => (Some(route), "settings_file".to_string()), From 97bf39760994d48da1e8723e4f01394b4d2c80dc Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:00:14 -0400 Subject: [PATCH 07/17] Limit Muse hook setup to verified events --- bt-daemon/src/setup.rs | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index 5cfd4861..ae12b7d3 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -25,6 +25,14 @@ const OPENCODE_PACKAGE_MANIFEST: &str = include_str!("../../src/plugins/opencode/content/package.json"); const ANTIGRAVITY_PLUGIN: &str = "braintrust-antigravity-tracing"; const MUSE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreLLMCall", + "PostLLMCall", + "Stop", + "SessionEnd", +]; +const LEGACY_MUSE_HOOK_EVENTS: &[&str] = &[ "SessionStart", "UserPromptSubmit", "PreToolUse", @@ -706,8 +714,14 @@ fn muse_hook_path(config_dir: &Path) -> PathBuf { fn muse_hook_config() -> Map { // Muse treats command-hook failures as non-fatal to the agent turn. - let command = "bt trace hook --source muse --flush-on-turn-end"; - let hooks = MUSE_HOOK_EVENTS + muse_hook_config_for( + MUSE_HOOK_EVENTS, + "bt trace hook --source muse --flush-on-turn-end", + ) +} + +fn muse_hook_config_for(events: &[&str], command: &str) -> Map { + let hooks = events .iter() .map(|event| { ( @@ -726,21 +740,10 @@ fn muse_hook_config() -> Map { } fn legacy_muse_hook_config() -> Map { - let mut config = muse_hook_config(); - if let Some(hooks) = config.get_mut("hooks").and_then(Value::as_object_mut) { - for entries in hooks.values_mut() { - if let Some(command) = entries - .get_mut(0) - .and_then(|entry| entry.get_mut("hooks")) - .and_then(|hooks| hooks.get_mut(0)) - .and_then(|hook| hook.get_mut("command")) - { - *command = - Value::String("bt trace hook --source muse --flush-on-turn-end || true".into()); - } - } - } - config + muse_hook_config_for( + LEGACY_MUSE_HOOK_EVENTS, + "bt trace hook --source muse --flush-on-turn-end || true", + ) } fn muse_version(output: &[u8]) -> anyhow::Result { @@ -1899,6 +1902,8 @@ mod tests { 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()); + assert_eq!(hooks["hooks"].as_object().unwrap().len(), 6); + assert!(hooks["hooks"].get("PreToolUse").is_none()); assert_eq!( hooks["hooks"]["PreLLMCall"][0]["hooks"][0]["command"], "bt trace hook --source muse --flush-on-turn-end" From b86e2b68736c1c8d8cc1c27b4867fbba4af5a253 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:02:37 -0400 Subject: [PATCH 08/17] Stream and validate Muse exports --- bt-daemon/src/transcript_import/muse.rs | 788 ++++++++++++++++++------ bt-daemon/tests/replay.rs | 38 +- 2 files changed, 645 insertions(+), 181 deletions(-) diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index 0d4dd189..40328323 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -1,239 +1,637 @@ -//! 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. +//! Versioned Muse Code export-v1 reader. Retained frames can contain large +//! serialized payloads, so export events are decoded one record at a time. use super::{envelope, validate_session_id}; use crate::wire::Envelope; -use anyhow::{bail, Context}; +use anyhow::{anyhow, bail, Context, Result}; +use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor}; +use serde::Deserialize; use serde_json::{json, Value}; use std::collections::HashMap; +use std::fmt; +use std::fs::File; +use std::io::BufReader; 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() - ); +#[derive(Deserialize)] +struct Header { + export_schema_version: u64, + exporter_version: Option, + sessions: Vec, + #[serde(default)] + session_terminated_abnormally: bool, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum Version { + String(String), + Object { semver: String }, +} + +impl Version { + fn into_string(self) -> String { + match self { + Self::String(value) | Self::Object { semver: value } => value, + } } - 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; +} + +#[derive(Clone, Deserialize)] +struct Session { + session_id: String, + trajectory_id: Option, + root_session_id: Option, + session_end: Option, +} + +#[derive(Clone, Deserialize)] +struct SessionEnd { + exit_reason: Option, +} + +#[derive(Deserialize)] +struct ExportRecord { + recorded_at: Option, + envelope: Option, +} + +#[derive(Deserialize)] +struct NativeEnvelope { + recorded_at: Option, + payload: Option, +} + +#[derive(Deserialize)] +struct NativePayload { + kind: Option, + run_id: Option, + event: Option, +} + +#[derive(Default)] +struct RunState { + active_model: Option, + last_assistant: Option, + model_ordinal: u64, +} + +struct ModelState { + request_id: String, + usage: Option, + text: Option, +} + +struct Reader<'a, F> { + path: &'a Path, + session: Session, + session_id: String, + version: Option, + terminated_abnormally: bool, + emit: F, + runs: HashMap, + first_ts: Option, + last_ts: Option, + session_end_ts: Option, + activity_count: usize, +} + +impl Reader<'_, F> +where + F: FnMut(Envelope) -> Result<()>, +{ + fn emit(&mut self, name: &str, ts: i64, payload: Value) -> Result<()> { + (self.emit)(envelope( + "muse", + self.version.clone(), + &self.session_id, + name, + ts, + payload, + )) + } + + fn ensure_start(&mut self, ts: i64) -> Result<()> { + if self.first_ts.is_none() { + self.first_ts = Some(ts); + self.emit( + "SessionStart", + ts, + json!({ + "session_id": self.session_id, + "hook_event_name": "SessionStart", + "source": "muse_export_v1", + "trajectory_id": self.session.trajectory_id, + "root_session_id": self.session.root_session_id, + "transcript_path": self.path.to_string_lossy(), + }), + )?; + } + self.last_ts = Some(self.last_ts.unwrap_or(ts).max(ts)); + Ok(()) + } + + fn start_model( + &mut self, + run_id: &str, + ts: i64, + request_id: Option<&str>, + messages: Value, + ) -> Result<()> { + let run = self.runs.entry(run_id.to_owned()).or_default(); + if run.active_model.is_some() { + bail!("Muse run {run_id} starts another model request before the previous response") + } + let request_id = match request_id { + Some(id) if !id.is_empty() => id.to_owned(), + Some(_) => bail!("Muse model_input_trace_recorded has an empty request_record_id"), + None => { + run.model_ordinal += 1; + format!("{run_id}:model:{}", run.model_ordinal) + } }; - let Some(payload) = envelope.get("payload") else { - continue; + run.active_model = Some(ModelState { + request_id: request_id.clone(), + usage: None, + text: None, + }); + self.emit( + "PreLLMCall", + ts, + json!({ + "session_id": self.session_id, + "hook_event_name": "PreLLMCall", + "turn_id": run_id, + "request_id": request_id, + "messages": messages, + "source": "muse_export_v1", + }), + ) + } + + fn ensure_run(&mut self, run_id: &str, ts: i64) -> Result<()> { + if self.runs.contains_key(run_id) { + return Ok(()); + } + self.runs.insert(run_id.to_owned(), RunState::default()); + self.emit( + "UserPromptSubmit", + ts, + json!({ + "session_id": self.session_id, + "hook_event_name": "UserPromptSubmit", + "turn_id": run_id, + "prompt": Value::Null, + "source": "muse_export_v1", + }), + ) + } + + fn finish_model(&mut self, run_id: &str, ts: i64, error: Option<&str>) -> Result<()> { + let Some(model) = self + .runs + .get_mut(run_id) + .and_then(|run| run.active_model.take()) + else { + return Ok(()); }; - 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; + self.emit( + "PostLLMCall", + ts, + json!({ + "session_id": self.session_id, + "hook_event_name": "PostLLMCall", + "turn_id": run_id, + "request_id": model.request_id, + "usage": model.usage, + "output_text_preview": model.text, + "error": error, + "source": "muse_export_v1", + }), + ) + } + + fn process(&mut self, item: ExportRecord) -> Result<()> { + let Some(native_envelope) = item.envelope else { + return Ok(()); }; - let ts = recorded_ms(item).unwrap_or(last_ts); - last_ts = last_ts.max(ts); + let Some(payload) = native_envelope.payload else { + return Ok(()); + }; + let kind = payload + .event + .as_ref() + .and_then(|event| event.get("kind")) + .and_then(Value::as_str); + if payload.kind.as_deref() == Some("run") && payload.event.is_some() && kind.is_none() { + bail!("Muse run event has no kind"); + } + let recognized = payload.kind.as_deref() == Some("session_end") + || (payload.kind.as_deref() == Some("run") + && matches!( + kind, + Some( + "started" + | "model_input_trace_recorded" + | "model_completed" + | "resource_usage_sampled" + | "assistant_message_committed" + | "terminal" + ) + )); + if !recognized { + return Ok(()); + } + let micros = item + .recorded_at + .or(native_envelope.recorded_at) + .ok_or_else(|| anyhow!("Muse {kind:?} record has no recorded_at timestamp"))?; + let ts = micros / 1_000; + self.ensure_start(ts)?; + if payload.kind.as_deref() == Some("session_end") { + self.session_end_ts = Some(ts); + return Ok(()); + } + let event = payload + .event + .as_ref() + .ok_or_else(|| anyhow!("Muse run record has no event"))?; + let kind = kind.expect("recognized run event has kind"); + let run_id = payload + .run_id + .as_deref() + .filter(|id| !id.is_empty()) + .ok_or_else(|| anyhow!("Muse {kind} record has no run_id"))? + .to_owned(); + self.activity_count += 1; + if kind != "started" { + self.ensure_run(&run_id, 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, + "started" => { + if self.runs.contains_key(&run_id) { + bail!("Muse run {run_id} has duplicate started event") + } + self.runs.insert(run_id.clone(), RunState::default()); + self.emit( "UserPromptSubmit", ts, json!({ - "session_id": session_id, + "session_id": self.session_id, "hook_event_name": "UserPromptSubmit", "turn_id": run_id, - "prompt": native.get("prompt").cloned().unwrap_or(Value::Null), + "prompt": event.get("prompt"), "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", + let request_id = event + .get("request_record_id") + .and_then(Value::as_str) + .ok_or_else(|| { + anyhow!("Muse model_input_trace_recorded has no request_record_id") + })?; + if self + .runs + .get(&run_id) + .and_then(|run| run.active_model.as_ref()) + .is_some() + { + self.finish_model(&run_id, ts, None)?; + } + self.start_model( + &run_id, 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", - }), - )); + Some(request_id), + event.get("bounded").cloned().unwrap_or(Value::Null), + )?; } "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), - ), - ); + if self + .runs + .get(&run_id) + .and_then(|run| run.active_model.as_ref()) + .is_none() + { + self.start_model(&run_id, ts, None, Value::Null)?; + } + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|run| run.active_model.as_mut()) + { + model.usage = event.get("usage").cloned(); + } } - "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", - }), - )); + "resource_usage_sampled" => { + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|run| run.active_model.as_mut()) + { + if model.usage.is_none() { + model.usage = event.get("usage").cloned(); } } } - "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", - }))); + "assistant_message_committed" => { + let text = event + .get("text") + .cloned() + .ok_or_else(|| anyhow!("Muse assistant_message_committed has no text"))?; + if self + .runs + .get(&run_id) + .and_then(|run| run.active_model.as_ref()) + .is_none() + { + self.start_model(&run_id, ts, None, Value::Null)?; } - out.push(event( - version.clone(), - session_id, + let run = self.runs.get_mut(&run_id).expect("model start creates run"); + run.last_assistant = Some(text.clone()); + run.active_model.as_mut().expect("model exists").text = Some(text); + } + "terminal" => { + let terminal = event + .get("terminal") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("Muse terminal record has no terminal status"))?; + let error = (terminal != "completed").then(|| { + event + .get("reason") + .filter(|reason| !reason.is_null()) + .map(|reason| { + reason + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| reason.to_string()) + }) + .unwrap_or_else(|| terminal.to_owned()) + }); + self.finish_model(&run_id, ts, error.as_deref())?; + let run = self.runs.remove(&run_id); + self.emit( "Stop", ts, json!({ - "session_id": session_id, + "session_id": self.session_id, "hook_event_name": "Stop", "turn_id": run_id, - "last_assistant_message": assistant.remove(run_id).unwrap_or(Value::Null), + "last_assistant_message": run.and_then(|run| run.last_assistant), + "terminal": terminal, + "error": error, "source": "muse_export_v1", }), - )); + )?; } - _ => {} + _ => unreachable!(), } + Ok(()) } - out.push(event( - version, - session_id, - "SessionEnd", - last_ts, - json!({ - "session_id": session_id, + + fn finish(mut self) -> Result<()> { + if self.activity_count == 0 { + bail!( + "Muse export {} has no traceable run events", + self.path.display() + ); + } + let end_ts = self + .session_end_ts + .or(self.last_ts) + .ok_or_else(|| anyhow!("Muse export {} has no timestamps", self.path.display()))?; + for run_id in self.runs.keys().cloned().collect::>() { + self.finish_model(&run_id, end_ts, Some("Run ended without a terminal event"))?; + let run = self.runs.remove(&run_id).expect("run exists"); + self.emit( + "Stop", + end_ts, + json!({ + "session_id": self.session_id, + "hook_event_name": "Stop", + "turn_id": run_id, + "last_assistant_message": run.last_assistant, + "error": "Run ended without a terminal event", + "source": "muse_export_v1", + }), + )?; + } + self.emit("SessionEnd", end_ts, json!({ + "session_id": self.session_id, "hook_event_name": "SessionEnd", - "reason": session.pointer("/session_end/exit_reason"), + "reason": self.session.session_end.as_ref().and_then(|end| end.exit_reason.as_deref()), + "error": self.terminated_abnormally.then_some("Session terminated abnormally"), "source": "muse_export_v1", - }), + })) + } +} + +pub(crate) fn envelopes(path: &Path) -> Result> { + let header = read_header(path)?; + let session = single_session(&header, path)?; + envelopes_for_session(path, &session.session_id) +} + +pub(crate) fn envelopes_for_session(path: &Path, expected_id: &str) -> Result> { + let mut events = Vec::new(); + for_each_envelope_for_session(path, expected_id, |event| { + events.push(event); + Ok(()) + })?; + Ok(events) +} + +pub(crate) fn for_each_envelope_for_session( + path: &Path, + expected_id: &str, + emit: F, +) -> Result<()> +where + F: FnMut(Envelope) -> Result<()>, +{ + validate_session_id(expected_id)?; + let header = read_header(path)?; + if header.export_schema_version != 1 { + bail!( + "unsupported Muse export schema {} in {}; expected version 1", + header.export_schema_version, + path.display() + ); + } + let session = single_session(&header, path)?.clone(); + if session.session_id != expected_id { + bail!( + "Muse export {} contains session {}, expected {expected_id}", + path.display(), + session.session_id + ); + } + let mut reader = Reader { + path, + session, + session_id: expected_id.to_owned(), + version: header.exporter_version.map(Version::into_string), + terminated_abnormally: header.session_terminated_abnormally, + emit, + runs: HashMap::new(), + first_ts: None, + last_ts: None, + session_end_ts: None, + activity_count: 0, + }; + let mut de = serde_json::Deserializer::from_reader(BufReader::new( + File::open(path).with_context(|| format!("read Muse export {}", path.display()))?, )); - Ok(out) + EventsSeed(&mut reader) + .deserialize(&mut de) + .with_context(|| format!("parse Muse export {}", path.display()))?; + de.end() + .with_context(|| format!("parse Muse export {}", path.display()))?; + reader.finish() } -fn recorded_ms(item: &Value) -> Option { - item.get("recorded_at")? - .as_i64() - .map(|micros| micros / 1_000) +fn read_header(path: &Path) -> Result
{ + serde_json::from_reader(BufReader::new( + File::open(path).with_context(|| format!("read Muse export {}", path.display()))?, + )) + .with_context(|| format!("parse Muse export header {}", path.display())) } -fn event( - version: Option, - session_id: &str, - name: &str, - ts: i64, - payload: Value, -) -> Envelope { - envelope("muse", version, session_id, name, ts, payload) +fn single_session<'a>(header: &'a Header, path: &Path) -> Result<&'a Session> { + if header.sessions.len() != 1 { + bail!( + "Muse export {} must contain exactly one session, found {}", + path.display(), + header.sessions.len() + ); + } + let session = &header.sessions[0]; + validate_session_id(&session.session_id)?; + Ok(session) +} + +struct EventsSeed<'a, 'b, F>(&'a mut Reader<'b, F>); + +impl<'de, F> DeserializeSeed<'de> for EventsSeed<'_, '_, F> +where + F: FnMut(Envelope) -> Result<()>, +{ + type Value = (); + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_map(EventsVisitor(self.0)) + } +} + +struct EventsVisitor<'a, 'b, F>(&'a mut Reader<'b, F>); + +impl<'de, F> Visitor<'de> for EventsVisitor<'_, '_, F> +where + F: FnMut(Envelope) -> Result<()>, +{ + type Value = (); + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a Muse export object") + } + + fn visit_map(self, mut map: M) -> std::result::Result + where + M: MapAccess<'de>, + { + let mut found_events = false; + while let Some(key) = map.next_key::()? { + if key == "events" { + if found_events { + return Err(serde::de::Error::custom("duplicate Muse events array")); + } + found_events = true; + map.next_value_seed(RecordsSeed(self.0))?; + } else { + map.next_value::()?; + } + } + if !found_events { + return Err(serde::de::Error::custom("Muse export has no events array")); + } + Ok(()) + } +} + +struct RecordsSeed<'a, 'b, F>(&'a mut Reader<'b, F>); + +impl<'de, F> DeserializeSeed<'de> for RecordsSeed<'_, '_, F> +where + F: FnMut(Envelope) -> Result<()>, +{ + type Value = (); + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_seq(RecordsVisitor(self.0)) + } +} + +struct RecordsVisitor<'a, 'b, F>(&'a mut Reader<'b, F>); + +impl<'de, F> Visitor<'de> for RecordsVisitor<'_, '_, F> +where + F: FnMut(Envelope) -> Result<()>, +{ + type Value = (); + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("Muse event records") + } + + fn visit_seq(self, mut seq: S) -> std::result::Result + where + S: SeqAccess<'de>, + { + let mut index = 0; + while let Some(item) = seq.next_element::()? { + self.0.process(item).map_err(|error| { + serde::de::Error::custom(format!("event {}: {error:#}", index + 1)) + })?; + index += 1; + } + Ok(()) + } } #[cfg(test)] mod tests { use super::*; - #[test] - fn export_v1_becomes_shared_muse_lifecycle_events() { + + fn write_export(value: Value) -> (tempfile::TempDir, std::path::PathBuf) { 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", + std::fs::write(&path, serde_json::to_vec(&value).unwrap()).unwrap(); + (temp, path) + } + + #[test] + fn real_export_shapes_preserve_usage_terminal_and_session_end_time() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "exporter_version": {"semver":"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"}}}} + {"recorded_at": 1_004_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}, + {"recorded_at": 1_009_000, "envelope":{"payload":{"kind":"session_end"}}} ] - })).unwrap()).unwrap(); - let events = envelopes(&path).unwrap(); + })); + let events = envelopes_for_session(&path, "session-1").unwrap(); assert_eq!( events .iter() @@ -248,7 +646,41 @@ mod tests { "SessionEnd" ] ); - assert_eq!(events[2].ts_ms, 1001); - assert_eq!(events[4].payload["last_assistant_message"], "world"); + assert_eq!(events[3].payload["usage"]["input_tokens"], 1); + assert_eq!(events[3].payload["output_text_preview"], "world"); + assert_eq!(events[5].ts_ms, 1009); + } + + #[test] + fn sparse_export_has_deterministic_model_and_failure() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "sessions": [{"session_id":"session-1"}], + "events": [ + {"recorded_at": 1_001_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"partial"}}}}, + {"recorded_at": 1_002_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"resource_usage_sampled","usage":{"output_tokens":2}}}}}, + {"recorded_at": 1_003_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"failed","reason":"provider failed"}}}} + ] + })); + let events = envelopes(&path).unwrap(); + assert_eq!(events[2].payload["request_id"], "run-1:model:1"); + assert_eq!(events[3].payload["usage"]["output_tokens"], 2); + assert_eq!(events[3].payload["error"], "provider failed"); + assert_eq!(events[4].payload["error"], "provider failed"); + } + + #[test] + fn identity_and_malformed_recognized_records_fail_with_context() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "sessions": [{"session_id":"session-1"}], + "events": [{"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","event":{"kind":"started"}}}}] + })); + assert!(envelopes_for_session(&path, "session-2") + .unwrap_err() + .to_string() + .contains("expected session-2")); + let error = envelopes(&path).unwrap_err(); + assert!(format!("{error:#}").contains("event 1: Muse started record has no run_id")); } } diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 10226f05..9d94ba03 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -38,7 +38,7 @@ async fn importing_muse_export_uses_the_production_translator() { {"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"}}}} + {"recorded_at": 1_004_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"terminal","terminal":"completed"}}}} ] })).unwrap()).unwrap(); let output = tmp.path().join("spans"); @@ -49,14 +49,46 @@ async fn importing_muse_export_uses_the_production_translator() { assert_eq!(inserted(&output_rows, "llm"), 1); assert!(output_rows .iter() - .any(|row| { row.pointer("/Insert/output").and_then(Value::as_str) == Some("traced") })); + .any(|row| { row.pointer("/Merge/output").and_then(Value::as_str) == Some("traced") })); assert!(output_rows.iter().any(|row| { - row.pointer("/Insert/metrics/input_tokens") + row.pointer("/Merge/metrics/prompt_tokens") .and_then(Value::as_i64) == Some(3) })); } +#[tokio::test] +async fn importing_sparse_failed_muse_export_preserves_model_and_error() { + let tmp = tempfile::tempdir().unwrap(); + let export = tmp.path().join("muse-sparse-export.json"); + std::fs::write(&export, serde_json::to_vec(&json!({ + "export_schema_version": 1, + "exporter_version": {"semver":"1.1.1"}, + "sessions": [{"session_id":"muse-sparse", "session_end":{"exit_reason":"error"}}], + "events": [ + {"recorded_at": 1_000_000, "envelope":{"payload":{"kind":"run", "run_id":"run-1", "event":{"kind":"assistant_message_committed", "text":"partial"}}}}, + {"recorded_at": 1_001_000, "envelope":{"payload":{"kind":"run", "run_id":"run-1", "event":{"kind":"resource_usage_sampled", "usage":{"input_tokens":3, "output_tokens":2}}}}}, + {"recorded_at": 1_002_000, "envelope":{"payload":{"kind":"run", "run_id":"run-1", "event":{"kind":"terminal", "terminal":"failed", "reason":"provider failed"}}}}, + {"recorded_at": 1_009_000, "envelope":{"payload":{"kind":"session_end"}}} + ] + })).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-sparse.ndjson")); + assert_eq!(inserted(&output_rows, "llm"), 1); + assert!(output_rows.iter().any(|row| { + row.pointer("/Merge/metrics/prompt_tokens") + .and_then(Value::as_u64) + == Some(3) + && row.pointer("/Merge/error").and_then(Value::as_str) == Some("provider failed") + })); + assert!(output_rows + .iter() + .any(|row| { row.pointer("/Merge/end_ms").and_then(Value::as_i64) == Some(1009) })); +} + #[tokio::test] async fn importing_large_codex_rollout_drains_translator_continuations() { const CALLS: usize = 32; From 28a24d6490a40da872b3eda98c59476b2e581d71 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:03:47 -0400 Subject: [PATCH 09/17] Ground Muse translation in captured hook events --- bt-daemon/src/translate/muse.rs | 620 ++++++++++++------ .../tests/fixtures/muse-hooks-1.1.1.jsonl | 6 + 2 files changed, 441 insertions(+), 185 deletions(-) create mode 100644 bt-daemon/tests/fixtures/muse-hooks-1.1.1.jsonl diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs index 4e69f1ef..97284aec 100644 --- a/bt-daemon/src/translate/muse.rs +++ b/bt-daemon/src/translate/muse.rs @@ -1,18 +1,19 @@ //! 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. +//! Muse hooks remain a thin transport. The session, turn and model payloads +//! used here were captured from Muse Code 1.1.1-R2514.1. Import feeds the +//! same event names through this translator. use super::git::GitMetadataCache; use super::{ - local_username, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, + local_username, root_tags, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, + TranslatorFactory, }; use crate::ids; use crate::wire::Envelope; -use serde_json::{json, Value}; -use std::collections::HashMap; +use serde::{Deserialize, Deserializer}; +use serde_json::{json, Map, Value}; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; pub struct MuseTranslatorFactory { @@ -36,10 +37,33 @@ impl TranslatorFactory for MuseTranslatorFactory { struct OpenOp { id: String, parents: Vec, - name: String, - start_ms: i64, - input: Value, - kind: SpanType, +} + +#[derive(Default, Deserialize)] +struct HookIds { + #[serde(default, deserialize_with = "string_or_none")] + turn_id: Option, + #[serde(default, deserialize_with = "string_or_none")] + request_id: Option, +} + +fn string_or_none<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + Ok(Value::deserialize(deserializer)? + .as_str() + .filter(|value| !value.is_empty()) + .map(str::to_owned)) +} + +impl HookIds { + fn turn(&self) -> Option<&str> { + self.turn_id.as_deref().filter(|id| !id.is_empty()) + } + fn request(&self) -> Option<&str> { + self.request_id.as_deref().filter(|id| !id.is_empty()) + } } struct MuseTranslator { @@ -48,10 +72,12 @@ struct MuseTranslator { effective_root_id: String, external_parent: Option, opened: bool, + session_metadata: Map, + session_version: Option, turns: HashMap, llms: HashMap, - tools: HashMap, - special: HashMap, + closed_turns: HashSet, + closed_llms: HashSet, last_ts: i64, git: Arc, } @@ -64,14 +90,37 @@ impl MuseTranslator { effective_root_id: String::new(), external_parent: None, opened: false, + session_metadata: Map::new(), + session_version: None, turns: HashMap::new(), llms: HashMap::new(), - tools: HashMap::new(), - special: HashMap::new(), + closed_turns: HashSet::new(), + closed_llms: HashSet::new(), last_ts: 0, git, } } + fn remember_session(&mut self, e: &Envelope) { + if let Some(version) = &e.source_version { + self.session_version = Some(version.clone()); + } + for key in [ + "cwd", + "model", + "permission_mode", + "trajectory_id", + "root_session_id", + ] { + if let Some(value) = e.payload.get(key) { + self.session_metadata.insert(key.into(), value.clone()); + } + } + } + + fn starts_trace(event: &str) -> bool { + matches!(event, "UserPromptSubmit" | "PreLLMCall") + } + fn ensure_root(&mut self, e: &Envelope, ctx: &SessionCtx) -> Vec { if self.opened { return Vec::new(); @@ -97,17 +146,14 @@ impl MuseTranslator { 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()); - } + if let Some(version) = &self.session_version { + metadata.insert("muse_version".into(), json!(version)); + } + if let Some(version) = &e.plugin_version { + metadata.insert("adapter_version".into(), json!(version)); + } + for (key, value) in &self.session_metadata { + metadata.insert(key.clone(), value.clone()); } vec![SpanOp::Insert(SpanRow { span_id: self.root_id.clone(), @@ -117,6 +163,7 @@ impl MuseTranslator { span_type: SpanType::Task, start_ms: Some(e.ts_ms), metadata: Some(Value::Object(metadata)), + tags: root_tags(ctx), ..Default::default() })] } @@ -132,15 +179,11 @@ impl MuseTranslator { let Some(open) = map.remove(key) else { return Vec::new(); }; - vec![SpanOp::Insert(SpanRow { + vec![SpanOp::Merge(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, @@ -149,22 +192,17 @@ impl MuseTranslator { } 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 (map, closed) in [ + (&mut self.llms, &mut self.closed_llms), + (&mut self.turns, &mut self.closed_turns), ] { - for (_, open) in std::mem::take(map) { - out.push(SpanOp::Insert(SpanRow { + for (key, open) in std::mem::take(map) { + closed.insert(key); + out.push(SpanOp::Merge(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() })); @@ -172,163 +210,183 @@ impl MuseTranslator { } out } + + fn parent(&self, ids: &HookIds) -> String { + ids.turn() + .and_then(|id| self.turns.get(id)) + .map(|op| op.id.clone()) + .unwrap_or_else(|| self.root_id.clone()) + } + + fn insert_open(map: &mut HashMap, key: &str, row: SpanRow) -> Vec { + if map.contains_key(key) { + return Vec::new(); + } + map.insert( + key.into(), + OpenOp { + id: row.span_id.clone(), + parents: row.parent_span_ids.clone(), + }, + ); + vec![SpanOp::Insert(row)] + } + + fn unpaired(&self, e: &Envelope, reason: &str) -> SpanOp { + // A malformed event must not share an `unknown` ID with another event. + let identity = format!("unpaired:{}:{}:{}", e.event, e.ts_ms, e.payload); + SpanOp::Insert(SpanRow { + span_id: ids::span_id(&self.session_id, &identity), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: vec![self.root_id.clone()], + name: e.event.clone(), + span_type: SpanType::Task, + start_ms: Some(e.ts_ms), + end_ms: Some(e.ts_ms), + metadata: Some(json!({"unpaired": true})), + error: Some(reason.into()), + ..Default::default() + }) + } } 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); + self.remember_session(e); + let ids = match serde_json::from_value::(e.payload.clone()) { + Ok(ids) => ids, + Err(error) => { + // Unknown or malformed optional fields must not terminate a session. + tracing::warn!(source = "muse", event = %e.event, %error, "invalid Muse hook identifiers"); + HookIds::default() + } + }; + if (e.event == "UserPromptSubmit" + && ids.turn().is_some_and(|id| self.closed_turns.contains(id))) + || (e.event == "PreLLMCall" + && ids + .request() + .is_some_and(|id| self.closed_llms.contains(id))) + { + return Ok(Vec::new()); + } + let mut out = if Self::starts_trace(&e.event) { + self.ensure_root(e, ctx) + } else { + Vec::new() + }; 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()], + if let Some(id) = ids.turn() { + out.extend(Self::insert_open( + &mut self.turns, + id, + SpanRow { + span_id: ids::span_id(&self.session_id, &format!("turn:{id}")), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: 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, + span_type: SpanType::Task, + start_ms: Some(e.ts_ms), + input: p.get("prompt").cloned(), + ..Default::default() }, - ); + )); + } else { + out.push(self.unpaired(e, "Muse UserPromptSubmit has no turn_id")); } } "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 }); + if let Some(id) = ids.request() { + let mut metadata = Map::new(); + for key in [ + "model", + "provider", + "step", + "attempt", + "message_count", + "tool_count", + ] { + if let Some(value) = p.get(key) { + metadata.insert(key.into(), value.clone()); + } + } + let input = json!({"messages": p.get("messages"), "tools": p.get("tools")}); + let parent = self.parent(&ids); + out.extend(Self::insert_open( + &mut self.llms, + id, + SpanRow { + span_id: ids::span_id(&self.session_id, &format!("llm:{id}")), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: vec![parent], + name: p + .get("model") + .and_then(Value::as_str) + .filter(|s| !s.is_empty() && *s != "unknown") + .unwrap_or("llm") + .into(), + span_type: SpanType::Llm, + start_ms: Some(e.ts_ms), + input: Some(input), + metadata: Some(Value::Object(metadata)), + ..Default::default() + }, + )); + } else { + out.push(self.unpaired(e, "Muse PreLLMCall has no request_id")); + } } "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, - )); + if let Some(id) = ids.request() { + let error = event_error(p); + let closed = Self::close( + &self.effective_root_id, + &mut self.llms, + id, + e.ts_ms, + p.get("output_text_preview").cloned(), + p.get("usage").map(usage_metrics), + error, + ); + if !closed.is_empty() { + self.closed_llms.insert(id.into()); + out.extend(closed); + } + } else if self.opened { + out.push(self.unpaired(e, "Muse PostLLMCall has no request_id")); + } } "Stop" => { - if let Some(id) = turn_id { - out.extend(Self::close( + if let Some(id) = ids.turn() { + let closed = Self::close( &self.effective_root_id, &mut self.turns, id, e.ts_ms, p.get("last_assistant_message").cloned(), None, - None, - )); + event_error(p), + ); + if !closed.is_empty() { + self.closed_turns.insert(id.into()); + out.extend(closed); + } + } else if self.opened { + out.push(self.unpaired(e, "Muse Stop has no turn_id")); } } "SessionEnd" => { - out.extend(self.close_all(e.ts_ms, None)); + out.extend(self.close_all(e.ts_ms, Some("Interrupted before completion"))); if self.opened { - out.push(SpanOp::Insert(SpanRow { + out.push(SpanOp::Merge(SpanRow { span_id: self.root_id.clone(), root_span_id: self.effective_root_id.clone(), end_ms: Some(e.ts_ms), + metadata: p.get("reason").map(|reason| json!({"end_reason":reason})), + error: event_error(p), ..Default::default() })); self.opened = false; @@ -340,10 +398,11 @@ impl AgentTranslator for MuseTranslator { .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 { + out.push(SpanOp::Merge(SpanRow { span_id: self.root_id.clone(), root_span_id: self.effective_root_id.clone(), end_ms: Some(self.last_ts), @@ -356,19 +415,58 @@ impl AgentTranslator for MuseTranslator { } } +fn event_error(payload: &Value) -> Option { + payload + .get("error") + .and_then(|error| { + error + .as_str() + .map(str::to_owned) + .or_else(|| (!error.is_null()).then(|| error.to_string())) + }) + .or_else(|| { + payload + .get("status") + .and_then(Value::as_str) + .filter(|status| !matches!(*status, "success" | "completed")) + .map(|status| format!("Muse operation {status}")) + }) +} + +fn usage_metrics(usage: &Value) -> Value { + let mut metrics = Map::new(); + for (native, standard) in [ + ("input_tokens", "prompt_tokens"), + ("output_tokens", "completion_tokens"), + ("cached_tokens", "prompt_cached_tokens"), + ("reasoning_tokens", "completion_reasoning_tokens"), + ] { + if let Some(value) = usage.get(native).and_then(Value::as_u64) { + metrics.insert(standard.into(), json!(value)); + } + } + if let (Some(input), Some(output)) = ( + usage.get("input_tokens").and_then(Value::as_u64), + usage.get("output_tokens").and_then(Value::as_u64), + ) { + metrics.insert("tokens".into(), json!(input.saturating_add(output))); + } + Value::Object(metrics) +} + #[cfg(test)] mod tests { use super::*; use crate::translate::Registry; - use crate::wire::Envelope; - fn event(name: &str, payload: Value) -> Envelope { + + fn event(name: &str, payload: Value, ts_ms: i64) -> Envelope { Envelope { source: "muse".into(), - source_version: Some("1.1.1".into()), + source_version: Some("1.1.1-R2514.1".into()), plugin_version: None, - session_id: "session-1".into(), + session_id: "muse-fixture-session".into(), event: name.into(), - ts_ms: 1000, + ts_ms, managed_run_id: None, capture: None, payload, @@ -376,26 +474,178 @@ mod tests { config: None, } } + + fn context() -> SessionCtx { + SessionCtx { + session_id: "muse-fixture-session".into(), + config: None, + } + } + + fn fixture() -> Vec { + include_str!("../../tests/fixtures/muse-hooks-1.1.1.jsonl") + .lines() + .enumerate() + .map(|(index, line)| { + let payload: Value = serde_json::from_str(line).unwrap(); + let name = payload["hook_event_name"].as_str().unwrap(); + event(name, payload.clone(), 1000 + index as i64) + }) + .collect() + } + #[test] - fn registry_and_llm_pairing_are_stable() { + fn captured_muse_1_1_1_lifecycle_has_visible_starts_and_merges() { 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, - }; + let mut translator = registry.create("muse", "muse-fixture-session"); + let ctx = context(); + let events = fixture(); + assert!(translator.handle(&events[0], &ctx).unwrap().is_empty()); + + let turn_start = translator.handle(&events[1], &ctx).unwrap(); + let root_id = turn_start + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.name == "Muse Code" => Some(row.span_id.clone()), + _ => None, + }) + .expect("session root"); + assert!(turn_start.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.name == "Muse Code" && row.metadata.as_ref().and_then(|m| m.get("muse_version")) == Some(&json!("1.1.1-R2514.1"))))); + assert!(turn_start.iter().any( + |op| matches!(op, SpanOp::Insert(row) if row.name == "Turn" && row.end_ms.is_none()) + )); + let llm_start = translator.handle(&events[2], &ctx).unwrap(); + let llm_id = llm_start + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.span_type == SpanType::Llm => Some(row.span_id.clone()), + _ => None, + }) + .expect("LLM start span"); + let llm_end = translator.handle(&events[3], &ctx).unwrap(); + assert!(llm_end.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == llm_id && row.output == Some(json!("echo: hello hook probe")) && row.metrics.as_ref().and_then(|m| m.get("prompt_tokens")) == Some(&json!(0))))); + let turn_end = translator.handle(&events[4], &ctx).unwrap(); + assert!(turn_end.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.output == Some(json!("echo: hello hook probe"))))); + assert!(translator.handle(&events[1], &ctx).unwrap().is_empty()); + assert!(translator.handle(&events[2], &ctx).unwrap().is_empty()); + assert!(translator.handle(&events[3], &ctx).unwrap().is_empty()); + let session_end = translator.handle(&events[5], &ctx).unwrap(); + assert!(session_end.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == root_id && row.end_ms == Some(1005)))); + assert!(translator.finalize(&ctx).unwrap().is_empty()); + } + + #[test] + fn passive_and_unknown_hooks_do_not_create_a_trace() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + for name in ["SessionStart", "SessionEnd", "UnrecognizedEvent", "Stop"] { + assert!(translator + .handle(&event(name, json!({}), 1000), &context()) + .unwrap() + .is_empty()); + } + assert!(translator.finalize(&context()).unwrap().is_empty()); + } + + #[test] + fn missing_ids_cannot_pair_distinct_operations_and_failure_is_visible() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + let ctx = context(); + let first = translator + .handle( + &event( + "PreLLMCall", + json!({"messages":["a"],"request_id":42}), + 1000, + ), + &ctx, + ) + .unwrap(); + let second = translator + .handle(&event("PreLLMCall", json!({"messages":["b"]}), 1001), &ctx) + .unwrap(); + let first_id = first + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.error.is_some() => Some(row.span_id.clone()), + _ => None, + }) + .unwrap(); + let second_id = second + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.error.is_some() => Some(row.span_id.clone()), + _ => None, + }) + .unwrap(); + assert_ne!(first_id, second_id); + assert!(translator.llms.is_empty()); + + let start = event( + "PreLLMCall", + json!({"request_id":"r1","turn_id":"t1"}), + 1002, + ); + translator.handle(&start, &ctx).unwrap(); + let completed = translator.handle(&event("PostLLMCall", json!({"request_id":"r1","status":"failed","usage":{"input_tokens":3,"output_tokens":1,"cached_tokens":2,"reasoning_tokens":1}}), 1003), &ctx).unwrap(); + assert!(completed.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.error.as_deref() == Some("Muse operation failed") && row.metrics.as_ref().and_then(|m| m.get("tokens")) == Some(&json!(4)) && row.metrics.as_ref().and_then(|m| m.get("prompt_cached_tokens")) == Some(&json!(2))))); + } + + #[test] + fn unsupported_hook_is_ignored_and_pending_work_is_interrupted() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + let ctx = context(); translator .handle( &event( "UserPromptSubmit", - json!({"session_id":"session-1","turn_id":"turn-1","prompt":"hello"}), + json!({"turn_id":"t1","prompt":"x"}), + 1000, + ), + &ctx, + ) + .unwrap(); + let unsupported = translator + .handle( + &event( + "PermissionRequest", + json!({"turn_id":"t1","request_id":"p1"}), + 1001, + ), + &ctx, + ) + .unwrap(); + assert!(unsupported.is_empty()); + translator + .handle( + &event( + "PreLLMCall", + json!({"turn_id":"t1","request_id":"r1"}), + 1002, + ), + &ctx, + ) + .unwrap(); + let ended = translator + .handle( + &event( + "SessionEnd", + json!({"reason":"other","error":"Session terminated abnormally"}), + 1003, ), &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"))))); + assert!(ended.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.error.as_deref() == Some("Interrupted before completion")))); + assert!(ended.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.error.as_deref() == Some("Session terminated abnormally")))); } } diff --git a/bt-daemon/tests/fixtures/muse-hooks-1.1.1.jsonl b/bt-daemon/tests/fixtures/muse-hooks-1.1.1.jsonl new file mode 100644 index 00000000..cd647f6e --- /dev/null +++ b/bt-daemon/tests/fixtures/muse-hooks-1.1.1.jsonl @@ -0,0 +1,6 @@ +{"hook_event_name":"SessionStart","session_id":"muse-fixture-session","cwd":"/workspace","model":"unknown","permission_mode":"on-request","source":"startup","transcript_path":"/tmp/fixture.jsonl"} +{"hook_event_name":"UserPromptSubmit","session_id":"muse-fixture-session","turn_id":"muse-fixture-turn","cwd":"/workspace","model":"unknown","permission_mode":"on-request","prompt":"hello hook probe","transcript_path":"/tmp/fixture.jsonl"} +{"hook_event_name":"PreLLMCall","session_id":"muse-fixture-session","turn_id":"muse-fixture-turn","request_id":"muse-fixture-turn:0:1","attempt":1,"step":0,"cwd":"/workspace","model":"unknown","provider":"model.unknown.response","permission_mode":"on-request","message_count":1,"tool_count":0,"messages":[{"role":"user","content":"hello hook probe"}],"tools":[],"options":{},"transcript_path":"/tmp/fixture.jsonl"} +{"hook_event_name":"PostLLMCall","session_id":"muse-fixture-session","turn_id":"muse-fixture-turn","request_id":"muse-fixture-turn:0:1","attempt":1,"step":0,"cwd":"/workspace","model":"unknown","provider":"model.unknown.response","permission_mode":"on-request","message_count":1,"tool_count":0,"tool_call_count":0,"messages":[{"role":"user","content":"hello hook probe"}],"tools":[],"options":{},"status":"success","error":null,"finish_reason":null,"response_id":null,"output_text_preview":"echo: hello hook probe","usage":{"input_tokens":0,"output_tokens":0,"cached_tokens":0,"reasoning_tokens":0},"transcript_path":"/tmp/fixture.jsonl"} +{"hook_event_name":"Stop","session_id":"muse-fixture-session","turn_id":"muse-fixture-turn","cwd":"/workspace","model":"unknown","permission_mode":"on-request","last_assistant_message":"echo: hello hook probe","stop_hook_active":false,"transcript_path":"/tmp/fixture.jsonl"} +{"hook_event_name":"SessionEnd","session_id":"muse-fixture-session","cwd":"/workspace","model":"unknown","permission_mode":"on-request","reason":"other","transcript_path":"/tmp/fixture.jsonl"} From c361da6539402ab0c2ca85374eff57c3314d456e Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:04:39 -0400 Subject: [PATCH 10/17] Bound Muse MSP and export import processing --- bt-daemon/README.md | 7 +- bt-daemon/src/lib.rs | 411 +++++++++++++++++++++++++++++++------------ 2 files changed, 302 insertions(+), 116 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index b1387ff4..93d8853d 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -284,7 +284,12 @@ 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 +completed session. Exports are parsed incrementally, one session at a time. +Muse Code 1.1.1 setup captures the six verified lifecycle and model hooks: +`SessionStart`, `UserPromptSubmit`, `PreLLMCall`, `PostLLMCall`, `Stop`, and +`SessionEnd`. Tool, permission, subagent, and compaction details are not +available from those hooks and are not inferred by the live translator. +`--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. diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index f361f6fd..2f1aa213 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -61,12 +61,15 @@ pub use translate::{ use anyhow::Context; use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; use clap::{Args, ValueEnum}; +use serde::Deserialize; use std::ffi::OsString; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines}; +use tokio::io::{ + AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines, +}; use wire::{ method, Envelope, ManagedRunFlushParams, SessionConfig, SessionRoute, StatusResult, PROTOCOL_VERSION, @@ -75,6 +78,9 @@ 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); +const MUSE_MSP_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const MUSE_MSP_MAX_PAGES: usize = 1_000; +const MUSE_MSP_DIAGNOSTIC_BYTES: usize = 64 * 1024; /// Arguments for `serve`. #[derive(Debug, Clone, Args)] @@ -788,9 +794,16 @@ pub async fn run_import( "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, + 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") + } + return import_muse_sessions_with_ledger( + &session_ids, opts, config, Some(paths::data_dir(None)), @@ -839,56 +852,44 @@ fn apply_import_span_plugins( Ok(()) } -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()?; +async fn export_muse_session(session_id: &str, path: &std::path::Path) -> anyhow::Result<()> { + transcript_import::validate_session_id(session_id)?; 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() - ); - } + let mut child = tokio::process::Command::new(&executable) + .kill_on_drop(true) + .args(["export", "--session", session_id, "--out"]) + .arg(path) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| { + format!( + "run {} export for Muse session {session_id}", + executable.to_string_lossy() + ) + })?; + let stderr = child.stderr.take().expect("piped stderr"); + let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); + let status = tokio::time::timeout(Duration::from_secs(300), child.wait()) + .await + .with_context(|| format!("Muse export for session {session_id} timed out"))? + .with_context(|| format!("wait for Muse export of session {session_id}"))?; + let stderr = diagnostics + .await + .context("collect Muse export diagnostics")??; + if !status.success() { + anyhow::bail!( + "Muse export for session {session_id} failed: {}", + stderr.trim() + ); } - Ok(directory) + Ok(()) } /// 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) @@ -905,73 +906,133 @@ async fn muse_msp_session_ids() -> anyhow::Result> { })?; let mut stdin = child.stdin.take().expect("piped stdin"); let stdout = child.stdout.take().expect("piped stdout"); + let stderr = child.stderr.take().expect("piped stderr"); + let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); 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( + let listing = async { + let mut request_id = 1u64; + muse_msp_request( &mut stdin, &mut reader, request_id, - "session/list", - serde_json::json!({"cursor": cursor, "limit": 200}), + "initialize", + serde_json::json!({"clientInfo":{"name":"braintrust_muse_import","version":"1"}}), ) .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; + // Muse requires the initialized notification before other requests. + muse_msp_notify(&mut stdin, "initialized").await?; + let mut cursor = None; + let mut seen_cursors = std::collections::HashSet::new(); + let mut sessions = Vec::new(); + for page_index in 0..MUSE_MSP_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?; + match muse_msp_append_page(result, page_index + 1, &mut sessions, &mut seen_cursors)? { + None => { + sessions.sort(); + sessions.dedup(); + return Ok(sessions); + } + Some(next) => cursor = Some(next), + } } + anyhow::bail!("Muse session/list exceeded {MUSE_MSP_MAX_PAGES} pages") } + .await; 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() - ); + drop(reader); + let status = match tokio::time::timeout(MUSE_MSP_SHUTDOWN_TIMEOUT, child.wait()).await { + Ok(status) => status.context("wait for Muse MSP host")?, + Err(_) => { + child + .start_kill() + .context("stop unresponsive Muse MSP host")?; + child + .wait() + .await + .context("reap unresponsive Muse MSP host")?; + anyhow::bail!("Muse MSP host did not exit after stdin closed") + } + }; + let stderr = diagnostics + .await + .context("collect Muse MSP diagnostics")??; + let sessions = listing?; + if !status.success() { + anyhow::bail!("Muse MSP session listing failed: {}", stderr.trim()); } - sessions.sort(); - sessions.dedup(); Ok(sessions) } +#[derive(Deserialize)] +struct MuseMspSessionPage { + sessions: Vec, + #[serde(rename = "nextCursor")] + next_cursor: Option, +} + +#[derive(Deserialize)] +struct MuseMspSessionEntry { + #[serde(rename = "sessionId", alias = "session_id")] + session_id: String, +} + +fn muse_msp_append_page( + result: serde_json::Value, + page_index: usize, + sessions: &mut Vec, + seen_cursors: &mut std::collections::HashSet, +) -> anyhow::Result> { + let page: MuseMspSessionPage = serde_json::from_value(result) + .with_context(|| format!("decode Muse session/list page {page_index}"))?; + for session in page.sessions { + transcript_import::validate_session_id(&session.session_id).with_context(|| { + format!("invalid Muse session id on session/list page {page_index}") + })?; + sessions.push(session.session_id); + } + if let Some(cursor) = &page.next_cursor { + if !seen_cursors.insert(cursor.clone()) { + anyhow::bail!("Muse session/list repeated cursor on page {page_index}") + } + } + Ok(page.next_cursor) +} + +async fn muse_drain_stderr(mut stderr: tokio::process::ChildStderr) -> anyhow::Result { + let mut output = Vec::new(); + let mut buffer = [0u8; 4096]; + loop { + let count = stderr.read(&mut buffer).await?; + if count == 0 { + break; + } + let remaining = MUSE_MSP_DIAGNOSTIC_BYTES.saturating_sub(output.len()); + output.extend_from_slice(&buffer[..count.min(remaining)]); + } + Ok(String::from_utf8_lossy(&output).into_owned()) +} + 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(()) + tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { + stdin.write_all(frame.as_bytes()).await?; + stdin.flush().await?; + Ok(()) + }) + .await + .map_err(|_| anyhow::anyhow!("Muse MSP timed out sending {method}"))? } async fn muse_msp_request( @@ -989,9 +1050,9 @@ where &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 { + stdin.write_all(request.as_bytes()).await?; + stdin.flush().await?; loop { let line = reader .next_line() @@ -999,8 +1060,17 @@ where .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) { + if frame.get("jsonrpc").and_then(serde_json::Value::as_str) != Some("2.0") { + anyhow::bail!( + "Muse MSP sent a frame without JSON-RPC 2.0 while waiting for {method}" + ) + } + let Some(response_id) = frame.get("id") else { + // Notifications can be interleaved with responses. continue; + }; + if response_id.as_u64() != Some(id) { + anyhow::bail!("Muse MSP returned an unexpected response id for {method}") } if let Some(error) = frame.get("error") { anyhow::bail!("Muse MSP {method} failed: {error}") @@ -1015,31 +1085,62 @@ where .map_err(|_| anyhow::anyhow!("Muse MSP timed out waiting for {method}"))? } -async fn import_muse_exports_with_ledger( - directory: &tempfile::TempDir, +async fn import_muse_sessions_with_ledger( + session_ids: &[String], 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); + for session_id in session_ids { + let directory = tempfile::Builder::new() + .prefix("bt-muse-import-") + .tempdir()?; + let export = directory.path().join("session.json"); + export_muse_session(session_id, &export).await?; + let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); + let export_for_parser = export.clone(); + let expected_id = session_id.clone(); + let parser = tokio::task::spawn_blocking(move || { + let mut batch = Vec::with_capacity(128); + transcript_import::muse::for_each_envelope_for_session( + &export_for_parser, + &expected_id, + |entry| { + batch.push(entry); + if batch.len() == 128 { + sender + .blocking_send(std::mem::take(&mut batch)) + .map_err(|_| { + anyhow::anyhow!("Muse import stopped while parsing {expected_id}") + })?; + } + Ok(()) + }, + )?; + if !batch.is_empty() { + sender.blocking_send(batch).map_err(|_| { + anyhow::anyhow!("Muse import stopped while parsing {expected_id}") + })?; } + Ok::<(), anyhow::Error>(()) + }); + let mut processing_error = None; + while let Some(batch) = receiver.recv().await { + if let Err(error) = processor.process(batch).await { + processing_error = Some(error); + break; + } + } + drop(receiver); + let parse_result = parser.await.context("Muse export parser task failed")?; + if let Some(error) = processing_error { + return Err(error); + } + parse_result?; + if let Some(summary) = processor.finish_session(session_id).await? { + summaries.push(summary); } } summaries.extend(processor.finish().await?); @@ -2114,6 +2215,86 @@ mod tests { server.await.unwrap(); } + #[test] + fn muse_msp_pages_reject_partial_or_cyclic_session_lists() { + let mut sessions = Vec::new(); + let mut cursors = std::collections::HashSet::new(); + assert_eq!( + muse_msp_append_page( + json!({"sessions":[{"sessionId":"session-a"}],"nextCursor":"next"}), + 1, + &mut sessions, + &mut cursors, + ) + .unwrap(), + Some("next".to_owned()) + ); + assert_eq!(sessions, ["session-a"]); + assert!(muse_msp_append_page( + json!({"sessions":[],"nextCursor":"next"}), + 2, + &mut sessions, + &mut cursors, + ) + .unwrap_err() + .to_string() + .contains("repeated cursor")); + assert!(muse_msp_append_page( + json!({"sessions":[{"status":"notLoaded"}]}), + 2, + &mut sessions, + &mut cursors, + ) + .unwrap_err() + .to_string() + .contains("decode Muse session/list page 2")); + } + + #[tokio::test] + async fn muse_msp_request_accepts_notifications_but_rejects_wrong_response_ids() { + async fn exchange(response: &str) -> anyhow::Result { + let (client, server) = tokio::io::duplex(4096); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, mut server_write) = tokio::io::split(server); + let response = response.to_owned(); + let server = tokio::spawn(async move { + let request = BufReader::new(server_read) + .lines() + .next_line() + .await + .unwrap() + .unwrap(); + assert_eq!( + serde_json::from_str::(&request).unwrap()["id"], + 7 + ); + server_write.write_all(response.as_bytes()).await.unwrap(); + }); + let result = muse_msp_request( + &mut client_write, + &mut BufReader::new(client_read).lines(), + 7, + "session/list", + json!({}), + ) + .await; + server.await.unwrap(); + result + } + + let result = exchange( + "{\"jsonrpc\":\"2.0\",\"method\":\"progress\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n", + ) + .await + .unwrap(); + assert_eq!(result, json!({"sessions": []})); + assert!(exchange("{\"jsonrpc\":\"2.0\",\"id\":8,\"result\":{}}\n") + .await + .unwrap_err() + .to_string() + .contains("unexpected response id")); + } + #[test] fn additional_metadata_overrides_a_route_only_with_a_json_object() { let mut route = SessionRoute { From 56f318a00e109f2db284e9aa099906fd244d6079 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:08:19 -0400 Subject: [PATCH 11/17] Stream Muse exports through direct import API --- bt-daemon/src/lib.rs | 102 ++++++++++++++---------- bt-daemon/src/transcript_import/muse.rs | 5 ++ 2 files changed, 67 insertions(+), 40 deletions(-) diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 2f1aa213..3427b4c8 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -1099,46 +1099,7 @@ async fn import_muse_sessions_with_ledger( .tempdir()?; let export = directory.path().join("session.json"); export_muse_session(session_id, &export).await?; - let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); - let export_for_parser = export.clone(); - let expected_id = session_id.clone(); - let parser = tokio::task::spawn_blocking(move || { - let mut batch = Vec::with_capacity(128); - transcript_import::muse::for_each_envelope_for_session( - &export_for_parser, - &expected_id, - |entry| { - batch.push(entry); - if batch.len() == 128 { - sender - .blocking_send(std::mem::take(&mut batch)) - .map_err(|_| { - anyhow::anyhow!("Muse import stopped while parsing {expected_id}") - })?; - } - Ok(()) - }, - )?; - if !batch.is_empty() { - sender.blocking_send(batch).map_err(|_| { - anyhow::anyhow!("Muse import stopped while parsing {expected_id}") - })?; - } - Ok::<(), anyhow::Error>(()) - }); - let mut processing_error = None; - while let Some(batch) = receiver.recv().await { - if let Err(error) = processor.process(batch).await { - processing_error = Some(error); - break; - } - } - drop(receiver); - let parse_result = parser.await.context("Muse export parser task failed")?; - if let Some(error) = processing_error { - return Err(error); - } - parse_result?; + process_muse_export(&mut processor, &export, session_id).await?; if let Some(summary) = processor.finish_session(session_id).await? { summaries.push(summary); } @@ -1147,6 +1108,53 @@ async fn import_muse_sessions_with_ledger( Ok(summaries) } +async fn process_muse_export( + processor: &mut ImportProcessor, + export: &std::path::Path, + session_id: &str, +) -> anyhow::Result<()> { + let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); + let export_for_parser = export.to_path_buf(); + let expected_id = session_id.to_owned(); + let parser = tokio::task::spawn_blocking(move || { + let mut batch = Vec::with_capacity(128); + transcript_import::muse::for_each_envelope_for_session( + &export_for_parser, + &expected_id, + |entry| { + batch.push(entry); + if batch.len() == 128 { + sender + .blocking_send(std::mem::take(&mut batch)) + .map_err(|_| { + anyhow::anyhow!("Muse import stopped while parsing {expected_id}") + })?; + } + Ok(()) + }, + )?; + if !batch.is_empty() { + sender + .blocking_send(batch) + .map_err(|_| anyhow::anyhow!("Muse import stopped while parsing {expected_id}"))?; + } + Ok::<(), anyhow::Error>(()) + }); + let mut processing_error = None; + while let Some(batch) = receiver.recv().await { + if let Err(error) = processor.process(batch).await { + processing_error = Some(error); + break; + } + } + drop(receiver); + let parse_result = parser.await.context("Muse export parser task failed")?; + if let Some(error) = processing_error { + return Err(error); + } + parse_result +} + 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"); @@ -1592,6 +1600,20 @@ async fn import_transcript_with_ledger( attach: bool, ledger_dir: Option, ) -> anyhow::Result> { + if source == ImportSource::Muse { + if attach { + anyhow::bail!("Muse import does not support --attach yet") + } + let session_id = transcript_import::muse::session_id(file)?; + let mut processor = ImportProcessor::new(opts, config, ledger_dir); + process_muse_export(&mut processor, file, &session_id).await?; + let mut summaries = Vec::new(); + if let Some(summary) = processor.finish_session(&session_id).await? { + summaries.push(summary); + } + summaries.extend(processor.finish().await?); + return Ok(summaries); + } let mut tail = transcript_import::TranscriptTail::new(file.to_path_buf(), source); let mut processor = ImportProcessor::new(opts, config, ledger_dir); let shutdown = tokio::signal::ctrl_c(); diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index 40328323..db27e30f 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -433,6 +433,11 @@ pub(crate) fn envelopes(path: &Path) -> Result> { envelopes_for_session(path, &session.session_id) } +pub(crate) fn session_id(path: &Path) -> Result { + let header = read_header(path)?; + Ok(single_session(&header, path)?.session_id.clone()) +} + pub(crate) fn envelopes_for_session(path: &Path, expected_id: &str) -> Result> { let mut events = Vec::new(); for_each_envelope_for_session(path, expected_id, |event| { From 3323cf9c1cc1104f4288f4b2ebd20dc66c2e77f5 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 23 Sep 2026 15:38:12 -0400 Subject: [PATCH 12/17] Clarify Muse attach constraints and resumed export ends --- bt-daemon/README.md | 13 +++++--- bt-daemon/src/transcript_import/muse.rs | 44 +++++++++++++++++++++---- 2 files changed, 46 insertions(+), 11 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 93d8853d..6533e994 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -289,10 +289,15 @@ Muse Code 1.1.1 setup captures the six verified lifecycle and model hooks: `SessionStart`, `UserPromptSubmit`, `PreLLMCall`, `PostLLMCall`, `Stop`, and `SessionEnd`. Tool, permission, subagent, and compaction details are not available from those hooks and are not inferred by the live translator. -`--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. +`--attach` is not available yet. MSP offers cursor-paged `view/page`, but its +view events omit the model request/response records needed by this translator; +`view/subscribe` also requires a session loaded on the same host. Repeated +full exports can see an active session but do not provide an incremental tail. +`run muse` is also deliberately unavailable: the observed process-local +managed-hook override replaces an existing managed hook file, and Muse's hook +environment drops the invocation route and duplicate-suppression variables. +Use persistent `enable muse` setup until those behaviors can be preserved +without changing unrelated hooks or tracing destinations. Add `--attach` to keep following an active Codex, Claude, or Antigravity transcript until Ctrl-C. `run [ARGS...]` launches the selected agent with diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index db27e30f..17dafd67 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -92,7 +92,7 @@ struct Reader<'a, F> { runs: HashMap, first_ts: Option, last_ts: Option, - session_end_ts: Option, + current_session_ended: bool, activity_count: usize, } @@ -249,9 +249,10 @@ where let ts = micros / 1_000; self.ensure_start(ts)?; if payload.kind.as_deref() == Some("session_end") { - self.session_end_ts = Some(ts); + self.current_session_ended = true; return Ok(()); } + self.current_session_ended = false; let event = payload .event .as_ref() @@ -398,9 +399,9 @@ where ); } let end_ts = self - .session_end_ts - .or(self.last_ts) + .last_ts .ok_or_else(|| anyhow!("Muse export {} has no timestamps", self.path.display()))?; + let current_session_ended = self.current_session_ended; for run_id in self.runs.keys().cloned().collect::>() { self.finish_model(&run_id, end_ts, Some("Run ended without a terminal event"))?; let run = self.runs.remove(&run_id).expect("run exists"); @@ -420,8 +421,14 @@ where self.emit("SessionEnd", end_ts, json!({ "session_id": self.session_id, "hook_event_name": "SessionEnd", - "reason": self.session.session_end.as_ref().and_then(|end| end.exit_reason.as_deref()), - "error": self.terminated_abnormally.then_some("Session terminated abnormally"), + "reason": current_session_ended.then(|| self.session.session_end.as_ref().and_then(|end| end.exit_reason.as_deref())).flatten(), + "error": if !current_session_ended { + Some("Session has no terminal end after its latest activity") + } else if self.terminated_abnormally { + Some("Session terminated abnormally") + } else { + None + }, "source": "muse_export_v1", })) } @@ -482,7 +489,7 @@ where runs: HashMap::new(), first_ts: None, last_ts: None, - session_end_ts: None, + current_session_ended: false, activity_count: 0, }; let mut de = serde_json::Deserializer::from_reader(BufReader::new( @@ -688,4 +695,27 @@ mod tests { let error = envelopes(&path).unwrap_err(); assert!(format!("{error:#}").contains("event 1: Muse started record has no run_id")); } + + #[test] + fn resumed_active_export_does_not_reuse_an_old_session_end() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "sessions": [{"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":"first"}}}}, + {"recorded_at": 1_001_000, "envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}, + {"recorded_at": 1_002_000, "envelope":{"payload":{"kind":"session_end"}}}, + {"recorded_at": 1_003_000, "envelope":{"payload":{"kind":"run","run_id":"run-2","event":{"kind":"started","prompt":"second"}}}} + ] + })); + let events = envelopes(&path).unwrap(); + let end = events.last().unwrap(); + assert_eq!(end.event, "SessionEnd"); + assert_eq!(end.ts_ms, 1003); + assert_eq!(end.payload["reason"], Value::Null); + assert_eq!( + end.payload["error"], + "Session has no terminal end after its latest activity" + ); + } } From 36ab99b54c10310616a180e4c0337af140d078b1 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 24 Sep 2026 11:47:22 -0400 Subject: [PATCH 13/17] Add invocation-local Muse managed runs --- bt-daemon/README.md | 18 ++- bt-daemon/src/lib.rs | 214 ++++++++++++++++++++++++++++++--- bt-daemon/src/main.rs | 2 +- bt-daemon/src/setup.rs | 193 ++++++++++++++++++++++++++++- bt-daemon/src/trace_command.rs | 12 ++ bt-daemon/tests/pipeline.rs | 176 +++++++++++++++++++++++++++ 6 files changed, 592 insertions(+), 23 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 6533e994..d8c97210 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -293,14 +293,20 @@ available from those hooks and are not inferred by the live translator. view events omit the model request/response records needed by this translator; `view/subscribe` also requires a session loaded on the same host. Repeated full exports can see an active session but do not provide an incremental tail. -`run muse` is also deliberately unavailable: the observed process-local -managed-hook override replaces an existing managed hook file, and Muse's hook -environment drops the invocation route and duplicate-suppression variables. -Use persistent `enable muse` setup until those behaviors can be preserved -without changing unrelated hooks or tracing destinations. +`run muse` uses a process-local managed-hook override in the verified Muse Code +1.1.1 release. It retains compatible unrelated managed hooks, replaces Braintrust's saved +managed hooks only for that invocation, and leaves ordinary Muse sessions and +settings unchanged. The injected hook reads a private invocation context file +because Muse clears its hook environment. A dedicated daemon is started before +Muse so environment-based Braintrust authentication never depends on that +cleared environment. Known duplicate Braintrust user or project hooks stop the +managed run with a diagnostic. A foreign managed file using Muse's strict +handler schema is also rejected rather than silently losing hooks. Muse plugins that independently trace the same +session should be disabled before using `run muse`; new Muse versions must be +verified before the undocumented hook-path override is enabled for them. Add `--attach` to keep following an active Codex, Claude, or Antigravity transcript until -Ctrl-C. `run [ARGS...]` launches the selected agent with +Ctrl-C. `run [ARGS...]` launches the selected agent with inherited stdio and injects Braintrust hooks or an adapter for that invocation, so it does not depend on the tracing plugin being installed or enabled. Managed runs suppress inherited Braintrust plugin hooks to avoid logging the same session diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 3427b4c8..d4529f35 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -154,6 +154,9 @@ pub struct HookArgs { /// not carry this flag and are suppressed for the managed child. #[arg(long, hide = true)] pub managed_run_hook: bool, + /// Private invocation context for Muse's environment-sanitized hooks. + #[arg(long, hide = true)] + pub managed_context: Option, } /// Arguments for `status`. @@ -315,6 +318,7 @@ pub struct RunHookCommand { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum RunSource { Codex, + Muse, #[value(name = "claude", alias = "claude-code")] Claude, #[value(name = "opencode", alias = "open-code")] @@ -351,6 +355,13 @@ fn build_hook_envelope( } } +#[derive(Debug, serde::Serialize, serde::Deserialize)] +struct MuseManagedContext { + route: SessionRoute, + managed_run_id: String, + socket: PathBuf, +} + pub(crate) fn should_flush_hook_event(event: &str, flush_on_turn_end: bool) -> bool { matches!(event, "SessionEnd" | "session_end") || (flush_on_turn_end @@ -405,14 +416,33 @@ pub async fn run_hook( mut route: SessionRoute, host: HostInfo, ) -> anyhow::Result<()> { + let managed_context = args + .managed_context + .as_ref() + .map(|path| -> anyhow::Result { + if args.source != "muse" || !args.managed_run_hook { + anyhow::bail!("managed context is only valid for an injected Muse hook"); + } + let context: MuseManagedContext = serde_json::from_slice(&std::fs::read(path)?)?; + if context.managed_run_id.is_empty() || context.route.destination.is_none() { + anyhow::bail!("Muse managed hook context is incomplete"); + } + Ok(context) + }) + .transpose()?; // A managed run injects its own hook definitions. Suppress an inherited // Braintrust plugin hook for the same child, but allow the injected hook // process, which carries the second marker. if std::env::var_os("_BT_TRACE_MANAGED_RUN").is_some() && !args.managed_run_hook { return Ok(()); } - let settings = settings::AgentSettings::load(&args.source); - if !settings.tracing_enabled() { + let settings = managed_context + .is_none() + .then(|| settings::AgentSettings::load(&args.source)); + if settings + .as_ref() + .is_some_and(|settings| !settings.tracing_enabled()) + { return Ok(()); } let mut payload = read_stdin_json()?; @@ -431,17 +461,31 @@ pub async fn run_hook( .or_else(|| json_str_field(&payload, &args.event_field)) .unwrap_or_default(); - if let Some(configured_route) = settings.route { + if let Some(context) = &managed_context { + route = context.route.clone(); + } else if let Some(configured_route) = settings.and_then(|settings| settings.route) { route = configured_route; } if args.flush_on_turn_end { route.flush_mode = wire::FlushMode::FlushOnTurnEnd; } apply_additional_metadata(&mut route, args.additional_metadata.as_deref())?; - let env = build_hook_envelope(&args, route, payload, session_id, event); + let mut env = build_hook_envelope(&args, route, payload, session_id, event); + if let Some(context) = &managed_context { + env.managed_run_id = Some(context.managed_run_id.clone()); + } - let socket = paths::socket_path(args.socket.as_deref()); - forward_envelope(&env, &socket, &host, args.no_spawn).await?; + let socket = managed_context + .as_ref() + .map(|context| context.socket.clone()) + .unwrap_or_else(|| paths::socket_path(args.socket.as_deref())); + forward_envelope( + &env, + &socket, + &host, + args.no_spawn || managed_context.is_some(), + ) + .await?; Ok(()) } @@ -1249,6 +1293,7 @@ pub async fn run_traced( args: RunArgs, hook_command: RunHookCommand, mut route: SessionRoute, + host: HostInfo, ) -> anyhow::Result { apply_run_span_plugins(&mut route, &args.plugin)?; if route.destination.is_none() { @@ -1269,23 +1314,57 @@ pub async fn run_traced( } let (executable_env, default_executable) = match args.source { RunSource::Codex => ("CODEX_BIN", "codex"), + RunSource::Muse => ("MUSE_BIN", "muse"), RunSource::Claude => ("CLAUDE_BIN", "claude"), RunSource::OpenCode => ("OPENCODE_BIN", "opencode"), RunSource::Pi => ("PI_BIN", "pi"), }; let executable = std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); + if args.source == RunSource::Muse { + setup::check_muse_project_hooks(&std::env::current_dir()?)?; + for (index, arg) in args.agent_args.iter().enumerate() { + let Some(arg) = arg.to_str() else { continue }; + let workspace = if arg == "--workspace" { + args.agent_args.get(index + 1).map(PathBuf::from) + } else { + arg.strip_prefix("--workspace=").map(PathBuf::from) + }; + if let Some(workspace) = workspace { + setup::check_muse_project_hooks(&workspace)?; + } + } + let output = tokio::process::Command::new(&executable) + .arg("--version") + .output() + .await + .context("check Muse Code version for managed hooks")?; + if !output.status.success() + || setup::muse_version(&output.stdout)? != semver::Version::new(1, 1, 1) + { + anyhow::bail!("Muse managed runs require the verified Muse Code 1.1.1 hook override"); + } + } let injected_args = managed_run_args(args.source, &hook_command)?; let managed_run_id = uuid::Uuid::new_v4().to_string(); // A caller-supplied socket is already an explicit daemon boundary (for // example an integration harness or a deliberately isolated runtime). // Only create our own boundary when environment auth would otherwise use // the ambient shared daemon. - let isolate_daemon = route.auth.effective_source() == wire::AuthSource::Environment + let isolate_daemon = (args.source == RunSource::Muse + || route.auth.effective_source() == wire::AuthSource::Environment) && std::env::var_os(paths::SOCKET_ENV).is_none(); let isolated_runtime = isolate_daemon .then(|| ManagedRunRuntime::new(&managed_run_id)) .transpose()?; + let muse_context_dir = (args.source == RunSource::Muse && isolated_runtime.is_none()) + .then(|| { + tempfile::Builder::new() + .prefix("bt-trace-run-muse-") + .tempdir() + }) + .transpose()?; + let muse_route = (args.source == RunSource::Muse).then(|| route.clone()); let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?; let mut command = tokio::process::Command::new(&executable); command @@ -1294,6 +1373,40 @@ pub async fn run_traced( .env("_BT_TRACE_MANAGED_RUN", "1") .env(MANAGED_RUN_ID_ENV, &managed_run_id) .env(settings::INVOCATION_SETTINGS_ENV, invocation_settings); + if args.source == RunSource::Muse { + let context_dir = isolated_runtime + .as_ref() + .map(|runtime| runtime.temp_dir.path()) + .or_else(|| muse_context_dir.as_ref().map(|dir| dir.path())) + .expect("Muse has an invocation directory"); + let context = MuseManagedContext { + route: muse_route.expect("Muse route is available"), + managed_run_id: managed_run_id.clone(), + socket: isolated_runtime + .as_ref() + .map(|runtime| runtime.socket.clone()) + .unwrap_or_else(|| paths::socket_path(None)), + }; + let context_path = context_dir.join("muse-context.json"); + std::fs::write(&context_path, serde_json::to_vec(&context)?)?; + let mut hook = hook_command.clone(); + hook.args.extend([ + OsString::from("--source"), + OsString::from("muse"), + OsString::from("--managed-run-hook"), + OsString::from("--flush-on-turn-end"), + OsString::from("--managed-context"), + context_path.into_os_string(), + ]); + let command_string = shell_command(&hook, cfg!(windows))?; + let hooks = setup::muse_managed_hook_config(&paths::muse_config_dir(), &command_string)?; + let hooks_path = context_dir.join("muse-hooks.json"); + std::fs::write(&hooks_path, serde_json::to_vec(&hooks)?)?; + command.env("TBH_MANAGED_HOOKS_PATH", hooks_path); + if let Some(runtime) = &isolated_runtime { + start_muse_managed_daemon(&host, runtime).await?; + } + } if let Some(runtime) = &isolated_runtime { command .env(paths::SOCKET_ENV, &runtime.socket) @@ -1305,9 +1418,20 @@ pub async fn run_traced( opencode_managed_config(std::env::var("OPENCODE_CONFIG_CONTENT").ok().as_deref())?, ); } - let mut child = command.spawn().map_err(|error| { - anyhow::anyhow!("failed to launch {}: {error}", executable.to_string_lossy()) - })?; + let mut child = match command.spawn() { + Ok(child) => child, + Err(error) => { + if args.source == RunSource::Muse { + if let Some(runtime) = &isolated_runtime { + let _ = shutdown_daemon(&runtime.socket).await; + } + } + return Err(anyhow::anyhow!( + "failed to launch {}: {error}", + executable.to_string_lossy() + )); + } + }; let interrupt = tokio::signal::ctrl_c(); tokio::pin!(interrupt); @@ -1383,6 +1507,18 @@ impl ManagedRunRuntime { } } +async fn start_muse_managed_daemon( + host: &HostInfo, + runtime: &ManagedRunRuntime, +) -> anyhow::Result<()> { + let stream = + client::ensure_daemon_in(&runtime.socket, host, false, Some(runtime.temp_dir.path())) + .await + .context("start isolated daemon for Muse managed run")?; + drop(stream); + Ok(()) +} + fn apply_run_span_plugins(route: &mut SessionRoute, plugins: &[PathBuf]) -> anyhow::Result<()> { route.span_plugins = resolve_span_plugin_paths(plugins)?; Ok(()) @@ -1407,6 +1543,7 @@ fn managed_run_args( ) -> anyhow::Result> { let source_name = match source { RunSource::Codex => "codex", + RunSource::Muse => "muse", RunSource::Claude => "claude", RunSource::OpenCode => "opencode", RunSource::Pi => "pi", @@ -1425,7 +1562,7 @@ fn managed_run_args( _ => unreachable!(), } } - RunSource::OpenCode => Ok(Vec::new()), + RunSource::Muse | RunSource::OpenCode => Ok(Vec::new()), RunSource::Pi => { let extension = match std::env::var_os("BT_TRACE_PI_PLUGIN_SPEC") { Some(extension) => extension, @@ -1469,15 +1606,25 @@ fn managed_hook_shell_command( argv.push(OsString::from("--source")); argv.push(OsString::from(source)); argv.push(OsString::from("--managed-run-hook")); + shell_argv(&argv, windows) +} + +fn shell_command(command: &RunHookCommand, windows: bool) -> anyhow::Result { + let mut argv = vec![command.program.clone()]; + argv.extend(command.args.iter().cloned()); + shell_argv(&argv, windows) +} + +fn shell_argv(argv: &[OsString], windows: bool) -> anyhow::Result { let mut rendered = Vec::with_capacity(argv.len()); for arg in argv { let arg = arg - .into_string() - .map_err(|_| anyhow::anyhow!("managed hook command contains non-Unicode argv"))?; + .to_str() + .ok_or_else(|| anyhow::anyhow!("managed hook command contains non-Unicode argv"))?; rendered.push(if windows { - quote_windows_command_arg(&arg) + quote_windows_command_arg(arg) } else { - quote_unix_shell_arg(&arg) + quote_unix_shell_arg(arg) }); } Ok(rendered.join(" ")) @@ -2684,6 +2831,10 @@ mod tests { }, test_run_hook_command(), SessionRoute::default(), + HostInfo { + serve_argv: Vec::new(), + version: "test".into(), + }, ) .await .unwrap_err(); @@ -2709,6 +2860,10 @@ mod tests { }), ..SessionRoute::default() }, + HostInfo { + serve_argv: Vec::new(), + version: "test".into(), + }, ) .await .unwrap_err(); @@ -2806,4 +2961,33 @@ mod tests { assert!(windows .contains("\"/opt/Braintrust CLI/bt\" \"agents\" \"hook\" \"--source\" \"claude\"")); } + + #[tokio::test] + async fn muse_managed_hook_rejects_missing_context_before_reading_stdin() { + let temp = tempfile::tempdir().unwrap(); + let args = HookCli::try_parse_from([ + "test", + "--source", + "muse", + "--managed-run-hook", + "--managed-context", + temp.path().join("missing.json").to_str().unwrap(), + ]) + .unwrap() + .args; + let error = run_hook( + args, + SessionRoute::default(), + HostInfo { + serve_argv: Vec::new(), + version: "test".into(), + }, + ) + .await + .unwrap_err(); + assert_eq!( + error.downcast_ref::().unwrap().kind(), + std::io::ErrorKind::NotFound + ); + } } diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index 87576031..a0c9ac9f 100644 --- a/bt-daemon/src/main.rs +++ b/bt-daemon/src/main.rs @@ -300,7 +300,7 @@ async fn main() { program: exe, args: vec![OsString::from("hook")], }; - match run_traced(args, hook_command, route.into_route()).await { + match run_traced(args, hook_command, route.into_route(), host_info()).await { Ok(status) if status.success() => {} Ok(status) => std::process::exit(status.code().unwrap_or(1)), Err(error) => { diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index ae12b7d3..96dc84ac 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -739,6 +739,110 @@ fn muse_hook_config_for(events: &[&str], command: &str) -> Map { ]) } +/// Build an invocation-local managed file without changing Muse's settings. +/// Preserve handlers administered by someone else, but never run two Muse +/// Braintrust forwarders with different destinations in the same session. +pub(crate) fn muse_managed_hook_config(config_dir: &Path, command: &str) -> anyhow::Result { + let settings = load_object(&config_dir.join("settings.json"))?; + if settings + .get("hooks") + .is_some_and(muse_contains_braintrust_hook) + { + bail!("Muse user hooks already include Braintrust tracing; remove that duplicate before `bt trace run muse`"); + } + let existing = match settings.get("managed_hooks_path").and_then(Value::as_str) { + Some(path) => { + let path = Path::new(path); + let path = if path.is_absolute() { + path.to_path_buf() + } else { + config_dir.join(path) + }; + if !path.exists() || muse_hook_is_owned(&path)? { + Map::from_iter([("schema_version".into(), Value::from(1))]) + } else { + let raw = std::fs::read(&path) + .with_context(|| format!("read Muse managed hooks {}", path.display()))?; + let value: Value = serde_json::from_slice(&raw)?; + if muse_contains_braintrust_hook(&value) { + bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); + } + if muse_has_strict_handler_fields(&value) { + bail!("Muse managed hooks use the strict handler schema; cannot safely merge invocation hooks"); + } + value + .as_object() + .cloned() + .context("Muse managed hooks must be a JSON object")? + } + } + None => Map::from_iter([("schema_version".into(), Value::from(1))]), + }; + let mut result = existing; + let mut hooks = result + .remove("hooks") + .unwrap_or_else(|| Value::Object(Map::new())); + let hooks = hooks + .as_object_mut() + .context("Muse managed hooks must be an object")?; + let generated = muse_hook_config_for(MUSE_HOOK_EVENTS, command); + for (event, handlers) in generated["hooks"].as_object().unwrap() { + let existing = hooks + .entry(event.clone()) + .or_insert_with(|| Value::Array(Vec::new())); + existing + .as_array_mut() + .with_context(|| format!("Muse managed hook event {event} must be an array"))? + .extend(handlers.as_array().unwrap().iter().cloned()); + } + result.insert("hooks".into(), hooks.clone().into()); + Ok(Value::Object(result)) +} + +pub(crate) fn check_muse_project_hooks(start: &Path) -> anyhow::Result<()> { + for dir in start.ancestors() { + let path = dir.join(".muse").join("hooks.json"); + if !path.exists() { + continue; + } + let hooks: Value = serde_json::from_slice(&std::fs::read(&path)?) + .with_context(|| format!("read Muse project hooks {}", path.display()))?; + if muse_contains_braintrust_hook(&hooks) { + bail!("Muse project hooks {} already include Braintrust tracing; cannot safely run a managed trace", path.display()); + } + break; + } + Ok(()) +} + +fn muse_contains_braintrust_hook(value: &Value) -> bool { + match value { + Value::Object(object) => object.iter().any(|(key, value)| { + (key == "command" + && value.as_str().is_some_and(|command| { + command.contains("--source") + && command.contains("muse") + && command.contains("hook") + })) + || muse_contains_braintrust_hook(value) + }), + Value::Array(items) => items.iter().any(muse_contains_braintrust_hook), + _ => false, + } +} + +fn muse_has_strict_handler_fields(value: &Value) -> bool { + match value { + Value::Object(object) => { + object.contains_key("name") + || object.contains_key("timeout_ms") + || object.values().any(muse_has_strict_handler_fields) + } + Value::Array(items) => items.iter().any(muse_has_strict_handler_fields), + _ => false, + } +} + fn legacy_muse_hook_config() -> Map { muse_hook_config_for( LEGACY_MUSE_HOOK_EVENTS, @@ -746,7 +850,7 @@ fn legacy_muse_hook_config() -> Map { ) } -fn muse_version(output: &[u8]) -> anyhow::Result { +pub(crate) fn muse_version(output: &[u8]) -> anyhow::Result { let text = std::str::from_utf8(output).context("Muse version is not UTF-8")?; text.split_whitespace() .find_map(|word| semver::Version::parse(word).ok()) @@ -1919,6 +2023,93 @@ mod tests { assert!(!muse_hook_path(&config_dir).exists()); } + #[test] + fn muse_managed_run_replaces_only_owned_braintrust_hooks() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + setup_muse_at(&config_dir).unwrap(); + let saved = std::fs::read(muse_hook_path(&config_dir)).unwrap(); + let generated = + muse_managed_hook_config(&config_dir, "'bt' 'hook' '--source' 'muse'").unwrap(); + assert_eq!( + generated["hooks"].as_object().unwrap().len(), + MUSE_HOOK_EVENTS.len() + ); + assert_eq!( + generated["hooks"]["SessionStart"][0]["hooks"][0]["command"], + "'bt' 'hook' '--source' 'muse'" + ); + assert_eq!(std::fs::read(muse_hook_path(&config_dir)).unwrap(), saved); + } + + #[test] + fn muse_managed_run_keeps_foreign_handlers_and_rejects_duplicate_user_hooks() { + 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.json"}"#, + ) + .unwrap(); + let original = r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"true"}]}]}}"#; + std::fs::write(config_dir.join("other.json"), original).unwrap(); + let generated = muse_managed_hook_config(&config_dir, "bt hook --source muse").unwrap(); + assert_eq!(generated["hooks"]["Stop"].as_array().unwrap().len(), 2); + assert_eq!(generated["hooks"]["Stop"][0]["hooks"][0]["command"], "true"); + assert_eq!( + std::fs::read_to_string(config_dir.join("other.json")).unwrap(), + original + ); + + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"bt trace hook --source muse"}]}]}}"#, + ) + .unwrap(); + assert!( + muse_managed_hook_config(&config_dir, "bt hook --source muse") + .unwrap_err() + .to_string() + .contains("user hooks already include Braintrust") + ); + + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"managed_hooks_path":"other.json"}"#, + ) + .unwrap(); + std::fs::write( + config_dir.join("other.json"), + r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","name":"foreign","timeout_ms":1000,"command":"true"}]}]}}"#, + ) + .unwrap(); + assert!( + muse_managed_hook_config(&config_dir, "bt hook --source muse") + .unwrap_err() + .to_string() + .contains("strict handler schema") + ); + } + + #[test] + fn muse_managed_run_rejects_duplicate_project_hooks() { + let temp = tempfile::tempdir().unwrap(); + let project = temp.path().join("project"); + let nested = project.join("nested"); + std::fs::create_dir_all(project.join(".muse")).unwrap(); + std::fs::create_dir_all(&nested).unwrap(); + std::fs::write( + project.join(".muse/hooks.json"), + r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"'bt' 'trace' 'hook' '--source' 'muse'"}]}]}}"#, + ) + .unwrap(); + assert!(check_muse_project_hooks(&nested) + .unwrap_err() + .to_string() + .contains("project hooks")); + } + #[test] fn muse_setup_refuses_another_managed_hook_owner() { let temp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs index eb308473..18399fcf 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -400,6 +400,18 @@ mod tests { } } + #[test] + fn run_accepts_muse() { + let parsed = Cli::try_parse_from(["bt", "run", "muse", "exec", "hello"]).unwrap(); + assert!(matches!( + parsed.trace.command, + TraceCommand::Run(RunArgs { + source: crate::RunSource::Muse, + .. + }) + )); + } + #[test] fn public_commands_preserve_repeated_plugin_order() { for args in [ diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 8a38f31d..f49e321d 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -2691,6 +2691,13 @@ impl EnvVarGuard { std::env::set_var(key, value); Self { key, previous } } + + #[cfg(unix)] + fn unset(key: &'static str) -> Self { + let previous = std::env::var_os(key); + std::env::remove_var(key); + Self { key, previous } + } } #[cfg(feature = "cli")] @@ -2770,6 +2777,7 @@ esac args: Vec::new(), }, route(), + dummy_host(), ) }; @@ -2807,6 +2815,174 @@ esac handle.await.unwrap(); } +#[cfg(all(feature = "cli", unix))] +#[tokio::test] +#[ignore = "requires a local Muse Code 1.1.1 installation"] +async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { + let output = match std::process::Command::new("muse").arg("--version").output() { + Ok(output) if output.status.success() => output, + _ => return, + }; + if !String::from_utf8_lossy(&output.stdout).contains("1.1.1") { + return; + } + let tmp = tempfile::tempdir().unwrap(); + let config_dir = tmp.path().join("muse"); + std::fs::create_dir_all(&config_dir).unwrap(); + let marker = tmp.path().join("foreign-hook.txt"); + let foreign_command = format!("printf foreign >> '{}'", marker.display()); + let managed = serde_json::json!({ + "schema_version": 1, + "hooks": {"SessionStart": [{"matcher": "*", "hooks": [{"type": "command", "command": foreign_command}]}]} + }); + std::fs::write(config_dir.join("global.json"), managed.to_string()).unwrap(); + let settings = r#"{"schema_version":1,"managed_hooks_path":"global.json"}"#; + std::fs::write(config_dir.join("settings.json"), settings).unwrap(); + let _config = EnvVarGuard::set("XDG_CONFIG_HOME", tmp.path()); + let (socket, daemon, flushes, daemon_dir) = + start_tracking_daemon(env!("CARGO_PKG_VERSION")).await; + let _socket = EnvVarGuard::set("BT_DAEMON_SOCKET", &socket); + let _data_dir = EnvVarGuard::set("BT_DAEMON_DATA_DIR", daemon_dir.path().join("data")); + let route = SessionRoute { + auth: AuthSelection { + source: AuthSource::Environment, + ..AuthSelection::default() + }, + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("muse-managed-test".into()), + }), + ..SessionRoute::default() + }; + let binary = env!("CARGO_BIN_EXE_bt-daemon"); + let status = run_traced( + RunArgs { + source: RunSource::Muse, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + agent_args: vec![ + "exec".into(), + "--provider".into(), + "echo".into(), + "--no-foreign-personal-context".into(), + "hello".into(), + ], + }, + RunHookCommand { + program: binary.into(), + args: vec!["hook".into()], + }, + route, + HostInfo { + serve_argv: vec![binary.into(), "serve".into(), "--debug-sink".into()], + version: env!("CARGO_PKG_VERSION").into(), + }, + ) + .await + .unwrap(); + assert!(status.success()); + assert_eq!(std::fs::read_to_string(&marker).unwrap(), "foreign"); + let saved: serde_json::Value = + serde_json::from_slice(&std::fs::read(config_dir.join("settings.json")).unwrap()).unwrap(); + assert_eq!(saved["managed_hooks_path"], "global.json"); + assert_eq!( + std::fs::read_to_string(config_dir.join("global.json")).unwrap(), + managed.to_string() + ); + let flushed = flushes.lock().unwrap().clone(); + assert!(!flushed.is_empty()); + let mut main_session = false; + for session in flushed.keys() { + let journal = source_journal_path(&daemon_dir.path().join("data"), "muse", session); + for line in std::fs::read_to_string(journal).unwrap().lines() { + let row: serde_json::Value = serde_json::from_str(line).unwrap(); + if row["source"] != "muse" { + continue; + } + assert_eq!( + row["route"]["destination"]["project_name"], + "muse-managed-test" + ); + assert!(row["managed_run_id"].as_str().is_some()); + main_session |= row["event"] == "SessionStart"; + } + } + assert!(main_session); + assert!(flushed.values().copied().sum::() >= 1); + let ordinary = std::process::Command::new("muse") + .args([ + "exec", + "--provider", + "echo", + "--no-foreign-personal-context", + "ordinary", + ]) + .output() + .unwrap(); + assert!(ordinary.status.success()); + assert_eq!(std::fs::read_to_string(&marker).unwrap(), "foreignforeign"); + shutdown(&socket).await; + daemon.await.unwrap(); +} + +#[cfg(all(feature = "cli", unix))] +#[tokio::test] +#[ignore = "requires a local Muse Code 1.1.1 installation"] +async fn muse_managed_run_starts_its_isolated_daemon_before_sanitized_hooks() { + let output = match std::process::Command::new("muse").arg("--version").output() { + Ok(output) if output.status.success() => output, + _ => return, + }; + if !String::from_utf8_lossy(&output.stdout).contains("1.1.1") { + return; + } + let tmp = tempfile::tempdir().unwrap(); + let _config = EnvVarGuard::set("XDG_CONFIG_HOME", tmp.path()); + let _socket = EnvVarGuard::unset("BT_DAEMON_SOCKET"); + let binary = env!("CARGO_BIN_EXE_bt-daemon"); + let status = run_traced( + RunArgs { + source: RunSource::Muse, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + agent_args: vec![ + "exec".into(), + "--provider".into(), + "echo".into(), + "--no-foreign-personal-context".into(), + "hello".into(), + ], + }, + RunHookCommand { + program: binary.into(), + args: vec!["hook".into()], + }, + SessionRoute { + auth: AuthSelection { + source: AuthSource::Environment, + ..AuthSelection::default() + }, + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("isolated-muse-test".into()), + }), + ..SessionRoute::default() + }, + HostInfo { + serve_argv: vec![binary.into(), "serve".into(), "--debug-sink".into()], + version: env!("CARGO_PKG_VERSION").into(), + }, + ) + .await + .unwrap(); + assert!(status.success()); +} + +#[cfg(all(feature = "cli", unix))] +type CreatedDebugSinks = Arc)>>>; + #[cfg(all(feature = "cli", unix))] struct FlushTrackingDebugSinkFactory { inner: DebugSinkFactory, From 001813d6aee203df9204182e0edf501c3f2892b1 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 24 Sep 2026 14:33:17 -0400 Subject: [PATCH 14/17] Add Muse full-export attach support --- bt-daemon/README.md | 13 +- bt-daemon/src/lib.rs | 232 +++++++++++++++++++++++- bt-daemon/src/transcript_import/muse.rs | 64 ++++++- bt-daemon/tests/replay.rs | 74 ++++++++ 4 files changed, 370 insertions(+), 13 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index d8c97210..053ed21d 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -289,10 +289,15 @@ Muse Code 1.1.1 setup captures the six verified lifecycle and model hooks: `SessionStart`, `UserPromptSubmit`, `PreLLMCall`, `PostLLMCall`, `Stop`, and `SessionEnd`. Tool, permission, subagent, and compaction details are not available from those hooks and are not inferred by the live translator. -`--attach` is not available yet. MSP offers cursor-paged `view/page`, but its -view events omit the model request/response records needed by this translator; -`view/subscribe` also requires a session loaded on the same host. Repeated -full exports can see an active session but do not provide an incremental tail. +`import muse --attach` follows an active session by polling full +exports. Each snapshot is parsed with the same reader and translator as +historical import; only new envelopes are delivered, and an active snapshot +does not invent `Stop` or `SessionEnd` events. The attach exits after a native +session end, or finalizes open spans on interruption. This reads the full +export each poll, so polling cost grows with long sessions. MSP offers +cursor-paged `view/page`, but its view events omit the model request/response +records needed by this translator; `view/subscribe` also requires a session +loaded on the same host. `run muse` uses a process-local managed-hook override in the verified Muse Code 1.1.1 release. It retains compatible unrelated managed hooks, replaces Braintrust's saved managed hooks only for that invocation, and leaves ordinary Muse sessions and diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index d4529f35..d60850d5 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -834,9 +834,13 @@ pub async fn run_import( 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" + return attach_muse_session( + &args.session_ids[0], + opts, + config, + Some(paths::data_dir(None)), ) + .await; } let session_ids = if args.all { muse_msp_session_ids().await? @@ -1157,15 +1161,36 @@ async fn process_muse_export( export: &std::path::Path, session_id: &str, ) -> anyhow::Result<()> { + process_muse_snapshot(processor, export, session_id, false, 0) + .await + .map(|_| ()) +} + +/// The file is fully parsed on every poll. Only envelopes beyond the previous +/// snapshot are sent to the long-lived translator and sink, so open model and +/// turn state is retained without replaying old starts into it. +async fn process_muse_snapshot( + processor: &mut ImportProcessor, + export: &std::path::Path, + session_id: &str, + attach_snapshot: bool, + skip_envelopes: usize, +) -> anyhow::Result<(usize, bool)> { let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); let export_for_parser = export.to_path_buf(); let expected_id = session_id.to_owned(); let parser = tokio::task::spawn_blocking(move || { let mut batch = Vec::with_capacity(128); - transcript_import::muse::for_each_envelope_for_session( + let mut envelope_count = 0usize; + let ended = transcript_import::muse::for_each_snapshot_envelope_for_session( &export_for_parser, &expected_id, + attach_snapshot, |entry| { + envelope_count += 1; + if envelope_count <= skip_envelopes { + return Ok(()); + } batch.push(entry); if batch.len() == 128 { sender @@ -1177,12 +1202,17 @@ async fn process_muse_export( Ok(()) }, )?; + if envelope_count < skip_envelopes { + anyhow::bail!( + "Muse export for {expected_id} shrank from {skip_envelopes} to {envelope_count} trace events" + ); + } if !batch.is_empty() { sender .blocking_send(batch) .map_err(|_| anyhow::anyhow!("Muse import stopped while parsing {expected_id}"))?; } - Ok::<(), anyhow::Error>(()) + Ok::<(usize, bool), anyhow::Error>((envelope_count, ended)) }); let mut processing_error = None; while let Some(batch) = receiver.recv().await { @@ -1199,6 +1229,62 @@ async fn process_muse_export( parse_result } +async fn attach_muse_session( + session_id: &str, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + transcript_import::validate_session_id(session_id)?; + let directory = tempfile::Builder::new() + .prefix("bt-muse-attach-") + .tempdir()?; + let export = directory.path().join("session.json"); + attach_muse_export(&export, session_id, true, opts, config, ledger_dir).await +} + +async fn attach_muse_export( + export: &std::path::Path, + session_id: &str, + reexport: bool, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + let mut processor = ImportProcessor::new(opts, config, ledger_dir); + let mut emitted = 0usize; + let shutdown = tokio::signal::ctrl_c(); + tokio::pin!(shutdown); + loop { + if reexport { + export_muse_session(session_id, export).await?; + } + let (count, ended) = + process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; + if count != emitted { + processor.flush_open().await?; + } + emitted = count; + if ended { + return processor.finish().await; + } + tokio::select! { + result = &mut shutdown => { + result?; + break; + } + _ = tokio::time::sleep(Duration::from_secs(1)) => {} + } + } + // Capture one final complete snapshot before interrupting any still-open + // spans, matching the finalization behavior of other attach sources. + if reexport { + export_muse_session(session_id, export).await?; + } + process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; + processor.finish().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"); @@ -1748,10 +1834,10 @@ async fn import_transcript_with_ledger( ledger_dir: Option, ) -> anyhow::Result> { if source == ImportSource::Muse { + let session_id = transcript_import::muse::session_id(file)?; if attach { - anyhow::bail!("Muse import does not support --attach yet") + return attach_muse_export(file, &session_id, false, opts, config, ledger_dir).await; } - let session_id = transcript_import::muse::session_id(file)?; let mut processor = ImportProcessor::new(opts, config, ledger_dir); process_muse_export(&mut processor, file, &session_id).await?; let mut summaries = Vec::new(); @@ -2037,6 +2123,14 @@ impl ImportProcessor { Ok(summaries) } + async fn flush_open(&mut self) -> anyhow::Result<()> { + for live in self.sessions.values_mut() { + live.sink.flush().await?; + live.pending_ops = 0; + } + Ok(()) + } + async fn finish_session(&mut self, session_id: &str) -> anyhow::Result> { let Some(mut live) = self.sessions.remove(session_id) else { return Ok(None); @@ -2181,6 +2275,94 @@ mod tests { args: HookArgs, } + #[tokio::test] + async fn muse_attach_snapshots_reuse_one_trace_and_keep_open_spans() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("export.json"); + let output = temp.path().join("spans"); + let mut processor = ImportProcessor::new( + ServeOptions { + version: "test".into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(DebugSinkFactory { + dir: output.clone(), + }), + auth_provider: None, + }, + None, + None, + ); + let mut events = vec![ + json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}), + json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","request_record_id":"request-1"}}}}), + ]; + let write = |events: &[serde_json::Value]| { + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-attach"}], + "events":events, + })) + .unwrap(), + ) + .unwrap(); + }; + write(&events); + let (count, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, 0) + .await + .unwrap(); + assert_eq!(count, 3); + assert!(!ended); + processor.flush_open().await.unwrap(); + let (same_count, ended) = + process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) + .await + .unwrap(); + assert_eq!(same_count, count); + assert!(!ended); + events.extend([ + json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_completed","usage":{"input_tokens":2}}}}}), + json!({"recorded_at":1_003_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"done"}}}}), + json!({"recorded_at":1_004_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), + json!({"recorded_at":1_005_000,"envelope":{"payload":{"kind":"session_end"}}}), + ]); + write(&events); + let (_, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) + .await + .unwrap(); + assert!(ended); + processor.finish().await.unwrap(); + let rows = std::fs::read_to_string(output.join("muse-attach.ndjson")) + .unwrap() + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + for identity in ["root", "turn:run-1", "llm:request-1"] { + let id = ids::span_id(&ids::session_namespace("muse", "muse-attach"), identity); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(|value| value.as_str()) + == Some(id.as_str())) + .count(), + 1, + "duplicate insert for {identity}" + ); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Merge/span_id") + .and_then(|value| value.as_str()) + == Some(id.as_str())) + .count(), + 1, + "missing or duplicate completion for {identity}" + ); + } + } + #[test] fn serve_defaults_to_short_journal_backed_session_retirement() { let args = ServeCli::try_parse_from(["test"]).unwrap().args; @@ -2812,6 +2994,44 @@ mod tests { ); } + #[tokio::test] + async fn repeated_muse_import_does_not_emit_a_second_trace() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("muse-export.json"); + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-ledger"}], + "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":"terminal","terminal":"completed"}}}}, + {"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}} + ] + })) + .unwrap(), + ) + .unwrap(); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let import = || { + import_transcript_with_ledger( + &export, + ImportSource::Muse, + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + false, + Some(ledger.clone()), + ) + }; + import().await.unwrap(); + let rows = output.join("spans/muse-ledger.ndjson"); + let first = std::fs::read_to_string(&rows).unwrap(); + assert!(first.contains("Muse Code")); + import().await.unwrap(); + assert_eq!(std::fs::read_to_string(&rows).unwrap(), first); + } + fn test_run_hook_command() -> RunHookCommand { RunHookCommand { program: OsString::from("/opt/Braintrust CLI/bt"), diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index 17dafd67..d0f0476e 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -391,8 +391,11 @@ where Ok(()) } - fn finish(mut self) -> Result<()> { + fn finish(mut self, attach_snapshot: bool) -> Result { if self.activity_count == 0 { + if attach_snapshot { + return Ok(self.current_session_ended); + } bail!( "Muse export {} has no traceable run events", self.path.display() @@ -402,6 +405,9 @@ where .last_ts .ok_or_else(|| anyhow!("Muse export {} has no timestamps", self.path.display()))?; let current_session_ended = self.current_session_ended; + if attach_snapshot && !current_session_ended { + return Ok(false); + } for run_id in self.runs.keys().cloned().collect::>() { self.finish_model(&run_id, end_ts, Some("Run ended without a terminal event"))?; let run = self.runs.remove(&run_id).expect("run exists"); @@ -430,7 +436,8 @@ where None }, "source": "muse_export_v1", - })) + }))?; + Ok(current_session_ended) } } @@ -459,6 +466,20 @@ pub(crate) fn for_each_envelope_for_session( expected_id: &str, emit: F, ) -> Result<()> +where + F: FnMut(Envelope) -> Result<()>, +{ + for_each_snapshot_envelope_for_session(path, expected_id, false, emit).map(|_| ()) +} + +/// Parse a complete export without manufacturing terminal events for an +/// active attach snapshot. Returns whether Muse actually ended the session. +pub(crate) fn for_each_snapshot_envelope_for_session( + path: &Path, + expected_id: &str, + attach_snapshot: bool, + emit: F, +) -> Result where F: FnMut(Envelope) -> Result<()>, { @@ -500,7 +521,7 @@ where .with_context(|| format!("parse Muse export {}", path.display()))?; de.end() .with_context(|| format!("parse Muse export {}", path.display()))?; - reader.finish() + reader.finish(attach_snapshot) } fn read_header(path: &Path) -> Result
{ @@ -718,4 +739,41 @@ mod tests { "Session has no terminal end after its latest activity" ); } + + #[test] + fn active_snapshot_does_not_invent_terminal_events() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "sessions": [{"session_id":"session-1"}], + "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"}}}} + ] + })); + let mut events = Vec::new(); + let ended = for_each_snapshot_envelope_for_session(&path, "session-1", true, |event| { + events.push(event.event); + Ok(()) + }) + .unwrap(); + assert!(!ended); + assert_eq!(events, ["SessionStart", "UserPromptSubmit", "PreLLMCall"]); + } + + #[test] + fn empty_active_snapshot_waits_for_first_run() { + let (_temp, path) = write_export(json!({ + "export_schema_version": 1, + "sessions": [{"session_id":"session-1"}], + "events": [] + })); + let mut events = Vec::new(); + let ended = for_each_snapshot_envelope_for_session(&path, "session-1", true, |event| { + events.push(event); + Ok(()) + }) + .unwrap(); + assert!(!ended); + assert!(events.is_empty()); + } } diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 9d94ba03..970e4f11 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -57,6 +57,80 @@ async fn importing_muse_export_uses_the_production_translator() { })); } +#[tokio::test] +async fn attaching_muse_export_waits_for_native_end_without_reinserting_root() { + let tmp = tempfile::tempdir().unwrap(); + let export = tmp.path().join("muse-export.json"); + let output = tmp.path().join("spans"); + let started = json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}); + let write_export = |events: Vec| { + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-follow"}], + "events":events, + })) + .unwrap(), + ) + .unwrap(); + }; + write_export(vec![started.clone()]); + let export_for_attach = export.clone(); + let output_for_attach = output.clone(); + let attach = tokio::spawn(async move { + import_transcript( + &export_for_attach, + ImportSource::Muse, + options(&output_for_attach), + None, + true, + ) + .await + }); + let rows_path = output.join("muse-follow.ndjson"); + tokio::time::timeout(std::time::Duration::from_secs(3), async { + loop { + if std::fs::read_to_string(&rows_path) + .is_ok_and(|contents| contents.contains("Muse Code")) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(!attach.is_finished()); + write_export(vec![ + started, + json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), + json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}}), + ]); + let summaries = tokio::time::timeout(std::time::Duration::from_secs(4), attach) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(summaries.len(), 1); + assert!(summaries[0].finalized); + let output_rows = rows(&rows_path); + assert_eq!( + output_rows + .iter() + .filter(|row| row.pointer("/Insert/name").and_then(Value::as_str) == Some("Muse Code")) + .count(), + 1 + ); + assert_eq!( + output_rows + .iter() + .filter(|row| row.pointer("/Merge/end_ms").and_then(Value::as_i64) == Some(1002)) + .count(), + 1 + ); +} + #[tokio::test] async fn importing_sparse_failed_muse_export_preserves_model_and_error() { let tmp = tempfile::tempdir().unwrap(); From bece2a03b03f1446b6b2f7d31ab3ca90f5ca1e2d Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 24 Sep 2026 14:44:55 -0400 Subject: [PATCH 15/17] Keep active Muse imports open for later replay --- bt-daemon/README.md | 2 + bt-daemon/src/lib.rs | 121 ++++++++++++++++++++++++++++++++++---- bt-daemon/tests/replay.rs | 3 +- 3 files changed, 114 insertions(+), 12 deletions(-) diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 053ed21d..6ac1a269 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -298,6 +298,8 @@ export each poll, so polling cost grows with long sessions. MSP offers cursor-paged `view/page`, but its view events omit the model request/response records needed by this translator; `view/subscribe` also requires a session loaded on the same host. +One-shot imports of active Muse sessions also leave spans open, allowing a +later import to add the native completion to the same deterministic trace. `run muse` uses a process-local managed-hook override in the verified Muse Code 1.1.1 release. It retains compatible unrelated managed hooks, replaces Braintrust's saved managed hooks only for that invocation, and leaves ordinary Muse sessions and diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index d60850d5..a4b3b0c9 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -1147,8 +1147,13 @@ async fn import_muse_sessions_with_ledger( .tempdir()?; let export = directory.path().join("session.json"); export_muse_session(session_id, &export).await?; - process_muse_export(&mut processor, &export, session_id).await?; - if let Some(summary) = processor.finish_session(session_id).await? { + let ended = process_muse_export(&mut processor, &export, session_id).await?; + let summary = if ended { + processor.finish_session(session_id).await? + } else { + processor.detach_session(session_id).await? + }; + if let Some(summary) = summary { summaries.push(summary); } } @@ -1160,10 +1165,10 @@ async fn process_muse_export( processor: &mut ImportProcessor, export: &std::path::Path, session_id: &str, -) -> anyhow::Result<()> { - process_muse_snapshot(processor, export, session_id, false, 0) +) -> anyhow::Result { + process_muse_snapshot(processor, export, session_id, true, 0) .await - .map(|_| ()) + .map(|(_, ended)| ended) } /// The file is fully parsed on every poll. Only envelopes beyond the previous @@ -1839,9 +1844,14 @@ async fn import_transcript_with_ledger( return attach_muse_export(file, &session_id, false, opts, config, ledger_dir).await; } let mut processor = ImportProcessor::new(opts, config, ledger_dir); - process_muse_export(&mut processor, file, &session_id).await?; + let ended = process_muse_export(&mut processor, file, &session_id).await?; let mut summaries = Vec::new(); - if let Some(summary) = processor.finish_session(&session_id).await? { + let summary = if ended { + processor.finish_session(&session_id).await? + } else { + processor.detach_session(&session_id).await? + }; + if let Some(summary) = summary { summaries.push(summary); } summaries.extend(processor.finish().await?); @@ -2117,7 +2127,7 @@ impl ImportProcessor { let ops = live.translator.flush(&live.ctx)?; Self::emit_translator_batches(&mut live, ops).await?; live.sink.flush().await?; - summaries.push(Self::summary(sid, live)); + summaries.push(Self::summary(sid, live, true)); } summaries.sort_by(|left, right| left.session_id.cmp(&right.session_id)); Ok(summaries) @@ -2138,16 +2148,27 @@ impl ImportProcessor { let ops = live.translator.flush(&live.ctx)?; Self::emit_translator_batches(&mut live, ops).await?; live.sink.flush().await?; - Ok(Some(Self::summary(session_id.to_string(), live))) + Ok(Some(Self::summary(session_id.to_string(), live, true))) + } + + /// A one-shot snapshot of an active session must not mark its spans + /// complete in the destination ledger. A later import can then extend the + /// same deterministic rows when Muse records the real terminal events. + async fn detach_session(&mut self, session_id: &str) -> anyhow::Result> { + let Some(mut live) = self.sessions.remove(session_id) else { + return Ok(None); + }; + live.sink.flush().await?; + Ok(Some(Self::summary(session_id.to_string(), live, false))) } - fn summary(session_id: String, live: ImportLive) -> ImportSummary { + fn summary(session_id: String, live: ImportLive, finalized: bool) -> ImportSummary { ImportSummary { session_id, destination: live.destination, root_span_id: live.root_span_id, span_count: live.span_ids.len(), - finalized: true, + finalized, } } } @@ -3032,6 +3053,84 @@ mod tests { assert_eq!(std::fs::read_to_string(&rows).unwrap(), first); } + #[tokio::test] + async fn later_muse_import_completes_an_active_snapshot_on_the_same_root() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("muse-export.json"); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let started = json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}); + let write = |events: Vec| { + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-active-import"}], + "events":events, + })) + .unwrap(), + ) + .unwrap(); + }; + let import = || { + import_transcript_with_ledger( + &export, + ImportSource::Muse, + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + false, + Some(ledger.clone()), + ) + }; + write(vec![started.clone()]); + let summaries = import().await.unwrap(); + assert!(!summaries[0].finalized); + let rows_path = output.join("spans/muse-active-import.ndjson"); + let first = std::fs::read_to_string(&rows_path).unwrap(); + assert!(!first.contains("\"Merge\"")); + + write(vec![ + started, + json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), + json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}}), + ]); + let summaries = import().await.unwrap(); + assert!(summaries[0].finalized); + let rows = std::fs::read_to_string(&rows_path).unwrap(); + let parsed = rows + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + let root_id = ids::span_id( + &ids::session_namespace("muse", "muse-active-import"), + "root", + ); + assert_eq!( + parsed + .iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(|value| value.as_str()) + == Some(root_id.as_str())) + .count(), + 2 + ); + assert_eq!( + parsed + .iter() + .filter(|row| row + .pointer("/Merge/span_id") + .and_then(|value| value.as_str()) + == Some(root_id.as_str())) + .count(), + 1 + ); + assert!(rows.contains("\"end_ms\":1002")); + let completed = rows.clone(); + import().await.unwrap(); + assert_eq!(std::fs::read_to_string(&rows_path).unwrap(), completed); + } + fn test_run_hook_command() -> RunHookCommand { RunHookCommand { program: OsString::from("/opt/Braintrust CLI/bt"), diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 970e4f11..6d2b406e 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -38,7 +38,8 @@ async fn importing_muse_export_uses_the_production_translator() { {"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","terminal":"completed"}}}} + {"recorded_at": 1_004_000, "envelope": {"payload": {"kind":"run", "run_id":"run-1", "event":{"kind":"terminal","terminal":"completed"}}}}, + {"recorded_at": 1_005_000, "envelope": {"payload": {"kind":"session_end"}}} ] })).unwrap()).unwrap(); let output = tmp.path().join("spans"); From b3364858b9c7acd4defeb4f00a5a3b7cb47339cc Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Sat, 17 Oct 2026 02:34:33 +0800 Subject: [PATCH 16/17] Adapt Muse integration to current agent capabilities --- bt-daemon/src/agents/mod.rs | 3 + bt-daemon/src/agents/muse.rs | 25 + bt-daemon/src/args.rs | 9 + bt-daemon/src/client.rs | 24 +- bt-daemon/src/hook.rs | 65 +- bt-daemon/src/lib.rs | 3284 +---------------- bt-daemon/src/main.rs | 4 +- bt-daemon/src/managed_run/mod.rs | 79 +- bt-daemon/src/managed_run/muse.rs | 97 + bt-daemon/src/setup.rs | 2277 ------------ bt-daemon/src/setup/common.rs | 2 +- bt-daemon/src/setup/mod.rs | 1 + bt-daemon/src/setup/muse.rs | 603 +++ bt-daemon/src/trace_command.rs | 135 +- bt-daemon/src/trace_runtime.rs | 22 +- bt-daemon/src/transcript_import/mod.rs | 751 ++-- bt-daemon/src/transcript_import/muse.rs | 20 +- .../src/transcript_import/muse_runtime.rs | 563 +++ bt-daemon/src/transcript_import/pipeline.rs | 276 +- bt-daemon/src/translate/mod.rs | 142 +- bt-daemon/src/translate/muse.rs | 45 +- bt-daemon/tests/pipeline.rs | 5 +- 22 files changed, 2522 insertions(+), 5910 deletions(-) create mode 100644 bt-daemon/src/agents/muse.rs create mode 100644 bt-daemon/src/managed_run/muse.rs delete mode 100644 bt-daemon/src/setup.rs create mode 100644 bt-daemon/src/setup/muse.rs create mode 100644 bt-daemon/src/transcript_import/muse_runtime.rs diff --git a/bt-daemon/src/agents/mod.rs b/bt-daemon/src/agents/mod.rs index cc86aed6..eca23c15 100644 --- a/bt-daemon/src/agents/mod.rs +++ b/bt-daemon/src/agents/mod.rs @@ -10,6 +10,7 @@ mod claude; mod codex; mod cursor; mod grok; +mod muse; mod opencode; mod pi; @@ -18,6 +19,7 @@ pub(crate) use claude::Claude; pub(crate) use codex::Codex; pub(crate) use cursor::Cursor; pub(crate) use grok::Grok; +pub(crate) use muse::Muse; pub(crate) use opencode::OpenCode; pub(crate) use pi::Pi; @@ -103,6 +105,7 @@ pub(crate) fn registrar() -> &'static Registrar { Arc::new(Codex).register(&mut registrar); Arc::new(Cursor).register(&mut registrar); Arc::new(Grok).register(&mut registrar); + Arc::new(Muse).register(&mut registrar); Arc::new(OpenCode).register(&mut registrar); Arc::new(Pi).register(&mut registrar); registrar diff --git a/bt-daemon/src/agents/muse.rs b/bt-daemon/src/agents/muse.rs new file mode 100644 index 00000000..559ca253 --- /dev/null +++ b/bt-daemon/src/agents/muse.rs @@ -0,0 +1,25 @@ +use super::{Agent, Identity, Register, Registrar}; +use std::sync::Arc; + +pub(crate) struct Muse; + +pub(crate) const IDENTITY: Identity = Identity { + id: "muse", + source: "muse", + aliases: &["muse-code"], + display_name: "Muse Code", +}; + +impl Agent for Muse { + fn identity(&self) -> &'static Identity { + &IDENTITY + } +} + +impl Register for Muse { + fn register(self: Arc, registrar: &mut Registrar) { + registrar.setup.insert(self.clone()); + registrar.run.insert(self.clone()); + registrar.translate.insert(self); + } +} diff --git a/bt-daemon/src/args.rs b/bt-daemon/src/args.rs index 1c8e0a4e..f57100a4 100644 --- a/bt-daemon/src/args.rs +++ b/bt-daemon/src/args.rs @@ -85,6 +85,9 @@ pub struct HookArgs { /// not carry this flag and are suppressed for the managed child. #[arg(long, hide = true)] pub managed_run_hook: bool, + /// Private route context for Muse, whose hooks clear their environment. + #[arg(long, hide = true)] + pub managed_context: Option, } /// Arguments for `status`. @@ -179,6 +182,8 @@ pub struct ImportArgs { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum ImportSource { Codex, + #[value(alias = "muse-code")] + Muse, Cursor, #[value(name = "claude", alias = "claude-code")] Claude, @@ -199,6 +204,7 @@ impl ImportSource { pub(crate) fn identity(self) -> &'static Identity { match self { Self::Codex => agents::Codex.identity(), + Self::Muse => agents::Muse.identity(), Self::Cursor => agents::Cursor.identity(), Self::Claude => agents::Claude.identity(), Self::Antigravity => agents::Antigravity.identity(), @@ -210,6 +216,7 @@ impl RunSource { pub(crate) fn identity(self) -> &'static Identity { match self { Self::Codex => agents::Codex.identity(), + Self::Muse => agents::Muse.identity(), Self::Cursor => agents::Cursor.identity(), Self::Claude => agents::Claude.identity(), Self::OpenCode => agents::OpenCode.identity(), @@ -269,6 +276,8 @@ pub struct RunHookCommand { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum RunSource { Codex, + #[value(alias = "muse-code")] + Muse, Cursor, #[value(name = "claude", alias = "claude-code")] Claude, diff --git a/bt-daemon/src/client.rs b/bt-daemon/src/client.rs index e0b8dc34..e9aa8a36 100644 --- a/bt-daemon/src/client.rs +++ b/bt-daemon/src/client.rs @@ -151,6 +151,15 @@ pub async fn ensure_daemon( socket: &Path, host: &HostInfo, no_spawn: bool, +) -> anyhow::Result { + ensure_daemon_in(socket, host, no_spawn, None).await +} + +pub(crate) async fn ensure_daemon_in( + socket: &Path, + host: &HostInfo, + no_spawn: bool, + data_dir: Option<&Path>, ) -> anyhow::Result { if let Ok(s) = connect(socket).await { return Ok(s); @@ -158,7 +167,7 @@ pub async fn ensure_daemon( if no_spawn { anyhow::bail!("no daemon at {} and --no-spawn is set", socket.display()); } - spawn_daemon(host, socket)?; + spawn_daemon(host, socket, data_dir)?; let deadline = tokio::time::Instant::now() + Duration::from_secs(5); loop { tokio::time::sleep(Duration::from_millis(20)).await; @@ -172,13 +181,15 @@ pub async fn ensure_daemon( anyhow::bail!("daemon did not come up at {}", socket.display()) } -fn spawn_daemon(host: &HostInfo, socket: &Path) -> anyhow::Result<()> { +fn spawn_daemon(host: &HostInfo, socket: &Path, data_dir_arg: Option<&Path>) -> anyhow::Result<()> { let (exe, rest) = host .serve_argv .split_first() .ok_or_else(|| anyhow::anyhow!("empty serve_argv"))?; - let data_dir = crate::paths::data_dir(None); + let data_dir = data_dir_arg + .map(Path::to_path_buf) + .unwrap_or_else(|| crate::paths::data_dir(None)); let _ = crate::paths::ensure_private_dir(&data_dir); let log = std::fs::OpenOptions::new() .create(true) @@ -191,6 +202,10 @@ fn spawn_daemon(host: &HostInfo, socket: &Path) -> anyhow::Result<()> { let mut args = rest.to_vec(); args.push("--socket".into()); args.push(socket.as_os_str().to_owned()); + if let Some(data_dir) = data_dir_arg { + args.push("--data-dir".into()); + args.push(data_dir.as_os_str().to_owned()); + } // A hook's stdout/stderr are pipes owned by the agent. Windows normally // inherits every inheritable handle, even when the daemon's standard // streams are redirected to serve.log. The daemon must not keep those @@ -204,6 +219,9 @@ fn spawn_daemon(host: &HostInfo, socket: &Path) -> anyhow::Result<()> { let mut cmd = crate::subprocess::detached_daemon_command(exe); cmd.args(rest); cmd.arg("--socket").arg(socket); + if let Some(data_dir) = data_dir_arg { + cmd.arg("--data-dir").arg(data_dir); + } cmd.stdin(Stdio::null()); match log { Some(f) => { diff --git a/bt-daemon/src/hook.rs b/bt-daemon/src/hook.rs index 099ffd0e..53e58856 100644 --- a/bt-daemon/src/hook.rs +++ b/bt-daemon/src/hook.rs @@ -96,6 +96,9 @@ pub async fn run_hook(args: HookArgs, route: SessionRoute, host: HostInfo) -> an if suppress_hook(&args) { return Ok(()); } + if args.managed_context.is_some() { + return run_hook_with_route(args, route, host).await; + } let settings = settings::AgentSettings::load_for_hook(&args.source)?; if !settings.tracing_enabled() { return Ok(()); @@ -185,6 +188,21 @@ async fn capture_hook( mut route: SessionRoute, host: HostInfo, ) -> anyhow::Result<()> { + let managed_context = if let Some(path) = &args.managed_context { + if crate::translate::canonical_source_name(&args.source) != "muse" || !args.managed_run_hook + { + anyhow::bail!("managed context is only valid for an injected Muse hook"); + } + let context: crate::managed_run::muse::MuseManagedContext = + serde_json::from_slice(&tokio::fs::read(path).await?)?; + if context.managed_run_id.is_empty() || context.route.destination.is_none() { + anyhow::bail!("Muse managed hook context is incomplete"); + } + route = context.route.clone(); + Some(context) + } else { + None + }; let mut payload = read_stdin_json(args.transcript_path_field.as_deref())?; resolve_dynamic_hook_versions(&mut args, &payload); @@ -208,10 +226,22 @@ async fn capture_hook( anyhow::bail!("trace destination is not configured for {}", args.source); } apply_additional_metadata(&mut route, args.additional_metadata.as_deref())?; - let env = build_hook_envelope(&args, route, payload, session_id, event); + let mut env = build_hook_envelope(&args, route, payload, session_id, event); + if let Some(context) = &managed_context { + env.managed_run_id = Some(context.managed_run_id.clone()); + } - let socket = paths::socket_path(args.socket.as_deref()); - forward_envelope(&env, &socket, &host, args.no_spawn).await?; + let socket = managed_context + .as_ref() + .map(|context| context.socket.clone()) + .unwrap_or_else(|| paths::socket_path(args.socket.as_deref())); + forward_envelope( + &env, + &socket, + &host, + args.no_spawn || managed_context.is_some(), + ) + .await?; Ok(()) } @@ -438,6 +468,35 @@ mod tests { use clap::Parser; use serde_json::json; + #[tokio::test] + async fn muse_managed_hook_rejects_missing_context_before_reading_stdin() { + let temp = tempfile::tempdir().unwrap(); + let args = HookCli::try_parse_from([ + "test", + "--source", + "muse", + "--managed-run-hook", + "--managed-context", + temp.path().join("missing.json").to_str().unwrap(), + ]) + .unwrap() + .args; + let error = run_hook( + args, + SessionRoute::default(), + HostInfo { + serve_argv: Vec::new(), + version: "test".into(), + }, + ) + .await + .unwrap_err(); + assert_eq!( + error.downcast_ref::().unwrap().kind(), + std::io::ErrorKind::NotFound + ); + } + #[derive(Debug, Parser)] struct HookCli { #[command(flatten)] diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index a4b3b0c9..d85d4f8c 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -12,39 +12,58 @@ pub mod paths; +mod agents; +mod args; mod client; mod command_output; mod correlation; mod delivery_ledger; mod dispatch; +mod env; +mod hook; mod ids; mod journal; mod managed_run; mod pipeline; mod plugin_diagnostics; pub(crate) mod process; +mod route; mod server; mod settings; mod setup; mod sink; mod span_processor; +mod subprocess; mod trace_command; mod trace_runtime; mod transcript_import; mod transcript_mirror; mod translate; mod transport; +#[cfg(windows)] +mod win_acl; pub mod wire; -pub use client::HostInfo; +pub use args::{ + HookArgs, ImportArgs, ImportSource, ParentObjectType, RunArgs, RunHookCommand, RunSource, + ServeArgs, StatusArgs, +}; +pub use client::{ + flush_managed_run, flush_session, forward_envelope, run_status, shutdown_daemon, HostInfo, +}; pub use command_output::{ - AuthDiagnostic, DoctorCommandOutput, ImportSummary, OutputFormat, SetupCommandOutput, - StatusCommandOutput, StopCommandOutput, TraceCommandOutput, + AuthDiagnostic, DaemonDiagnostic, DaemonStatus, DoctorCommandOutput, ImportSummary, + OutputFormat, SetupCommandOutput, StatusCommandOutput, StopCommandOutput, TraceCommandOutput, }; +pub use hook::run_hook; #[doc(hidden)] pub use journal::source_journal_path; +pub use managed_run::{run_traced, run_traced_with_host}; pub use plugin_diagnostics::PluginDiagnostic; -pub use server::{AuthLease, AuthProvider, AuthResolveReason, ServeOptions}; +pub use server::{ + braintrust_serve_options, debug_serve_options, run_serve, AuthLease, AuthProvider, + AuthResolveReason, ServeOptions, +}; pub use setup::{run_disable, run_enable, run_setup, run_update}; pub use sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, Sink, SinkFactory}; pub use trace_command::{ @@ -54,3259 +73,8 @@ pub use trace_command::{ pub use trace_runtime::{ run_trace, HostRouteSelection, RouteRequirements, TraceHostContext, TraceHostServices, }; +pub use transcript_import::{import_transcript, import_transcripts, run_import}; pub use translate::{ - AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, -}; - -use anyhow::Context; -use braintrust_sdk_rust::{SpanComponents, SpanObjectType}; -use clap::{Args, ValueEnum}; -use serde::Deserialize; -use std::ffi::OsString; -use std::path::PathBuf; -use std::process::Stdio; -use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::io::{ - AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines, -}; -use wire::{ - method, Envelope, ManagedRunFlushParams, SessionConfig, SessionRoute, StatusResult, - PROTOCOL_VERSION, + AgentTranslator, OriginSnapshot, Registry, SessionCtx, SpanOp, SpanRow, SpanType, + TURN_SPAN_ID_KEY, }; - -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); -const MUSE_MSP_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); -const MUSE_MSP_MAX_PAGES: usize = 1_000; -const MUSE_MSP_DIAGNOSTIC_BYTES: usize = 64 * 1024; - -/// Arguments for `serve`. -#[derive(Debug, Clone, Args)] -pub struct ServeArgs { - /// Socket path override (default: see docs/protocol.md). - #[arg(long)] - pub socket: Option, - /// Data/journal directory override. - #[arg(long)] - pub data_dir: Option, - /// Exit after this many seconds idle (no activity, empty queues). 0 - /// disables the watchdog. - #[arg(long, default_value_t = 300)] - pub idle_timeout_secs: u64, - /// Retire a session's in-memory state after this many seconds without - /// traffic. A later event rebuilds it from the journal. 0 disables - /// retirement. - #[arg(long, default_value_t = 30)] - pub session_idle_timeout_secs: u64, -} - -/// Arguments for `hook`. -#[derive(Debug, Clone, Args)] -pub struct HookArgs { - /// Which translator should interpret this event's payload. - #[arg(long)] - pub source: String, - /// Optional agent version, forwarded for payload-drift handling. - #[arg(long)] - pub source_version: Option, - /// Optional instrumentation package version, forwarded independently from - /// the coding agent version. - #[arg(long)] - pub plugin_version: Option, - /// Socket path override. - #[arg(long)] - pub socket: Option, - /// JSON field in the payload holding the session id. - #[arg(long, default_value = "session_id")] - pub session_id_field: String, - /// JSON field in the payload holding the event name. - #[arg(long, default_value = "hook_event_name")] - pub event_field: String, - /// Explicit event name (overrides `--event-field` lookup). - #[arg(long)] - pub event: Option, - /// JSON field holding a transcript path. When present, capture the file - /// length observed by this hook so deterministic journal replay cannot - /// read transcript records written by later lifecycle events. - #[arg(long)] - pub transcript_path_field: Option, - /// Fail instead of spawning a daemon if none is running. - #[arg(long)] - pub no_spawn: bool, - /// Ask the daemon to flush the session after a turn-ending event. The - /// flush is scheduled out-of-band; hook capture still returns immediately - /// after the durable journal write. - #[arg(long)] - pub flush_on_turn_end: bool, - /// Deprecated compatibility option. Hook capture never waits for daemon - /// translation, reporting, or flushing. - #[arg(long, default_value_t = 10_000, hide = true)] - pub flush_timeout_ms: u64, - /// JSON object merged into root-span metadata. Deliberately not read from - /// the environment: a hook fires automatically on every event, so its - /// configuration must come only from the persisted agent route (or, for a - /// managed child, the invocation settings `run` injected). - #[arg(long)] - pub additional_metadata: Option, - /// Marks the hook definition injected by `run`; inherited plugin hooks do - /// not carry this flag and are suppressed for the managed child. - #[arg(long, hide = true)] - pub managed_run_hook: bool, - /// Private invocation context for Muse's environment-sanitized hooks. - #[arg(long, hide = true)] - pub managed_context: Option, -} - -/// Arguments for `status`. -#[derive(Debug, Clone, Args)] -pub struct StatusArgs { - #[arg(long)] - pub socket: Option, - /// Limit to one session. - #[arg(long)] - pub session_id: Option, -} - -/// Arguments for importing a past coding-agent session. -#[derive(Debug, Clone, Args)] -pub struct ImportArgs { - /// Agent that produced the session. - #[arg(value_enum)] - pub source: ImportSource, - /// 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( - long, - value_name = "DESTINATION", - conflicts_with_all = ["parent", "parent_span_id", "parent_root_span_id", "parent_object_type", "parent_object_id", "parent_project"] - )] - pub destination: Option, - /// Attach below the opaque value returned by a Braintrust SDK span.export(). - #[arg( - long, - value_name = "SPAN_EXPORT", - conflicts_with_all = ["destination", "parent_span_id", "parent_root_span_id", "parent_object_type", "parent_object_id", "parent_project"] - )] - pub parent: Option, - /// Attach below this span ID. Also requires --parent-root-span-id and - /// --parent-object-type, plus --parent-object-id or --parent-project. - #[arg( - long, - value_name = "SPAN_ID", - requires_all = ["parent_root_span_id", "parent_object_type"], - conflicts_with_all = ["destination", "parent"] - )] - pub parent_span_id: Option, - /// Root span ID for --parent-span-id. - #[arg(long, value_name = "ROOT_SPAN_ID", requires = "parent_span_id")] - pub parent_root_span_id: Option, - /// Braintrust object type for --parent-span-id. - #[arg(long, value_enum, requires = "parent_span_id")] - pub parent_object_type: Option, - /// Braintrust object ID for --parent-span-id. - #[arg( - long, - value_name = "OBJECT_ID", - requires = "parent_span_id", - conflicts_with = "parent_project" - )] - pub parent_object_id: Option, - /// Braintrust project name for a project-logs parent. - #[arg( - long, - value_name = "PROJECT", - requires = "parent_span_id", - conflicts_with = "parent_object_id" - )] - pub parent_project: Option, - /// Keep following the transcript until Ctrl-C, importing new turns as the - /// coding-agent session grows. - #[arg(long, conflicts_with = "all")] - pub attach: bool, - /// JSON object merged into every imported root span's metadata. - #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] - pub additional_metadata: Option, - /// Tag applied to each imported root span. May be repeated or comma-separated. - #[arg(long = "tag", env = "BRAINTRUST_TAGS", value_delimiter = ',')] - pub tags: Vec, - /// JavaScript span transform for this import. Repeat to compose an isolated - /// transform chain; persistent setup plugins are not included. - #[arg(long, value_name = "PATH")] - pub plugin: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] -pub enum ImportSource { - Codex, - Muse, - #[value(name = "claude", alias = "claude-code")] - Claude, - #[value(name = "antigravity", alias = "agy")] - Antigravity, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] -pub enum ParentObjectType { - Experiment, - #[value(name = "project_logs", alias = "project-logs")] - ProjectLogs, - #[value(name = "playground_logs", alias = "playground-logs")] - PlaygroundLogs, -} - -impl From for SpanObjectType { - fn from(value: ParentObjectType) -> Self { - match value { - ParentObjectType::Experiment => Self::Experiment, - ParentObjectType::ProjectLogs => Self::ProjectLogs, - ParentObjectType::PlaygroundLogs => Self::PlaygroundLogs, - } - } -} - -impl ImportArgs { - pub fn has_destination_override(&self) -> bool { - self.destination.is_some() || self.parent.is_some() || self.parent_span_id.is_some() - } -} - -/// Arguments for launching a coding agent with invocation-local live hooks. -#[derive(Debug, Clone, Args)] -#[command(trailing_var_arg = true)] -pub struct RunArgs { - /// Coding agent to launch. - #[arg(value_enum)] - pub source: RunSource, - /// JSON object merged into root-span metadata for this invocation. - #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] - pub additional_metadata: Option, - /// Tag applied to each root span for this invocation. May be repeated or comma-separated. - #[arg(long = "tag", env = "BRAINTRUST_TAGS", value_delimiter = ',')] - pub tags: Vec, - /// JavaScript span transform for this invocation. Repeat to compose an - /// isolated transform chain; persistent setup plugins are not included. - #[arg(long, value_name = "PATH")] - pub plugin: Vec, - /// Arguments forwarded verbatim to the coding agent. - #[arg(allow_hyphen_values = true)] - pub agent_args: Vec, -} - -/// Front-end command used by a managed agent run to forward one hook payload. -/// -/// The standalone binary uses `[bt-daemon, hook]`; the embedded `bt` front-end -/// uses its own equivalent prefix. `run_traced` appends `--source `. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct RunHookCommand { - pub program: OsString, - pub args: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] -pub enum RunSource { - Codex, - Muse, - #[value(name = "claude", alias = "claude-code")] - Claude, - #[value(name = "opencode", alias = "open-code")] - OpenCode, - Pi, -} - -/// Run the daemon until shutdown. -pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { - server::run(args, opts).await -} - -fn build_hook_envelope( - args: &HookArgs, - route: SessionRoute, - payload: serde_json::Value, - session_id: String, - event: String, -) -> Envelope { - Envelope { - source: args.source.clone(), - source_version: args.source_version.clone(), - plugin_version: args.plugin_version.clone(), - session_id, - event, - ts_ms: now_ms(), - managed_run_id: std::env::var(MANAGED_RUN_ID_ENV) - .ok() - .filter(|value| !value.is_empty()), - capture: None, - payload, - route: Some(route), - config: None, - } -} - -#[derive(Debug, serde::Serialize, serde::Deserialize)] -struct MuseManagedContext { - route: SessionRoute, - managed_run_id: String, - socket: PathBuf, -} - -pub(crate) fn should_flush_hook_event(event: &str, flush_on_turn_end: bool) -> bool { - matches!(event, "SessionEnd" | "session_end") - || (flush_on_turn_end - && matches!( - event, - "Stop" - | "stop" - | "StopFailure" - | "stop_failure" - | "StopCancelled" - | "stop_cancelled" - | "SubagentStop" - | "subagent_stop" - )) -} - -pub(crate) fn should_flush_ingress_event(env: &wire::Envelope) -> bool { - let native = env.payload.get("event").unwrap_or(&env.payload); - let flush_on_turn_end = matches!( - env.route.as_ref().map(|route| route.flush_mode), - Some(wire::FlushMode::FlushOnTurnEnd) - ); - should_flush_hook_event(&env.event, flush_on_turn_end) - || match env.source.as_str() { - "pi" => match env.event.as_str() { - // Preserve Pi's explicit lifecycle flushes even in batched mode. - "session_shutdown" | "session_compact" | "session_tree" => true, - // OMP can end an attempt while keeping the same turn open for retry. - "agent_end" => { - flush_on_turn_end - && native.get("willRetry").and_then(serde_json::Value::as_bool) - != Some(true) - } - _ => false, - }, - "opencode" => matches!( - env.event.as_str(), - "session.idle" | "session.deleted" | "session.error" | "server.instance.disposed" - ), - _ => false, - } -} - -/// Capture one hook event from `stdin` and forward it to the daemon. -/// -/// `route` contains only non-secret profile and destination selection. -/// Returns `Ok` once the daemon has durably journaled the event. Callers that -/// must never fail the agent's turn should treat any `Err` as non-fatal and -/// exit 0. -pub async fn run_hook( - mut args: HookArgs, - mut route: SessionRoute, - host: HostInfo, -) -> anyhow::Result<()> { - let managed_context = args - .managed_context - .as_ref() - .map(|path| -> anyhow::Result { - if args.source != "muse" || !args.managed_run_hook { - anyhow::bail!("managed context is only valid for an injected Muse hook"); - } - let context: MuseManagedContext = serde_json::from_slice(&std::fs::read(path)?)?; - if context.managed_run_id.is_empty() || context.route.destination.is_none() { - anyhow::bail!("Muse managed hook context is incomplete"); - } - Ok(context) - }) - .transpose()?; - // A managed run injects its own hook definitions. Suppress an inherited - // Braintrust plugin hook for the same child, but allow the injected hook - // process, which carries the second marker. - if std::env::var_os("_BT_TRACE_MANAGED_RUN").is_some() && !args.managed_run_hook { - return Ok(()); - } - let settings = managed_context - .is_none() - .then(|| settings::AgentSettings::load(&args.source)); - if settings - .as_ref() - .is_some_and(|settings| !settings.tracing_enabled()) - { - return Ok(()); - } - let mut payload = read_stdin_json()?; - - resolve_dynamic_hook_versions(&mut args, &payload); - - if let Some(field) = &args.transcript_path_field { - add_transcript_observation(&mut payload, field); - } - - let session_id = json_str_field(&payload, &args.session_id_field) - .ok_or_else(|| anyhow::anyhow!("no `{}` field in hook payload", args.session_id_field))?; - let event = args - .event - .clone() - .or_else(|| json_str_field(&payload, &args.event_field)) - .unwrap_or_default(); - - if let Some(context) = &managed_context { - route = context.route.clone(); - } else if let Some(configured_route) = settings.and_then(|settings| settings.route) { - route = configured_route; - } - if args.flush_on_turn_end { - route.flush_mode = wire::FlushMode::FlushOnTurnEnd; - } - apply_additional_metadata(&mut route, args.additional_metadata.as_deref())?; - let mut env = build_hook_envelope(&args, route, payload, session_id, event); - if let Some(context) = &managed_context { - env.managed_run_id = Some(context.managed_run_id.clone()); - } - - let socket = managed_context - .as_ref() - .map(|context| context.socket.clone()) - .unwrap_or_else(|| paths::socket_path(args.socket.as_deref())); - forward_envelope( - &env, - &socket, - &host, - args.no_spawn || managed_context.is_some(), - ) - .await?; - - Ok(()) -} - -fn resolve_dynamic_hook_versions(args: &mut HookArgs, payload: &serde_json::Value) { - if args.source_version.is_none() { - args.source_version = source_version_from_env(&args.source) - .or_else(|| source_version_from_payload(&args.source, payload)); - } - if args.plugin_version.is_none() { - args.plugin_version = plugin_version_from_plugin_root(&args.source); - } -} - -fn source_version_from_env(source: &str) -> Option { - let key = match source { - "claude-code" => "CLAUDE_CODE_VERSION", - "codex" => "CODEX_VERSION", - "grok" => "GROK_VERSION", - "antigravity" => "AGY_VERSION", - _ => return None, - }; - std::env::var(key) - .ok() - .filter(|value| !value.trim().is_empty()) -} - -fn source_version_from_payload(source: &str, payload: &serde_json::Value) -> Option { - let fields: &[&str] = match source { - "codex" => &["cli_version", "version"], - "claude-code" | "grok" | "antigravity" => &["version", "cli_version"], - _ => &[], - }; - fields - .iter() - .find_map(|field| json_str_field(payload, field)) -} - -fn plugin_version_from_plugin_root(source: &str) -> Option { - let root = match source { - "claude-code" => std::env::var_os("CLAUDE_PLUGIN_ROOT"), - "codex" => std::env::var_os("PLUGIN_ROOT"), - "grok" => std::env::var_os("GROK_PLUGIN_ROOT"), - // Antigravity runs plugin hooks with the plugin root as cwd. - "antigravity" => std::env::current_dir() - .ok() - .map(|path| path.into_os_string()), - _ => None, - }?; - let manifest = match source { - "claude-code" => ".claude-plugin/plugin.json", - "codex" => ".codex-plugin/plugin.json", - "grok" => ".grok-plugin/plugin.json", - "antigravity" => "plugin.json", - _ => return None, - }; - plugin_version_from_manifest(&PathBuf::from(root).join(manifest)) -} - -fn plugin_version_from_manifest(path: &std::path::Path) -> Option { - serde_json::from_slice::(&std::fs::read(path).ok()?) - .ok()? - .get("version")? - .as_str() - .map(str::to_owned) -} - -/// Apply one invocation-local JSON metadata override to a non-secret route. -/// -/// The route is then carried unchanged through live hooks, managed runs, and -/// transcript import. Keeping validation here gives every public command the -/// same contract and prevents individual agent shims from parsing JSON. -pub(crate) fn apply_additional_metadata( - route: &mut SessionRoute, - additional_metadata: Option<&str>, -) -> anyhow::Result<()> { - let Some(metadata) = additional_metadata else { - return Ok(()); - }; - let value: serde_json::Value = serde_json::from_str(metadata) - .map_err(|e| anyhow::anyhow!("invalid --additional-metadata JSON: {e}"))?; - if !value.is_object() { - anyhow::bail!("--additional-metadata must be a JSON object"); - } - route.additional_metadata = Some(value); - Ok(()) -} - -/// Apply invocation-local root-span tags to a route. Tags are normalized once -/// at the CLI boundary so hook shims and translators only receive valid values. -pub(crate) fn apply_tags(route: &mut SessionRoute, tags: &[String]) -> anyhow::Result<()> { - if tags.is_empty() { - return Ok(()); - } - let mut normalized = Vec::with_capacity(tags.len()); - for tag in tags { - let tag = tag.trim(); - if tag.is_empty() { - anyhow::bail!("--tag must not be empty"); - } - if !normalized.iter().any(|existing| existing == tag) { - normalized.push(tag.to_string()); - } - } - route.tags = normalized; - Ok(()) -} - -fn initialize_params(env: &Envelope, daemon_version: &str) -> serde_json::Value { - serde_json::json!({ - "protocol_version": PROTOCOL_VERSION, - "client": { - "source": env.source, - "daemon_version": daemon_version, - "plugin_version": env.plugin_version, - "pid": std::process::id() - } - }) -} - -/// Ensure a daemon is up and forward one already-built [`Envelope`] to it -/// (`initialize` handshake + `event.log`). Also the seam in-process clients and -/// tests use to send events without going through stdin. -pub async fn forward_envelope( - env: &Envelope, - socket: &std::path::Path, - host: &HostInfo, - no_spawn: bool, -) -> anyhow::Result<()> { - const MAX_HANDOVER_ATTEMPTS: usize = 3; - let mut handover_attempts = 0; - let mut capture_retry_deadline = None; - - loop { - if capture_retry_deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) { - anyhow::bail!( - "replacement daemon did not accept the event at {}", - socket.display() - ); - } - - let stream = match client::connect(socket).await { - Ok(stream) => stream, - Err(_) if capture_retry_deadline.is_some() && no_spawn => { - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - continue; - } - Err(_) => client::ensure_daemon(socket, host, no_spawn).await?, - }; - let mut conn = client::Conn::new(stream); - let initialized = match conn - .request(method::INITIALIZE, initialize_params(env, &host.version)) - .await - { - Ok(initialized) => initialized, - Err(_) if capture_retry_deadline.is_some() => { - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - continue; - } - Err(error) => return Err(error), - }; - let initialized: wire::InitializeResult = serde_json::from_value(initialized)?; - if daemon_needs_upgrade(&initialized.daemon_version, &host.version) { - if no_spawn { - anyhow::bail!( - "daemon version {} is older than client {} and --no-spawn is set", - initialized.daemon_version, - host.version - ); - } - if handover_attempts == MAX_HANDOVER_ATTEMPTS { - anyhow::bail!( - "daemon version {} is still older than client {} after {MAX_HANDOVER_ATTEMPTS} handover attempts", - initialized.daemon_version, - host.version - ); - } - conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) - .await?; - handover_attempts += 1; - drop(conn); - wait_for_daemon_exit(socket).await; - continue; - } - let result = match conn.request(method::EVENT_LOG, env).await { - Ok(result) => result, - Err(_) if capture_retry_deadline.is_some() => { - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - continue; - } - Err(error) => return Err(error), - }; - let result: wire::EventLogResult = serde_json::from_value(result)?; - if result.accepted { - return Ok(()); - } - capture_retry_deadline - .get_or_insert_with(|| tokio::time::Instant::now() + std::time::Duration::from_secs(5)); - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - } -} - -async fn wait_for_daemon_exit(socket: &std::path::Path) { - for _ in 0..100 { - if client::connect(socket).await.is_err() { - return; - } - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - } -} - -/// A shared daemon only moves forward. Released versions use semver; opaque -/// development versions retain the old exact-match handover behavior. -fn daemon_needs_upgrade(daemon_version: &str, client_version: &str) -> bool { - if daemon_version == client_version { - return false; - } - compare_daemon_versions(daemon_version, client_version).is_none_or(|ordering| ordering.is_lt()) -} - -pub(crate) fn compare_daemon_versions(left: &str, right: &str) -> Option { - Some( - semver::Version::parse(left) - .ok()? - .cmp(&semver::Version::parse(right).ok()?), - ) -} - -/// Ask the daemon to flush a session, bounded by `timeout_ms`. A reliable -/// barrier: on return, every event enqueued before the call has been processed -/// and its spans emitted to the sink. -pub async fn flush_session( - session_id: &str, - socket: &std::path::Path, - timeout_ms: u64, -) -> anyhow::Result { - let stream = client::connect(socket).await?; - let mut conn = client::Conn::new(stream); - conn.request( - method::INITIALIZE, - serde_json::json!({ "protocol_version": PROTOCOL_VERSION, "client": { "source": "flush" } }), - ) - .await?; - let params = wire::FlushParams { - session_id: session_id.to_string(), - timeout_ms, - }; - let value = conn.request(method::SESSION_FLUSH, params).await?; - Ok(serde_json::from_value(value)?) -} - -/// Flush every daemon session accepted from one managed child process tree. -/// A missing daemon is not itself a failure: the daemon may have idle-exited -/// after draining every session, so acceptance is checked against the -/// persisted managed-run record instead. -pub async fn flush_managed_run( - managed_run_id: &str, - socket: &std::path::Path, - timeout_ms: u64, -) -> anyhow::Result { - flush_managed_run_in(managed_run_id, socket, timeout_ms, &paths::data_dir(None)).await -} - -async fn flush_managed_run_in( - managed_run_id: &str, - socket: &std::path::Path, - timeout_ms: u64, - data_dir: &std::path::Path, -) -> anyhow::Result { - let stream = match client::connect(socket).await { - Ok(stream) => stream, - Err(_) => { - let accepted_sessions = journal::read_managed_run_keys(data_dir, managed_run_id) - .await - .len() as u64; - return Ok(wire::FlushResult { - flushed: true, - pending: 0, - accepted_sessions, - }); - } - }; - let mut conn = client::Conn::new(stream); - conn.request( - method::INITIALIZE, - serde_json::json!({ - "protocol_version": PROTOCOL_VERSION, - "client": { "source": "managed-run-flush" } - }), - ) - .await?; - let params = ManagedRunFlushParams { - managed_run_id: managed_run_id.to_string(), - timeout_ms, - }; - let value = conn.request(method::MANAGED_RUN_FLUSH, params).await?; - Ok(serde_json::from_value(value)?) -} - -/// Query daemon status. `Ok(None)` means no daemon is running. -pub async fn run_status(args: StatusArgs) -> anyhow::Result> { - let socket = paths::socket_path(args.socket.as_deref()); - let stream = match client::connect(&socket).await { - Ok(s) => s, - Err(_) => return Ok(None), - }; - let mut conn = client::Conn::new(stream); - conn.request( - method::INITIALIZE, - serde_json::json!({ - "protocol_version": PROTOCOL_VERSION, - "client": { "source": "status" } - }), - ) - .await?; - let params = wire::StatusParams { - session_id: args.session_id.clone(), - }; - let value = conn.request(method::STATUS_GET, params).await?; - Ok(Some(serde_json::from_value(value)?)) -} - -/// Request a graceful daemon shutdown. Primarily useful for lifecycle -/// management and transport integration tests. -pub async fn shutdown_daemon(socket: &std::path::Path) -> anyhow::Result<()> { - let stream = client::connect(socket).await?; - let mut conn = client::Conn::new(stream); - conn.request(method::DAEMON_SHUTDOWN, serde_json::json!({})) - .await?; - Ok(()) -} - -/// Import a native coding-agent transcript through the normal translators and -/// sink. This is separate from daemon journal recovery: import creates traces -/// for a past session, while recovery rebuilds live correlation state. -pub async fn run_import( - args: ImportArgs, - opts: ServeOptions, - mut config: Option, -) -> anyhow::Result> { - validate_import_selection(&args)?; - apply_import_span_plugins(&mut config, &args.plugin)?; - let destination = import_parent_components(&args)? - .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 { - return attach_muse_session( - &args.session_ids[0], - opts, - config, - Some(paths::data_dir(None)), - ) - .await; - } - 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") - } - return import_muse_sessions_with_ledger( - &session_ids, - 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 { - return import_transcript_with_ledger( - &files[0], - args.source, - opts, - config, - true, - Some(ledger_dir), - ) - .await; - } - import_transcripts_with_ledger(&files, args.source, opts, config, Some(ledger_dir)).await -} - -fn apply_import_span_plugins( - config: &mut Option, - plugins: &[PathBuf], -) -> anyhow::Result<()> { - let plugins = resolve_span_plugin_paths(plugins)?; - if let Some(config) = config.as_mut() { - config.span_plugins = plugins; - } else if !plugins.is_empty() { - *config = Some(SessionConfig { - auth: wire::BackendAuth { - token: String::new(), - api_url: None, - app_url: None, - org_name: None, - org_id: None, - }, - destination: None, - flush_mode: wire::FlushMode::FireAndForget, - additional_metadata: None, - tags: Vec::new(), - span_plugins: plugins, - }); - } - Ok(()) -} - -async fn export_muse_session(session_id: &str, path: &std::path::Path) -> anyhow::Result<()> { - transcript_import::validate_session_id(session_id)?; - 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) - .args(["export", "--session", session_id, "--out"]) - .arg(path) - .stdout(Stdio::null()) - .stderr(Stdio::piped()) - .spawn() - .with_context(|| { - format!( - "run {} export for Muse session {session_id}", - executable.to_string_lossy() - ) - })?; - let stderr = child.stderr.take().expect("piped stderr"); - let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); - let status = tokio::time::timeout(Duration::from_secs(300), child.wait()) - .await - .with_context(|| format!("Muse export for session {session_id} timed out"))? - .with_context(|| format!("wait for Muse export of session {session_id}"))?; - let stderr = diagnostics - .await - .context("collect Muse export diagnostics")??; - if !status.success() { - anyhow::bail!( - "Muse export for session {session_id} failed: {}", - stderr.trim() - ); - } - Ok(()) -} - -/// 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> { - 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 stderr = child.stderr.take().expect("piped stderr"); - let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); - let mut reader = BufReader::new(stdout).lines(); - let listing = async { - 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 requires the initialized notification before other requests. - muse_msp_notify(&mut stdin, "initialized").await?; - let mut cursor = None; - let mut seen_cursors = std::collections::HashSet::new(); - let mut sessions = Vec::new(); - for page_index in 0..MUSE_MSP_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?; - match muse_msp_append_page(result, page_index + 1, &mut sessions, &mut seen_cursors)? { - None => { - sessions.sort(); - sessions.dedup(); - return Ok(sessions); - } - Some(next) => cursor = Some(next), - } - } - anyhow::bail!("Muse session/list exceeded {MUSE_MSP_MAX_PAGES} pages") - } - .await; - drop(stdin); - drop(reader); - let status = match tokio::time::timeout(MUSE_MSP_SHUTDOWN_TIMEOUT, child.wait()).await { - Ok(status) => status.context("wait for Muse MSP host")?, - Err(_) => { - child - .start_kill() - .context("stop unresponsive Muse MSP host")?; - child - .wait() - .await - .context("reap unresponsive Muse MSP host")?; - anyhow::bail!("Muse MSP host did not exit after stdin closed") - } - }; - let stderr = diagnostics - .await - .context("collect Muse MSP diagnostics")??; - let sessions = listing?; - if !status.success() { - anyhow::bail!("Muse MSP session listing failed: {}", stderr.trim()); - } - Ok(sessions) -} - -#[derive(Deserialize)] -struct MuseMspSessionPage { - sessions: Vec, - #[serde(rename = "nextCursor")] - next_cursor: Option, -} - -#[derive(Deserialize)] -struct MuseMspSessionEntry { - #[serde(rename = "sessionId", alias = "session_id")] - session_id: String, -} - -fn muse_msp_append_page( - result: serde_json::Value, - page_index: usize, - sessions: &mut Vec, - seen_cursors: &mut std::collections::HashSet, -) -> anyhow::Result> { - let page: MuseMspSessionPage = serde_json::from_value(result) - .with_context(|| format!("decode Muse session/list page {page_index}"))?; - for session in page.sessions { - transcript_import::validate_session_id(&session.session_id).with_context(|| { - format!("invalid Muse session id on session/list page {page_index}") - })?; - sessions.push(session.session_id); - } - if let Some(cursor) = &page.next_cursor { - if !seen_cursors.insert(cursor.clone()) { - anyhow::bail!("Muse session/list repeated cursor on page {page_index}") - } - } - Ok(page.next_cursor) -} - -async fn muse_drain_stderr(mut stderr: tokio::process::ChildStderr) -> anyhow::Result { - let mut output = Vec::new(); - let mut buffer = [0u8; 4096]; - loop { - let count = stderr.read(&mut buffer).await?; - if count == 0 { - break; - } - let remaining = MUSE_MSP_DIAGNOSTIC_BYTES.saturating_sub(output.len()); - output.extend_from_slice(&buffer[..count.min(remaining)]); - } - Ok(String::from_utf8_lossy(&output).into_owned()) -} - -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'); - tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { - stdin.write_all(frame.as_bytes()).await?; - stdin.flush().await?; - Ok(()) - }) - .await - .map_err(|_| anyhow::anyhow!("Muse MSP timed out sending {method}"))? -} - -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'); - tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { - stdin.write_all(request.as_bytes()).await?; - stdin.flush().await?; - 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("jsonrpc").and_then(serde_json::Value::as_str) != Some("2.0") { - anyhow::bail!( - "Muse MSP sent a frame without JSON-RPC 2.0 while waiting for {method}" - ) - } - let Some(response_id) = frame.get("id") else { - // Notifications can be interleaved with responses. - continue; - }; - if response_id.as_u64() != Some(id) { - anyhow::bail!("Muse MSP returned an unexpected response id for {method}") - } - 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_sessions_with_ledger( - session_ids: &[String], - opts: ServeOptions, - config: Option, - ledger_dir: Option, -) -> anyhow::Result> { - let mut processor = ImportProcessor::new(opts, config, ledger_dir); - let mut summaries = Vec::new(); - for session_id in session_ids { - let directory = tempfile::Builder::new() - .prefix("bt-muse-import-") - .tempdir()?; - let export = directory.path().join("session.json"); - export_muse_session(session_id, &export).await?; - let ended = process_muse_export(&mut processor, &export, session_id).await?; - let summary = if ended { - processor.finish_session(session_id).await? - } else { - processor.detach_session(session_id).await? - }; - if let Some(summary) = summary { - summaries.push(summary); - } - } - summaries.extend(processor.finish().await?); - Ok(summaries) -} - -async fn process_muse_export( - processor: &mut ImportProcessor, - export: &std::path::Path, - session_id: &str, -) -> anyhow::Result { - process_muse_snapshot(processor, export, session_id, true, 0) - .await - .map(|(_, ended)| ended) -} - -/// The file is fully parsed on every poll. Only envelopes beyond the previous -/// snapshot are sent to the long-lived translator and sink, so open model and -/// turn state is retained without replaying old starts into it. -async fn process_muse_snapshot( - processor: &mut ImportProcessor, - export: &std::path::Path, - session_id: &str, - attach_snapshot: bool, - skip_envelopes: usize, -) -> anyhow::Result<(usize, bool)> { - let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); - let export_for_parser = export.to_path_buf(); - let expected_id = session_id.to_owned(); - let parser = tokio::task::spawn_blocking(move || { - let mut batch = Vec::with_capacity(128); - let mut envelope_count = 0usize; - let ended = transcript_import::muse::for_each_snapshot_envelope_for_session( - &export_for_parser, - &expected_id, - attach_snapshot, - |entry| { - envelope_count += 1; - if envelope_count <= skip_envelopes { - return Ok(()); - } - batch.push(entry); - if batch.len() == 128 { - sender - .blocking_send(std::mem::take(&mut batch)) - .map_err(|_| { - anyhow::anyhow!("Muse import stopped while parsing {expected_id}") - })?; - } - Ok(()) - }, - )?; - if envelope_count < skip_envelopes { - anyhow::bail!( - "Muse export for {expected_id} shrank from {skip_envelopes} to {envelope_count} trace events" - ); - } - if !batch.is_empty() { - sender - .blocking_send(batch) - .map_err(|_| anyhow::anyhow!("Muse import stopped while parsing {expected_id}"))?; - } - Ok::<(usize, bool), anyhow::Error>((envelope_count, ended)) - }); - let mut processing_error = None; - while let Some(batch) = receiver.recv().await { - if let Err(error) = processor.process(batch).await { - processing_error = Some(error); - break; - } - } - drop(receiver); - let parse_result = parser.await.context("Muse export parser task failed")?; - if let Some(error) = processing_error { - return Err(error); - } - parse_result -} - -async fn attach_muse_session( - session_id: &str, - opts: ServeOptions, - config: Option, - ledger_dir: Option, -) -> anyhow::Result> { - transcript_import::validate_session_id(session_id)?; - let directory = tempfile::Builder::new() - .prefix("bt-muse-attach-") - .tempdir()?; - let export = directory.path().join("session.json"); - attach_muse_export(&export, session_id, true, opts, config, ledger_dir).await -} - -async fn attach_muse_export( - export: &std::path::Path, - session_id: &str, - reexport: bool, - opts: ServeOptions, - config: Option, - ledger_dir: Option, -) -> anyhow::Result> { - let mut processor = ImportProcessor::new(opts, config, ledger_dir); - let mut emitted = 0usize; - let shutdown = tokio::signal::ctrl_c(); - tokio::pin!(shutdown); - loop { - if reexport { - export_muse_session(session_id, export).await?; - } - let (count, ended) = - process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; - if count != emitted { - processor.flush_open().await?; - } - emitted = count; - if ended { - return processor.finish().await; - } - tokio::select! { - result = &mut shutdown => { - result?; - break; - } - _ = tokio::time::sleep(Duration::from_secs(1)) => {} - } - } - // Capture one final complete snapshot before interrupting any still-open - // spans, matching the finalization behavior of other attach sources. - if reexport { - export_muse_session(session_id, export).await?; - } - process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; - processor.finish().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"); - } - import_parent_components(args)?; - Ok(()) -} - -fn import_parent_components(args: &ImportArgs) -> anyhow::Result> { - if let Some(parent) = &args.parent { - parent - .to_parent_span_info() - .map_err(|error| anyhow::anyhow!("invalid --parent value: {error}"))?; - return Ok(Some(parent.clone())); - } - - let Some(span_id) = args.parent_span_id.as_deref() else { - return Ok(None); - }; - let span_id = span_id.trim(); - if span_id.is_empty() { - anyhow::bail!("--parent-span-id must not be empty"); - } - let root_span_id = args - .parent_root_span_id - .as_deref() - .ok_or_else(|| anyhow::anyhow!("--parent-span-id requires --parent-root-span-id"))?; - let root_span_id = root_span_id.trim(); - if root_span_id.is_empty() { - anyhow::bail!("--parent-root-span-id must not be empty"); - } - let object_type = args - .parent_object_type - .ok_or_else(|| anyhow::anyhow!("--parent-span-id requires --parent-object-type"))?; - - let object_id = args - .parent_object_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string); - let parent_project = args - .parent_project - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - let compute_object_metadata_args = match object_type { - ParentObjectType::ProjectLogs if object_id.is_none() => { - let project = parent_project.ok_or_else(|| { - anyhow::anyhow!( - "project_logs parent requires --parent-object-id or --parent-project" - ) - })?; - Some(serde_json::Map::from_iter([( - "project_name".to_string(), - serde_json::Value::String(project.to_string()), - )])) - } - ParentObjectType::ProjectLogs => None, - _ if parent_project.is_some() => { - anyhow::bail!("--parent-project is only valid with --parent-object-type project_logs") - } - _ if object_id.is_none() => { - anyhow::bail!("this parent object type requires --parent-object-id") - } - _ => None, - }; - - let components = SpanComponents { - object_type: object_type.into(), - object_id, - compute_object_metadata_args, - row_id: None, - span_id: Some(span_id.to_string()), - root_span_id: Some(root_span_id.to_string()), - span_parents: None, - propagated_event: None, - }; - components - .to_parent_span_info() - .map_err(|error| anyhow::anyhow!("invalid parent identifiers: {error}"))?; - Ok(Some(components)) -} - -/// Launch a coding agent with inherited stdio and inject Braintrust hooks for -/// this invocation, without requiring the tracing plugin to be installed or -/// enabled globally. -pub async fn run_traced( - args: RunArgs, - hook_command: RunHookCommand, - mut route: SessionRoute, - host: HostInfo, -) -> anyhow::Result { - apply_run_span_plugins(&mut route, &args.plugin)?; - if route.destination.is_none() { - anyhow::bail!( - "managed run requires a trace destination; select a project, object destination, or parent span" - ); - } - if args.source == RunSource::Codex - && args.agent_args.iter().any(|arg| { - let arg = arg.to_string_lossy(); - arg == "--dangerously-bypass-hook-trust" - || arg.starts_with("--dangerously-bypass-hook-trust=") - }) - { - anyhow::bail!( - "bt trace run codex cannot be combined with --dangerously-bypass-hook-trust; managed tracing preserves Codex hook trust. Remove the flag, or run Codex directly" - ); - } - let (executable_env, default_executable) = match args.source { - RunSource::Codex => ("CODEX_BIN", "codex"), - RunSource::Muse => ("MUSE_BIN", "muse"), - RunSource::Claude => ("CLAUDE_BIN", "claude"), - RunSource::OpenCode => ("OPENCODE_BIN", "opencode"), - RunSource::Pi => ("PI_BIN", "pi"), - }; - let executable = - std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); - if args.source == RunSource::Muse { - setup::check_muse_project_hooks(&std::env::current_dir()?)?; - for (index, arg) in args.agent_args.iter().enumerate() { - let Some(arg) = arg.to_str() else { continue }; - let workspace = if arg == "--workspace" { - args.agent_args.get(index + 1).map(PathBuf::from) - } else { - arg.strip_prefix("--workspace=").map(PathBuf::from) - }; - if let Some(workspace) = workspace { - setup::check_muse_project_hooks(&workspace)?; - } - } - let output = tokio::process::Command::new(&executable) - .arg("--version") - .output() - .await - .context("check Muse Code version for managed hooks")?; - if !output.status.success() - || setup::muse_version(&output.stdout)? != semver::Version::new(1, 1, 1) - { - anyhow::bail!("Muse managed runs require the verified Muse Code 1.1.1 hook override"); - } - } - let injected_args = managed_run_args(args.source, &hook_command)?; - let managed_run_id = uuid::Uuid::new_v4().to_string(); - // A caller-supplied socket is already an explicit daemon boundary (for - // example an integration harness or a deliberately isolated runtime). - // Only create our own boundary when environment auth would otherwise use - // the ambient shared daemon. - let isolate_daemon = (args.source == RunSource::Muse - || route.auth.effective_source() == wire::AuthSource::Environment) - && std::env::var_os(paths::SOCKET_ENV).is_none(); - let isolated_runtime = isolate_daemon - .then(|| ManagedRunRuntime::new(&managed_run_id)) - .transpose()?; - let muse_context_dir = (args.source == RunSource::Muse && isolated_runtime.is_none()) - .then(|| { - tempfile::Builder::new() - .prefix("bt-trace-run-muse-") - .tempdir() - }) - .transpose()?; - let muse_route = (args.source == RunSource::Muse).then(|| route.clone()); - let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?; - let mut command = tokio::process::Command::new(&executable); - command - .args(injected_args) - .args(args.agent_args) - .env("_BT_TRACE_MANAGED_RUN", "1") - .env(MANAGED_RUN_ID_ENV, &managed_run_id) - .env(settings::INVOCATION_SETTINGS_ENV, invocation_settings); - if args.source == RunSource::Muse { - let context_dir = isolated_runtime - .as_ref() - .map(|runtime| runtime.temp_dir.path()) - .or_else(|| muse_context_dir.as_ref().map(|dir| dir.path())) - .expect("Muse has an invocation directory"); - let context = MuseManagedContext { - route: muse_route.expect("Muse route is available"), - managed_run_id: managed_run_id.clone(), - socket: isolated_runtime - .as_ref() - .map(|runtime| runtime.socket.clone()) - .unwrap_or_else(|| paths::socket_path(None)), - }; - let context_path = context_dir.join("muse-context.json"); - std::fs::write(&context_path, serde_json::to_vec(&context)?)?; - let mut hook = hook_command.clone(); - hook.args.extend([ - OsString::from("--source"), - OsString::from("muse"), - OsString::from("--managed-run-hook"), - OsString::from("--flush-on-turn-end"), - OsString::from("--managed-context"), - context_path.into_os_string(), - ]); - let command_string = shell_command(&hook, cfg!(windows))?; - let hooks = setup::muse_managed_hook_config(&paths::muse_config_dir(), &command_string)?; - let hooks_path = context_dir.join("muse-hooks.json"); - std::fs::write(&hooks_path, serde_json::to_vec(&hooks)?)?; - command.env("TBH_MANAGED_HOOKS_PATH", hooks_path); - if let Some(runtime) = &isolated_runtime { - start_muse_managed_daemon(&host, runtime).await?; - } - } - if let Some(runtime) = &isolated_runtime { - command - .env(paths::SOCKET_ENV, &runtime.socket) - .env(paths::DATA_DIR_ENV, runtime.temp_dir.path()); - } - if args.source == RunSource::OpenCode { - command.env( - "OPENCODE_CONFIG_CONTENT", - opencode_managed_config(std::env::var("OPENCODE_CONFIG_CONTENT").ok().as_deref())?, - ); - } - let mut child = match command.spawn() { - Ok(child) => child, - Err(error) => { - if args.source == RunSource::Muse { - if let Some(runtime) = &isolated_runtime { - let _ = shutdown_daemon(&runtime.socket).await; - } - } - return Err(anyhow::anyhow!( - "failed to launch {}: {error}", - executable.to_string_lossy() - )); - } - }; - let interrupt = tokio::signal::ctrl_c(); - tokio::pin!(interrupt); - - let status = tokio::select! { - status = child.wait() => status.map_err(anyhow::Error::from), - result = &mut interrupt => { - match result { - Ok(()) => { - let kill_result = child.start_kill(); - let wait_result = child.wait().await; - kill_result - .map_err(anyhow::Error::from) - .and_then(|()| wait_result.map_err(anyhow::Error::from)) - } - Err(error) => Err(error.into()), - } - } - }; - let socket = isolated_runtime - .as_ref() - .map(|runtime| runtime.socket.clone()) - .unwrap_or_else(|| paths::socket_path(None)); - let data_dir = isolated_runtime - .as_ref() - .map(|runtime| runtime.temp_dir.path().to_path_buf()) - .unwrap_or_else(|| paths::data_dir(None)); - match flush_managed_run_in( - &managed_run_id, - &socket, - MANAGED_RUN_FLUSH_TIMEOUT_MS, - &data_dir, - ) - .await - { - Ok(result) if result.accepted_sessions == 0 => tracing::warn!( - managed_run_id, - "managed run produced no accepted trace events" - ), - Ok(result) if result.flushed => {} - Ok(result) => tracing::warn!( - managed_run_id, - pending = result.pending, - "managed run trace flush timed out" - ), - Err(error) => tracing::warn!(managed_run_id, %error, "managed run trace flush failed"), - } - if let Some(runtime) = &isolated_runtime { - let _ = shutdown_daemon(&socket).await; - if let Err(error) = - crate::plugin_diagnostics::merge(runtime.temp_dir.path(), &paths::data_dir(None)) - { - tracing::warn!(managed_run_id, %error, "failed to preserve managed-run plugin diagnostics"); - } - } - status -} - -struct ManagedRunRuntime { - temp_dir: tempfile::TempDir, - socket: std::path::PathBuf, -} - -impl ManagedRunRuntime { - fn new(_managed_run_id: &str) -> anyhow::Result { - let temp_dir = tempfile::Builder::new().prefix("bt-trace-run-").tempdir()?; - #[cfg(unix)] - let socket = temp_dir.path().join("daemon.sock"); - #[cfg(windows)] - let socket = std::path::PathBuf::from(format!( - r"\\.\pipe\braintrust-bt-daemon-managed-{_managed_run_id}" - )); - Ok(Self { temp_dir, socket }) - } -} - -async fn start_muse_managed_daemon( - host: &HostInfo, - runtime: &ManagedRunRuntime, -) -> anyhow::Result<()> { - let stream = - client::ensure_daemon_in(&runtime.socket, host, false, Some(runtime.temp_dir.path())) - .await - .context("start isolated daemon for Muse managed run")?; - drop(stream); - Ok(()) -} - -fn apply_run_span_plugins(route: &mut SessionRoute, plugins: &[PathBuf]) -> anyhow::Result<()> { - route.span_plugins = resolve_span_plugin_paths(plugins)?; - Ok(()) -} - -pub(crate) fn resolve_span_plugin_paths(paths: &[PathBuf]) -> anyhow::Result> { - let paths: Vec<_> = paths - .iter() - .map(|path| { - path.canonicalize().map_err(|error| { - anyhow::anyhow!("could not resolve span plugin {}: {error}", path.display()) - }) - }) - .collect::>()?; - crate::span_processor::validate(&paths)?; - Ok(paths) -} - -fn managed_run_args( - source: RunSource, - hook_command: &RunHookCommand, -) -> anyhow::Result> { - let source_name = match source { - RunSource::Codex => "codex", - RunSource::Muse => "muse", - RunSource::Claude => "claude", - RunSource::OpenCode => "opencode", - RunSource::Pi => "pi", - }; - match source { - RunSource::Codex | RunSource::Claude => { - let unix_command = managed_hook_shell_command(hook_command, source_name, false)?; - let windows_command = managed_hook_shell_command(hook_command, source_name, true)?; - match source { - RunSource::Codex => Ok(codex_managed_run_args(&unix_command, &windows_command)), - RunSource::Claude => Ok(claude_managed_run_args(if cfg!(windows) { - &windows_command - } else { - &unix_command - })?), - _ => unreachable!(), - } - } - RunSource::Muse | RunSource::OpenCode => Ok(Vec::new()), - RunSource::Pi => { - let extension = match std::env::var_os("BT_TRACE_PI_PLUGIN_SPEC") { - Some(extension) => extension, - None => OsString::from(crate::setup::pi_plugin_spec()), - }; - Ok(vec![OsString::from("-e"), extension]) - } - } -} - -fn opencode_managed_config(existing: Option<&str>) -> anyhow::Result { - let mut config = match existing { - Some(raw) => serde_json::from_str::(raw) - .map_err(|error| anyhow::anyhow!("invalid OPENCODE_CONFIG_CONTENT: {error}"))?, - None => serde_json::json!({}), - }; - let object = config - .as_object_mut() - .ok_or_else(|| anyhow::anyhow!("OPENCODE_CONFIG_CONTENT must be a JSON object"))?; - let plugins = object - .entry("plugin") - .or_insert_with(|| serde_json::Value::Array(Vec::new())) - .as_array_mut() - .ok_or_else(|| anyhow::anyhow!("OPENCODE_CONFIG_CONTENT.plugin must be an array"))?; - let plugin = std::env::var("BT_TRACE_OPENCODE_PLUGIN_SPEC") - .unwrap_or_else(|_| "@braintrust/trace-opencode/tracing".to_string()); - if !plugins.iter().any(|value| value.as_str() == Some(&plugin)) { - plugins.push(serde_json::Value::String(plugin)); - } - Ok(serde_json::to_string(&config)?) -} - -fn managed_hook_shell_command( - hook_command: &RunHookCommand, - source: &str, - windows: bool, -) -> anyhow::Result { - let mut argv = Vec::with_capacity(hook_command.args.len() + 4); - argv.push(hook_command.program.clone()); - argv.extend(hook_command.args.iter().cloned()); - argv.push(OsString::from("--source")); - argv.push(OsString::from(source)); - argv.push(OsString::from("--managed-run-hook")); - shell_argv(&argv, windows) -} - -fn shell_command(command: &RunHookCommand, windows: bool) -> anyhow::Result { - let mut argv = vec![command.program.clone()]; - argv.extend(command.args.iter().cloned()); - shell_argv(&argv, windows) -} - -fn shell_argv(argv: &[OsString], windows: bool) -> anyhow::Result { - let mut rendered = Vec::with_capacity(argv.len()); - for arg in argv { - let arg = arg - .to_str() - .ok_or_else(|| anyhow::anyhow!("managed hook command contains non-Unicode argv"))?; - rendered.push(if windows { - quote_windows_command_arg(arg) - } else { - quote_unix_shell_arg(arg) - }); - } - Ok(rendered.join(" ")) -} - -fn quote_unix_shell_arg(arg: &str) -> String { - format!("'{}'", arg.replace('\'', "'\"'\"'")) -} - -fn quote_windows_command_arg(arg: &str) -> String { - format!("\"{}\"", arg.replace('\\', "/").replace('"', "\"\"")) -} - -const CODEX_RUN_HOOK_EVENTS: &[&str] = &[ - "SessionStart", - "UserPromptSubmit", - "PreToolUse", - "PermissionRequest", - "PostToolUse", - "PreCompact", - "PostCompact", - "SubagentStart", - "SubagentStop", - "Stop", - "SessionEnd", -]; - -const CLAUDE_RUN_HOOK_EVENTS: &[&str] = &[ - "SessionStart", - "Setup", - "UserPromptSubmit", - "UserPromptExpansion", - "PreToolUse", - "PermissionRequest", - "PermissionDenied", - "PostToolUse", - "PostToolUseFailure", - "PostToolBatch", - "PreCompact", - "PostCompact", - "Notification", - "MessageDisplay", - "SubagentStart", - "SubagentStop", - "TaskCreated", - "TaskCompleted", - "Stop", - "StopFailure", - "SessionEnd", -]; - -fn codex_managed_run_args(unix_command: &str, windows_command: &str) -> Vec { - let unix_command = serde_json::to_string(unix_command).expect("serialize hook command"); - let windows_command = - serde_json::to_string(windows_command).expect("serialize Windows hook command"); - let mut args = vec![OsString::from("--enable"), OsString::from("hooks")]; - for event in CODEX_RUN_HOOK_EVENTS { - args.push(OsString::from("-c")); - args.push(OsString::from(format!( - "hooks.{event}=[{{hooks=[{{type=\"command\",command={unix_command},commandWindows={windows_command}}}]}}]" - ))); - } - args -} - -fn claude_managed_run_args(command: &str) -> anyhow::Result> { - let matcher_group = serde_json::json!([{ - "hooks": [{ - "type": "command", - "command": command, - "async": false - }] - }]); - let hooks = CLAUDE_RUN_HOOK_EVENTS - .iter() - .map(|event| ((*event).to_string(), matcher_group.clone())) - .collect::>(); - Ok(vec![ - OsString::from("--settings"), - OsString::from(serde_json::to_string( - &serde_json::json!({ "hooks": hooks }), - )?), - ]) -} - -fn apply_import_destination( - config: &mut Option, - destination: Option, -) -> anyhow::Result<()> { - if let Some(destination) = destination { - let config = config.as_mut().ok_or_else(|| { - anyhow::anyhow!( - "import destination requires a resolved Braintrust session configuration; \ - use `bt trace import` instead of the standalone `bt-daemon import` command" - ) - })?; - config.destination = Some(destination); - } - Ok(()) -} - -/// Import a native transcript from a known path. Front-ends should normally -/// expose [`run_import`] so users only need the agent's session id; this lower- -/// level entry point is useful for embedding and isolated tests. -pub async fn import_transcript( - file: &std::path::Path, - source: ImportSource, - opts: ServeOptions, - config: Option, - attach: bool, -) -> anyhow::Result> { - import_transcript_with_ledger(file, source, opts, config, attach, None).await -} - -async fn import_transcript_with_ledger( - file: &std::path::Path, - source: ImportSource, - opts: ServeOptions, - config: Option, - attach: bool, - ledger_dir: Option, -) -> anyhow::Result> { - if source == ImportSource::Muse { - let session_id = transcript_import::muse::session_id(file)?; - if attach { - return attach_muse_export(file, &session_id, false, opts, config, ledger_dir).await; - } - let mut processor = ImportProcessor::new(opts, config, ledger_dir); - let ended = process_muse_export(&mut processor, file, &session_id).await?; - let mut summaries = Vec::new(); - let summary = if ended { - processor.finish_session(&session_id).await? - } else { - processor.detach_session(&session_id).await? - }; - if let Some(summary) = summary { - summaries.push(summary); - } - summaries.extend(processor.finish().await?); - return Ok(summaries); - } - let mut tail = transcript_import::TranscriptTail::new(file.to_path_buf(), source); - let mut processor = ImportProcessor::new(opts, config, ledger_dir); - let shutdown = tokio::signal::ctrl_c(); - tokio::pin!(shutdown); - let mut finalizing = !attach; - loop { - let entries = tail.poll(finalizing)?; - if tail.take_translator_reset() { - processor.reset_translators(); - } - processor.process(entries).await?; - if finalizing { - break; - } - tokio::select! { - result = &mut shutdown => { - result?; - finalizing = true; - } - _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} - } - } - 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> { - import_transcripts_with_ledger(files, source, opts, config, None).await -} - -async fn import_transcripts_with_ledger( - files: &[PathBuf], - source: ImportSource, - opts: ServeOptions, - config: Option, - ledger_dir: Option, -) -> anyhow::Result> { - let mut processor = ImportProcessor::new(opts, config, ledger_dir); - let mut summaries = Vec::new(); - 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()))?; - 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) -} - -struct ImportLive { - source: String, - translator: Box, - sink: Box, - ctx: SessionCtx, - pending_ops: usize, - span_ids: std::collections::HashSet, - root_span_id: Option, - destination: Option, - diagnostics_dir: Option, -} - -struct ImportProcessor { - sessions: std::collections::HashMap, - opts: ServeOptions, - config: Option, - ledger_dir: Option, -} - -impl ImportProcessor { - fn new(opts: ServeOptions, config: Option, ledger_dir: Option) -> Self { - Self { - sessions: std::collections::HashMap::new(), - opts, - config, - ledger_dir, - } - } - - async fn process(&mut self, entries: Vec) -> anyhow::Result<()> { - for mut env in entries { - env.config = self.config.clone(); - let sid = env.session_id.clone(); - let live = match self.sessions.get_mut(&sid) { - Some(live) => live, - None => { - let translator = self.opts.translators.create(&env.source, &sid); - let sink = self.opts.sink_factory.create( - &sid, - &env.source, - env.plugin_version.as_deref(), - )?; - let sink: Box = match &self.ledger_dir { - Some(ledger_dir) => Box::new( - delivery_ledger::LedgerSink::new( - sink, - ledger_dir, - &env.source, - &sid, - env.config.as_ref(), - ) - .await, - ), - None => sink, - }; - self.sessions.insert( - sid.clone(), - ImportLive { - source: env.source.clone(), - translator, - sink, - ctx: SessionCtx { - session_id: sid.clone(), - config: None, - }, - pending_ops: 0, - span_ids: std::collections::HashSet::new(), - root_span_id: self - .config - .as_ref() - .and_then(|config| config.attached_span_ids().1), - destination: self - .config - .as_ref() - .and_then(|config| config.destination.clone()), - diagnostics_dir: self.ledger_dir.clone(), - }, - ); - self.sessions.get_mut(&sid).unwrap() - } - }; - if let Some(cfg) = &env.config { - live.sink.configure(cfg); - live.ctx.config = Some(cfg.clone()); - } - let ops = live.translator.handle(&env, &live.ctx)?; - Self::emit_translator_batches(live, ops).await?; - } - Ok(()) - } - - fn reset_translators(&mut self) { - for (session_id, live) in &mut self.sessions { - live.translator = self.opts.translators.create(&live.source, session_id); - } - } - - async fn emit_translator_batches( - live: &mut ImportLive, - first: Vec, - ) -> anyhow::Result<()> { - let mut next = Some(first); - while let Some(ops) = next { - // Imports can contain tens of thousands of SDK log commands. Bound the - // number queued between drains without serializing one network flush - // for every native turn boundary. - const FLUSH_OPS: usize = 500; - for chunk in ops.chunks(FLUSH_OPS) { - for op in chunk { - let row = match op { - SpanOp::Insert(row) | SpanOp::Merge(row) => row, - }; - live.span_ids.insert(row.span_id.clone()); - if live.root_span_id.is_none() && !row.root_span_id.is_empty() { - live.root_span_id = Some(row.root_span_id.clone()); - } - } - let plugins = live - .ctx - .config - .as_ref() - .map(|config| config.span_plugins.as_slice()) - .unwrap_or_default(); - let mut transformed = Vec::with_capacity(chunk.len()); - for op in chunk { - match crate::span_processor::process( - plugins, - op, - &live.source, - &live.ctx.session_id, - ) { - Ok(result) => { - if let Some(failure) = result.failure { - if failure.newly_seen { - if let Some(data_dir) = &live.diagnostics_dir { - if let Err(error) = crate::plugin_diagnostics::record( - data_dir, - &live.source, - &failure.path, - &failure.message, - ) { - tracing::warn!( - session_id = %live.ctx.session_id, - "failed to persist span plugin diagnostic: {error}" - ); - } - } - tracing::warn!( - session_id = %live.ctx.session_id, - plugin = %failure.path.display(), - error = %failure.message, - "span plugin failed during import; span operations are being discarded" - ); - } - } - if let Some(op) = result.op { - transformed.push(op); - } - } - Err(error) => { - if let (Some(data_dir), Some(plugin)) = - (&live.diagnostics_dir, plugins.first()) - { - if let Err(diagnostic_error) = crate::plugin_diagnostics::record( - data_dir, - &live.source, - plugin, - &error.to_string(), - ) { - tracing::warn!( - session_id = %live.ctx.session_id, - "failed to persist span plugin diagnostic: {diagnostic_error}" - ); - } - } - tracing::warn!( - session_id = %live.ctx.session_id, - %error, - "span plugin processor failed during import; span operation discarded" - ); - } - } - } - if !transformed.is_empty() { - live.sink.emit(&transformed).await?; - } - live.pending_ops += transformed.len(); - if live.pending_ops >= FLUSH_OPS { - live.sink.flush().await?; - live.pending_ops = 0; - } - } - next = live.translator.drain_pending(&live.ctx)?; - } - Ok(()) - } - - async fn finish(self) -> anyhow::Result> { - let mut summaries = Vec::new(); - for (sid, mut live) in self.sessions { - let ops = live.translator.flush(&live.ctx)?; - Self::emit_translator_batches(&mut live, ops).await?; - live.sink.flush().await?; - summaries.push(Self::summary(sid, live, true)); - } - summaries.sort_by(|left, right| left.session_id.cmp(&right.session_id)); - Ok(summaries) - } - - async fn flush_open(&mut self) -> anyhow::Result<()> { - for live in self.sessions.values_mut() { - live.sink.flush().await?; - live.pending_ops = 0; - } - Ok(()) - } - - async fn finish_session(&mut self, session_id: &str) -> anyhow::Result> { - let Some(mut live) = self.sessions.remove(session_id) else { - return Ok(None); - }; - let ops = live.translator.flush(&live.ctx)?; - Self::emit_translator_batches(&mut live, ops).await?; - live.sink.flush().await?; - Ok(Some(Self::summary(session_id.to_string(), live, true))) - } - - /// A one-shot snapshot of an active session must not mark its spans - /// complete in the destination ledger. A later import can then extend the - /// same deterministic rows when Muse records the real terminal events. - async fn detach_session(&mut self, session_id: &str) -> anyhow::Result> { - let Some(mut live) = self.sessions.remove(session_id) else { - return Ok(None); - }; - live.sink.flush().await?; - Ok(Some(Self::summary(session_id.to_string(), live, false))) - } - - fn summary(session_id: String, live: ImportLive, finalized: bool) -> ImportSummary { - ImportSummary { - session_id, - destination: live.destination, - root_span_id: live.root_span_id, - span_count: live.span_ids.len(), - finalized, - } - } -} - -/// Build debug [`ServeOptions`]: the production translator registry plus a -/// debug sink writing NDJSON under `/spans/`. -pub fn debug_serve_options(version: impl Into, data_dir: &std::path::Path) -> ServeOptions { - ServeOptions { - version: version.into(), - translators: Arc::new(Registry::default_agents()), - sink_factory: Arc::new(DebugSinkFactory { - dir: data_dir.join("spans"), - }), - auth_provider: None, - } -} - -/// Build [`ServeOptions`] with the Braintrust sink. `translators` lets the -/// caller choose the translator registry (debug-only until Phase 3 adds the -/// Codex/Claude translators). Clients are built lazily per session URL, so this -/// is cheap and infallible. -pub fn braintrust_serve_options( - version: impl Into, - sink_config: BraintrustSinkConfig, - translators: Arc, -) -> ServeOptions { - ServeOptions { - version: version.into(), - translators, - sink_factory: Arc::new(BraintrustSinkFactory::new(sink_config)), - auth_provider: None, - } -} - -fn now_ms() -> i64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0) -} - -fn read_stdin_json() -> anyhow::Result { - use std::io::Read; - let mut buf = String::new(); - std::io::stdin().read_to_string(&mut buf)?; - if buf.trim().is_empty() { - anyhow::bail!("empty stdin (expected a JSON hook payload)"); - } - Ok(serde_json::from_str(&buf)?) -} - -/// Read a string-ish field (`session_id` / event name) from the payload, -/// coercing a JSON number to its string form. -fn json_str_field(payload: &serde_json::Value, field: &str) -> Option { - match payload.get(field) { - Some(serde_json::Value::String(s)) => Some(s.clone()), - Some(serde_json::Value::Number(n)) => Some(n.to_string()), - _ => None, - } -} - -/// Stamp the transcript boundary visible when a blocking hook runs. Agent -/// transcripts are append-only, while daemon journal replay may happen after -/// the session has advanced. Recording byte lengths keeps translation causally -/// aligned with each native hook without copying transcript contents into the -/// journal. -fn add_transcript_observation(payload: &mut serde_json::Value, field: &str) { - let Some(path) = json_str_field(payload, field) else { - return; - }; - let transcript = std::path::Path::new(&path); - let mut observation = serde_json::Map::new(); - observation.insert("path".into(), serde_json::Value::String(path.clone())); - if let Ok(metadata) = std::fs::metadata(transcript) { - observation.insert( - "observed_bytes".into(), - serde_json::Value::Number(metadata.len().into()), - ); - } - - if transcript.file_name().and_then(|name| name.to_str()) == Some("transcript.jsonl") { - let full = transcript.with_file_name("transcript_full.jsonl"); - if let Ok(metadata) = std::fs::metadata(&full) { - observation.insert( - "full_path".into(), - serde_json::Value::String(full.to_string_lossy().into_owned()), - ); - observation.insert( - "full_observed_bytes".into(), - serde_json::Value::Number(metadata.len().into()), - ); - } - } - - if let Some(object) = payload.as_object_mut() { - object.insert( - "_bt_transcript_observation".into(), - serde_json::Value::Object(observation), - ); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use clap::Parser; - use serde_json::json; - use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; - - #[derive(Debug, Parser)] - struct ImportCli { - #[command(flatten)] - args: ImportArgs, - } - - #[derive(Debug, Parser)] - struct ServeCli { - #[command(flatten)] - args: ServeArgs, - } - - #[derive(Debug, Parser)] - struct HookCli { - #[command(flatten)] - args: HookArgs, - } - - #[tokio::test] - async fn muse_attach_snapshots_reuse_one_trace_and_keep_open_spans() { - let temp = tempfile::tempdir().unwrap(); - let export = temp.path().join("export.json"); - let output = temp.path().join("spans"); - let mut processor = ImportProcessor::new( - ServeOptions { - version: "test".into(), - translators: Arc::new(Registry::default_agents()), - sink_factory: Arc::new(DebugSinkFactory { - dir: output.clone(), - }), - auth_provider: None, - }, - None, - None, - ); - let mut events = vec![ - json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}), - json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","request_record_id":"request-1"}}}}), - ]; - let write = |events: &[serde_json::Value]| { - std::fs::write( - &export, - serde_json::to_vec(&json!({ - "export_schema_version":1, - "sessions":[{"session_id":"muse-attach"}], - "events":events, - })) - .unwrap(), - ) - .unwrap(); - }; - write(&events); - let (count, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, 0) - .await - .unwrap(); - assert_eq!(count, 3); - assert!(!ended); - processor.flush_open().await.unwrap(); - let (same_count, ended) = - process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) - .await - .unwrap(); - assert_eq!(same_count, count); - assert!(!ended); - events.extend([ - json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_completed","usage":{"input_tokens":2}}}}}), - json!({"recorded_at":1_003_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"done"}}}}), - json!({"recorded_at":1_004_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), - json!({"recorded_at":1_005_000,"envelope":{"payload":{"kind":"session_end"}}}), - ]); - write(&events); - let (_, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) - .await - .unwrap(); - assert!(ended); - processor.finish().await.unwrap(); - let rows = std::fs::read_to_string(output.join("muse-attach.ndjson")) - .unwrap() - .lines() - .map(|line| serde_json::from_str::(line).unwrap()) - .collect::>(); - for identity in ["root", "turn:run-1", "llm:request-1"] { - let id = ids::span_id(&ids::session_namespace("muse", "muse-attach"), identity); - assert_eq!( - rows.iter() - .filter(|row| row - .pointer("/Insert/span_id") - .and_then(|value| value.as_str()) - == Some(id.as_str())) - .count(), - 1, - "duplicate insert for {identity}" - ); - assert_eq!( - rows.iter() - .filter(|row| row - .pointer("/Merge/span_id") - .and_then(|value| value.as_str()) - == Some(id.as_str())) - .count(), - 1, - "missing or duplicate completion for {identity}" - ); - } - } - - #[test] - fn serve_defaults_to_short_journal_backed_session_retirement() { - let args = ServeCli::try_parse_from(["test"]).unwrap().args; - assert_eq!(args.session_idle_timeout_secs, 30); - } - - #[test] - fn hook_plugin_version_reaches_envelope_and_every_initialize_attempt() { - let args = HookCli::try_parse_from([ - "test", - "--source", - "grok", - "--source-version", - "1.0.13", - "--plugin-version", - "0.1.0", - ]) - .unwrap() - .args; - assert_eq!(args.source_version.as_deref(), Some("1.0.13")); - assert_eq!(args.plugin_version.as_deref(), Some("0.1.0")); - - let payload = serde_json::json!({"native": "unchanged"}); - let env = build_hook_envelope( - &args, - SessionRoute::default(), - payload.clone(), - "session-1".into(), - "SessionEnd".into(), - ); - assert_eq!(env.source_version.as_deref(), Some("1.0.13")); - assert_eq!(env.plugin_version.as_deref(), Some("0.1.0")); - assert_eq!(env.payload, payload); - - // Both the initial connection and the post-restart retry use this - // shared parameter builder. - let initialize = initialize_params(&env, "1.0.13"); - assert_eq!(initialize["client"]["source"], "grok"); - assert_eq!(initialize["client"]["daemon_version"], "1.0.13"); - assert_eq!(initialize["client"]["plugin_version"], "0.1.0"); - assert_ne!(initialize["client"]["plugin_version"], "1.0.13"); - } - - #[test] - fn hook_versions_are_discovered_without_hook_configuration_literals() { - let payload = serde_json::json!({"cli_version": "2.3.4"}); - assert_eq!( - source_version_from_payload("codex", &payload).as_deref(), - Some("2.3.4") - ); - - let temp = tempfile::tempdir().unwrap(); - let manifest = temp.path().join("plugin.json"); - std::fs::write(&manifest, br#"{"version":"5.6.7"}"#).unwrap(); - assert_eq!( - plugin_version_from_manifest(&manifest).as_deref(), - Some("5.6.7") - ); - } - - #[test] - fn hook_flush_recognizes_native_and_documented_terminal_events() { - for event in ["session_end", "SessionEnd"] { - assert!(should_flush_hook_event(event, false)); - } - for event in [ - "stop", - "Stop", - "stop_failure", - "StopFailure", - "stop_cancelled", - "StopCancelled", - "subagent_stop", - "SubagentStop", - ] { - assert!(!should_flush_hook_event(event, false)); - assert!(should_flush_hook_event(event, true)); - } - assert!(!should_flush_hook_event("turn_completed", true)); - } - - #[test] - fn pi_agent_end_flushes_only_completed_turns_when_enabled() { - let mut env: wire::Envelope = serde_json::from_value(serde_json::json!({ - "source": "pi", "session_id": "pi-turn", "event": "agent_end", - "ts_ms": 1, "payload": {"event": {}}, - "route": {"flush_mode": "flush_on_turn_end"} - })) - .unwrap(); - for native in [ - serde_json::json!({}), - serde_json::json!({"willRetry": false}), - ] { - env.payload = serde_json::json!({"event": native}); - assert!(should_flush_ingress_event(&env)); - } - for payload in [ - serde_json::json!({"event": {"willRetry": true}}), - serde_json::json!({"willRetry": true}), - ] { - env.payload = payload; - assert!(!should_flush_ingress_event(&env)); - } - env.payload = serde_json::json!({"event": {}}); - env.source = "opencode".into(); - assert!(!should_flush_ingress_event(&env)); - env.source = "pi".into(); - env.event = "message_end".into(); - assert!(!should_flush_ingress_event(&env)); - env.event = "agent_end".into(); - env.route.as_mut().unwrap().flush_mode = wire::FlushMode::FireAndForget; - assert!(!should_flush_ingress_event(&env)); - env.route = None; - assert!(!should_flush_ingress_event(&env)); - for event in ["session_shutdown", "session_compact", "session_tree"] { - env.event = event.into(); - assert!(should_flush_ingress_event(&env)); - env.source = "opencode".into(); - assert!(!should_flush_ingress_event(&env)); - env.source = "pi".into(); - } - } - - #[test] - fn opencode_lifecycle_events_flush_from_ingress() { - let mut env: wire::Envelope = serde_json::from_value(serde_json::json!({ - "source": "opencode", "session_id": "opencode-session", "event": "session.idle", - "ts_ms": 1, "payload": {}, - "route": {"flush_mode": "fire_and_forget"} - })) - .unwrap(); - for event in [ - "session.idle", - "session.deleted", - "session.error", - "server.instance.disposed", - ] { - env.event = event.into(); - assert!(should_flush_ingress_event(&env)); - } - env.event = "session.compacted".into(); - assert!(!should_flush_ingress_event(&env)); - env.event = "session.idle".into(); - env.source = "pi".into(); - assert!(!should_flush_ingress_event(&env)); - } - - #[test] - fn span_plugin_paths_are_canonicalized_before_use() { - let dir = tempfile::Builder::new() - .prefix("span-plugin-path-") - .tempdir_in(".") - .unwrap(); - let plugin = dir.path().join("plugin.mjs"); - std::fs::write(&plugin, "export default span => span").unwrap(); - let relative = PathBuf::from(dir.path().file_name().unwrap()).join("plugin.mjs"); - - let resolved = resolve_span_plugin_paths(&[relative]).unwrap(); - - assert_eq!(resolved, [plugin.canonicalize().unwrap()]); - assert!(resolved[0].is_absolute()); - } - - #[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 muse_msp_pages_reject_partial_or_cyclic_session_lists() { - let mut sessions = Vec::new(); - let mut cursors = std::collections::HashSet::new(); - assert_eq!( - muse_msp_append_page( - json!({"sessions":[{"sessionId":"session-a"}],"nextCursor":"next"}), - 1, - &mut sessions, - &mut cursors, - ) - .unwrap(), - Some("next".to_owned()) - ); - assert_eq!(sessions, ["session-a"]); - assert!(muse_msp_append_page( - json!({"sessions":[],"nextCursor":"next"}), - 2, - &mut sessions, - &mut cursors, - ) - .unwrap_err() - .to_string() - .contains("repeated cursor")); - assert!(muse_msp_append_page( - json!({"sessions":[{"status":"notLoaded"}]}), - 2, - &mut sessions, - &mut cursors, - ) - .unwrap_err() - .to_string() - .contains("decode Muse session/list page 2")); - } - - #[tokio::test] - async fn muse_msp_request_accepts_notifications_but_rejects_wrong_response_ids() { - async fn exchange(response: &str) -> anyhow::Result { - let (client, server) = tokio::io::duplex(4096); - let (client_read, mut client_write) = tokio::io::split(client); - let (server_read, mut server_write) = tokio::io::split(server); - let response = response.to_owned(); - let server = tokio::spawn(async move { - let request = BufReader::new(server_read) - .lines() - .next_line() - .await - .unwrap() - .unwrap(); - assert_eq!( - serde_json::from_str::(&request).unwrap()["id"], - 7 - ); - server_write.write_all(response.as_bytes()).await.unwrap(); - }); - let result = muse_msp_request( - &mut client_write, - &mut BufReader::new(client_read).lines(), - 7, - "session/list", - json!({}), - ) - .await; - server.await.unwrap(); - result - } - - let result = exchange( - "{\"jsonrpc\":\"2.0\",\"method\":\"progress\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n", - ) - .await - .unwrap(); - assert_eq!(result, json!({"sessions": []})); - assert!(exchange("{\"jsonrpc\":\"2.0\",\"id\":8,\"result\":{}}\n") - .await - .unwrap_err() - .to_string() - .contains("unexpected response id")); - } - - #[test] - fn additional_metadata_overrides_a_route_only_with_a_json_object() { - let mut route = SessionRoute { - additional_metadata: Some(serde_json::json!({"saved": true})), - ..SessionRoute::default() - }; - apply_additional_metadata(&mut route, Some(r#"{"run_id":"123"}"#)).unwrap(); - assert_eq!( - route.additional_metadata, - Some(serde_json::json!({"run_id": "123"})) - ); - - let error = apply_additional_metadata(&mut route, Some("[]")).unwrap_err(); - assert!(error.to_string().contains("must be a JSON object")); - let error = apply_additional_metadata(&mut route, Some("not-json")).unwrap_err(); - assert!(error - .to_string() - .contains("invalid --additional-metadata JSON")); - } - - #[test] - fn tags_override_a_route_and_are_normalized() { - let mut route = SessionRoute { - tags: vec!["saved".into()], - ..SessionRoute::default() - }; - apply_tags(&mut route, &[" ci ".into(), "ci".into(), "docs".into()]).unwrap(); - assert_eq!(route.tags, ["ci", "docs"]); - - let error = apply_tags(&mut route, &[" ".into()]).unwrap_err(); - assert!(error.to_string().contains("must not be empty")); - } - - #[test] - fn managed_run_plugins_replace_inherited_plugins() { - let temp = tempfile::tempdir().unwrap(); - let plugin = temp.path().join("run.mjs"); - std::fs::write(&plugin, "export default span => span").unwrap(); - let mut route = SessionRoute { - span_plugins: vec![PathBuf::from("persisted.mjs")], - ..SessionRoute::default() - }; - - apply_run_span_plugins(&mut route, &[]).unwrap(); - assert!(route.span_plugins.is_empty()); - - apply_run_span_plugins(&mut route, std::slice::from_ref(&plugin)).unwrap(); - assert_eq!(route.span_plugins, [plugin.canonicalize().unwrap()]); - } - - #[test] - fn import_plugins_replace_inherited_plugins() { - let temp = tempfile::tempdir().unwrap(); - let plugin = temp.path().join("import.mjs"); - std::fs::write(&plugin, "export default span => span").unwrap(); - let mut config = Some(SessionConfig { - auth: wire::BackendAuth { - token: String::new(), - api_url: None, - app_url: None, - org_name: None, - org_id: None, - }, - destination: None, - flush_mode: wire::FlushMode::FireAndForget, - additional_metadata: None, - tags: Vec::new(), - span_plugins: vec![PathBuf::from("persisted.mjs")], - }); - - apply_import_span_plugins(&mut config, &[]).unwrap(); - assert!(config.as_ref().unwrap().span_plugins.is_empty()); - - apply_import_span_plugins(&mut config, std::slice::from_ref(&plugin)).unwrap(); - assert_eq!( - config.unwrap().span_plugins, - [plugin.canonicalize().unwrap()] - ); - } - - #[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); - - 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()); - } - - #[test] - fn import_parent_accepts_complete_cli_identifiers_without_guessing() { - let args = ImportCli::try_parse_from([ - "test", - "codex", - "session-a", - "--parent-span-id", - "span-1", - "--parent-root-span-id", - "root-1", - "--parent-object-type", - "project_logs", - "--parent-project", - "Agents", - ]) - .unwrap() - .args; - - let components = import_parent_components(&args).unwrap().unwrap(); - assert_eq!(components.object_type, SpanObjectType::ProjectLogs); - assert_eq!(components.span_id.as_deref(), Some("span-1")); - assert_eq!(components.root_span_id.as_deref(), Some("root-1")); - assert_eq!( - components - .compute_object_metadata_args - .as_ref() - .and_then(|value| value.get("project_name")) - .and_then(serde_json::Value::as_str), - Some("Agents") - ); - } - - #[test] - fn import_parent_rejects_incomplete_or_conflicting_identifiers() { - assert!(ImportCli::try_parse_from([ - "test", - "codex", - "session-a", - "--parent-span-id", - "span-1", - ]) - .is_err()); - assert!(ImportCli::try_parse_from([ - "test", - "codex", - "session-a", - "--destination", - "project_logs:project-id", - "--parent-span-id", - "span-1", - "--parent-root-span-id", - "root-1", - "--parent-object-type", - "project_logs", - "--parent-project", - "Agents", - ]) - .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, - parent_span_id: None, - parent_root_span_id: None, - parent_object_type: None, - parent_object_id: None, - parent_project: None, - attach: true, - additional_metadata: None, - tags: Vec::new(), - plugin: Vec::new(), - }; - assert!(validate_import_selection(&args) - .unwrap_err() - .to_string() - .contains("exactly one")); - } - - #[test] - fn json_string_fields_accept_strings_and_numbers_only() { - let payload = serde_json::json!({ - "string": "session", - "number": 42, - "boolean": true, - "null": null - }); - assert_eq!( - json_str_field(&payload, "string").as_deref(), - Some("session") - ); - assert_eq!(json_str_field(&payload, "number").as_deref(), Some("42")); - assert_eq!(json_str_field(&payload, "boolean"), None); - assert_eq!(json_str_field(&payload, "null"), None); - assert_eq!(json_str_field(&payload, "missing"), None); - } - - #[test] - fn clock_returns_a_positive_epoch_timestamp() { - assert!(now_ms() > 0); - } - - #[test] - fn transcript_observation_captures_compact_and_full_boundaries() { - let dir = tempfile::tempdir().unwrap(); - let compact = dir.path().join("transcript.jsonl"); - let full = dir.path().join("transcript_full.jsonl"); - std::fs::write(&compact, b"compact\n").unwrap(); - std::fs::write(&full, b"complete record\n").unwrap(); - let mut payload = serde_json::json!({ - "transcriptPath": compact.to_string_lossy() - }); - - add_transcript_observation(&mut payload, "transcriptPath"); - - let observed = &payload["_bt_transcript_observation"]; - assert_eq!(observed["path"], compact.to_string_lossy().as_ref()); - assert_eq!(observed["observed_bytes"], 8); - assert_eq!(observed["full_path"], full.to_string_lossy().as_ref()); - assert_eq!(observed["full_observed_bytes"], 16); - } - - #[test] - fn import_destination_without_session_config_fails_fast() { - let mut config = None; - let destination = wire::TraceDestination::ProjectLogs { - project_id: Some("project-id".to_string()), - project_name: None, - }; - - let error = apply_import_destination(&mut config, Some(destination)).unwrap_err(); - - assert!(error - .to_string() - .contains("import destination requires a resolved Braintrust session configuration")); - } - - fn import_test_config(project_id: &str) -> SessionConfig { - SessionConfig { - auth: wire::BackendAuth { - token: "test".into(), - api_url: Some("https://api.example.test".into()), - app_url: None, - org_name: Some("test-org".into()), - org_id: Some("test-org-id".into()), - }, - destination: Some(wire::TraceDestination::ProjectLogs { - project_id: Some(project_id.into()), - project_name: None, - }), - flush_mode: wire::FlushMode::FireAndForget, - additional_metadata: None, - tags: Vec::new(), - span_plugins: Vec::new(), - } - } - - fn imported_terminal_span_ids(path: &std::path::Path) -> Vec { - std::fs::read_to_string(path) - .unwrap() - .lines() - .filter_map(|line| serde_json::from_str::(line).ok()) - .filter_map(|op| { - let row = op.get("Insert").or_else(|| op.get("Merge"))?; - row.get("end_ms")?.as_i64()?; - row.get("span_id")?.as_str().map(str::to_string) - }) - .collect() - } - - #[tokio::test] - async fn imports_share_destination_delivery_ledger_and_replay_to_new_destinations() { - let temp = tempfile::tempdir().unwrap(); - let transcript = temp.path().join("session.jsonl"); - std::fs::write( - &transcript, - format!( - "{}\n{}\n{}\n{}\n", - json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"ledger-import","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":"turn-1"}}), - json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"done"}}), - ), - ) - .unwrap(); - let ledger_dir = temp.path().join("ledger"); - let first_output = temp.path().join("first"); - import_transcript_with_ledger( - &transcript, - ImportSource::Codex, - debug_serve_options("test", &first_output), - Some(import_test_config("project-a")), - false, - Some(ledger_dir.clone()), - ) - .await - .unwrap(); - let spans = first_output.join("spans/ledger-import.ndjson"); - let first_terminals = imported_terminal_span_ids(&spans); - assert!(!first_terminals.is_empty()); - let first_operations = std::fs::read_to_string(&spans).unwrap(); - - import_transcript_with_ledger( - &transcript, - ImportSource::Codex, - debug_serve_options("test", &first_output), - Some(import_test_config("project-a")), - false, - Some(ledger_dir.clone()), - ) - .await - .unwrap(); - assert_eq!(imported_terminal_span_ids(&spans), first_terminals); - assert_eq!(std::fs::read_to_string(&spans).unwrap(), first_operations); - - let second_output = temp.path().join("second"); - import_transcript_with_ledger( - &transcript, - ImportSource::Codex, - debug_serve_options("test", &second_output), - Some(import_test_config("project-b")), - false, - Some(ledger_dir), - ) - .await - .unwrap(); - assert_eq!( - imported_terminal_span_ids(&second_output.join("spans/ledger-import.ndjson")), - first_terminals - ); - } - - #[tokio::test] - async fn repeated_muse_import_does_not_emit_a_second_trace() { - let temp = tempfile::tempdir().unwrap(); - let export = temp.path().join("muse-export.json"); - std::fs::write( - &export, - serde_json::to_vec(&json!({ - "export_schema_version":1, - "sessions":[{"session_id":"muse-ledger"}], - "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":"terminal","terminal":"completed"}}}}, - {"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}} - ] - })) - .unwrap(), - ) - .unwrap(); - let output = temp.path().join("output"); - let ledger = temp.path().join("ledger"); - let import = || { - import_transcript_with_ledger( - &export, - ImportSource::Muse, - debug_serve_options("test", &output), - Some(import_test_config("project-a")), - false, - Some(ledger.clone()), - ) - }; - import().await.unwrap(); - let rows = output.join("spans/muse-ledger.ndjson"); - let first = std::fs::read_to_string(&rows).unwrap(); - assert!(first.contains("Muse Code")); - import().await.unwrap(); - assert_eq!(std::fs::read_to_string(&rows).unwrap(), first); - } - - #[tokio::test] - async fn later_muse_import_completes_an_active_snapshot_on_the_same_root() { - let temp = tempfile::tempdir().unwrap(); - let export = temp.path().join("muse-export.json"); - let output = temp.path().join("output"); - let ledger = temp.path().join("ledger"); - let started = json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}); - let write = |events: Vec| { - std::fs::write( - &export, - serde_json::to_vec(&json!({ - "export_schema_version":1, - "sessions":[{"session_id":"muse-active-import"}], - "events":events, - })) - .unwrap(), - ) - .unwrap(); - }; - let import = || { - import_transcript_with_ledger( - &export, - ImportSource::Muse, - debug_serve_options("test", &output), - Some(import_test_config("project-a")), - false, - Some(ledger.clone()), - ) - }; - write(vec![started.clone()]); - let summaries = import().await.unwrap(); - assert!(!summaries[0].finalized); - let rows_path = output.join("spans/muse-active-import.ndjson"); - let first = std::fs::read_to_string(&rows_path).unwrap(); - assert!(!first.contains("\"Merge\"")); - - write(vec![ - started, - json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), - json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}}), - ]); - let summaries = import().await.unwrap(); - assert!(summaries[0].finalized); - let rows = std::fs::read_to_string(&rows_path).unwrap(); - let parsed = rows - .lines() - .map(|line| serde_json::from_str::(line).unwrap()) - .collect::>(); - let root_id = ids::span_id( - &ids::session_namespace("muse", "muse-active-import"), - "root", - ); - assert_eq!( - parsed - .iter() - .filter(|row| row - .pointer("/Insert/span_id") - .and_then(|value| value.as_str()) - == Some(root_id.as_str())) - .count(), - 2 - ); - assert_eq!( - parsed - .iter() - .filter(|row| row - .pointer("/Merge/span_id") - .and_then(|value| value.as_str()) - == Some(root_id.as_str())) - .count(), - 1 - ); - assert!(rows.contains("\"end_ms\":1002")); - let completed = rows.clone(); - import().await.unwrap(); - assert_eq!(std::fs::read_to_string(&rows_path).unwrap(), completed); - } - - fn test_run_hook_command() -> RunHookCommand { - RunHookCommand { - program: OsString::from("/opt/Braintrust CLI/bt"), - args: vec![OsString::from("agents"), OsString::from("hook")], - } - } - - #[tokio::test] - async fn managed_run_rejects_a_missing_destination_before_launch() { - let error = run_traced( - RunArgs { - source: RunSource::Codex, - additional_metadata: None, - tags: Vec::new(), - plugin: Vec::new(), - agent_args: Vec::new(), - }, - test_run_hook_command(), - SessionRoute::default(), - HostInfo { - serve_argv: Vec::new(), - version: "test".into(), - }, - ) - .await - .unwrap_err(); - - assert!(error.to_string().contains("requires a trace destination")); - } - - #[tokio::test] - async fn codex_managed_run_rejects_global_hook_trust_bypass_before_launch() { - let error = run_traced( - RunArgs { - source: RunSource::Codex, - additional_metadata: None, - tags: Vec::new(), - plugin: Vec::new(), - agent_args: vec![OsString::from("--dangerously-bypass-hook-trust")], - }, - test_run_hook_command(), - SessionRoute { - destination: Some(wire::TraceDestination::ProjectLogs { - project_id: Some("project-id".into()), - project_name: None, - }), - ..SessionRoute::default() - }, - HostInfo { - serve_argv: Vec::new(), - version: "test".into(), - }, - ) - .await - .unwrap_err(); - - assert!(error - .to_string() - .contains("cannot be combined with --dangerously-bypass-hook-trust")); - } - - #[test] - fn codex_managed_run_injects_live_hooks() { - let args = managed_run_args(RunSource::Codex, &test_run_hook_command()).unwrap(); - assert_eq!(args[0], "--enable"); - assert_eq!(args[1], "hooks"); - assert!(!args - .iter() - .any(|arg| arg == "--dangerously-bypass-hook-trust")); - assert_eq!( - args.iter().filter(|arg| *arg == "-c").count(), - CODEX_RUN_HOOK_EVENTS.len() - ); - let config = args - .iter() - .find_map(|arg| { - let arg = arg.to_str()?; - arg.starts_with("hooks.SessionStart=").then_some(arg) - }) - .unwrap(); - assert!(config.contains("--managed-run-hook")); - assert!(config.contains("agents")); - assert!(config.contains("hook")); - assert!(config.contains("--source")); - assert!(config.contains("codex")); - assert!(!config.contains("transcript")); - } - - #[test] - fn claude_managed_run_injects_live_hooks() { - let args = managed_run_args(RunSource::Claude, &test_run_hook_command()).unwrap(); - assert_eq!(args[0], "--settings"); - let settings: serde_json::Value = serde_json::from_str(args[1].to_str().unwrap()).unwrap(); - let hooks = settings["hooks"].as_object().unwrap(); - assert_eq!(hooks.len(), CLAUDE_RUN_HOOK_EVENTS.len()); - for event in ["SessionStart", "PreToolUse"] { - let matcher_groups = hooks[event].as_array().unwrap(); - assert_eq!(matcher_groups.len(), 1); - assert!(matcher_groups[0]["hooks"].is_array()); - assert!(matcher_groups[0]["hooks"][0]["hooks"].is_null()); - } - let command = hooks["SessionStart"][0]["hooks"][0]["command"] - .as_str() - .unwrap(); - assert!(command.contains("--managed-run-hook")); - assert!(command.contains("agents")); - assert!(command.contains("hook")); - assert!(command.contains("--source")); - assert!(command.contains("claude")); - assert!(!command.contains("transcript")); - } - - #[test] - fn opencode_managed_run_preserves_inline_config_and_adds_plugin() { - let config = - opencode_managed_config(Some(r#"{"model":"test/model","plugin":["other"]}"#)).unwrap(); - let config: serde_json::Value = serde_json::from_str(&config).unwrap(); - assert_eq!(config["model"], "test/model"); - assert_eq!( - config["plugin"], - serde_json::json!(["other", "@braintrust/trace-opencode/tracing"]) - ); - assert!( - managed_run_args(RunSource::OpenCode, &test_run_hook_command()) - .unwrap() - .is_empty() - ); - } - - #[test] - fn pi_managed_run_loads_the_npm_extension() { - assert_eq!( - managed_run_args(RunSource::Pi, &test_run_hook_command()).unwrap(), - vec![ - OsString::from("-e"), - OsString::from(crate::setup::pi_plugin_spec()), - ] - ); - } - - #[test] - fn managed_hook_commands_quote_frontend_paths() { - let hook = test_run_hook_command(); - let unix = managed_hook_shell_command(&hook, "codex", false).unwrap(); - assert!(unix.contains("'/opt/Braintrust CLI/bt' 'agents' 'hook' '--source' 'codex'")); - let windows = managed_hook_shell_command(&hook, "claude", true).unwrap(); - assert!(windows - .contains("\"/opt/Braintrust CLI/bt\" \"agents\" \"hook\" \"--source\" \"claude\"")); - } - - #[tokio::test] - async fn muse_managed_hook_rejects_missing_context_before_reading_stdin() { - let temp = tempfile::tempdir().unwrap(); - let args = HookCli::try_parse_from([ - "test", - "--source", - "muse", - "--managed-run-hook", - "--managed-context", - temp.path().join("missing.json").to_str().unwrap(), - ]) - .unwrap() - .args; - let error = run_hook( - args, - SessionRoute::default(), - HostInfo { - serve_argv: Vec::new(), - version: "test".into(), - }, - ) - .await - .unwrap_err(); - assert_eq!( - error.downcast_ref::().unwrap().kind(), - std::io::ErrorKind::NotFound - ); - } -} diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index a0c9ac9f..2af7390b 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 coding-agent session by its resume id. + /// Import a past Codex or Claude Code session by its resume id. Import(ImportArgs), /// Launch a coding agent with live tracing hooks for this invocation. Run { @@ -300,7 +300,7 @@ async fn main() { program: exe, args: vec![OsString::from("hook")], }; - match run_traced(args, hook_command, route.into_route(), host_info()).await { + match run_traced(args, hook_command, route.into_route()).await { Ok(status) if status.success() => {} Ok(status) => std::process::exit(status.code().unwrap_or(1)), Err(error) => { diff --git a/bt-daemon/src/managed_run/mod.rs b/bt-daemon/src/managed_run/mod.rs index 1f3716f4..a35ac5dd 100644 --- a/bt-daemon/src/managed_run/mod.rs +++ b/bt-daemon/src/managed_run/mod.rs @@ -4,6 +4,7 @@ mod claude; mod codex; mod cursor; +pub(crate) mod muse; mod opencode; mod pi; @@ -39,6 +40,22 @@ pub(crate) trait ManagedRun: Agent { hook_command: &RunHookCommand, managed_run_id: &str, ) -> anyhow::Result; + + /// Agents that sanitize hook environments need an already-running daemon + /// and a private context file instead of relying on inherited settings. + fn needs_daemon_host(&self) -> bool { + false + } + + fn inject_with_context( + &self, + hook: &RunHookCommand, + id: &str, + _route: &SessionRoute, + _socket: &std::path::Path, + ) -> anyhow::Result { + self.inject(hook, id) + } } /// Whether the managed run's stdin and stdout are terminals. @@ -62,9 +79,28 @@ pub(crate) struct Injection { /// this invocation, without requiring the tracing plugin to be installed or /// enabled globally. pub async fn run_traced( + args: RunArgs, + hook_command: RunHookCommand, + route: SessionRoute, +) -> anyhow::Result { + run_traced_inner(args, hook_command, route, None).await +} + +/// Launch an agent whose sanitized hooks require the embedding daemon host. +pub async fn run_traced_with_host( + args: RunArgs, + hook_command: RunHookCommand, + route: SessionRoute, + host: crate::HostInfo, +) -> anyhow::Result { + run_traced_inner(args, hook_command, route, Some(host)).await +} + +async fn run_traced_inner( args: RunArgs, hook_command: RunHookCommand, mut route: SessionRoute, + host: Option, ) -> anyhow::Result { apply_run_span_plugins(&mut route, &args.plugin)?; let name = args.source.identity().id; @@ -88,16 +124,39 @@ pub async fn run_traced( let executable = std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); let managed_run_id = uuid::Uuid::new_v4().to_string(); - let injection = agent.inject(&hook_command, &managed_run_id)?; // A caller-supplied socket is already an explicit daemon boundary (for // example an integration harness or a deliberately isolated runtime). // Only create our own boundary when environment auth would otherwise use // the ambient shared daemon. - let isolate_daemon = route.auth.effective_source() == wire::AuthSource::Environment + let isolate_daemon = (agent.needs_daemon_host() + || route.auth.effective_source() == wire::AuthSource::Environment) && env::var_os(paths::SOCKET_ENV).is_none(); let isolated_runtime = isolate_daemon .then(|| ManagedRunRuntime::new(&managed_run_id)) .transpose()?; + let socket = isolated_runtime + .as_ref() + .map(|runtime| runtime.socket.clone()) + .unwrap_or_else(|| paths::socket_path(None)); + let injection = agent.inject_with_context(&hook_command, &managed_run_id, &route, &socket)?; + if agent.needs_daemon_host() { + if let Some(runtime) = &isolated_runtime { + let host = host.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "Muse managed runs require run_traced_with_host to start the invocation daemon" + ) + })?; + drop( + crate::client::ensure_daemon_in( + &runtime.socket, + host, + false, + Some(runtime.temp_dir.path()), + ) + .await?, + ); + } + } let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?; let mut command: tokio::process::Command = subprocess::interactive_command(&executable).into(); command @@ -113,9 +172,18 @@ pub async fn run_traced( } command.envs(injection.env); let _guards = injection.guards; - let mut child = command.spawn().map_err(|error| { - anyhow::anyhow!("failed to launch {}: {error}", executable.to_string_lossy()) - })?; + let mut child = match command.spawn() { + Ok(child) => child, + Err(error) => { + if agent.needs_daemon_host() && isolated_runtime.is_some() { + let _ = shutdown_daemon(&socket).await; + } + return Err(anyhow::anyhow!( + "failed to launch {}: {error}", + executable.to_string_lossy() + )); + } + }; let interrupt = tokio::signal::ctrl_c(); tokio::pin!(interrupt); @@ -240,6 +308,7 @@ mod tests { fn every_run_source_has_a_registered_managed_run() { for source in [ RunSource::Codex, + RunSource::Muse, RunSource::Cursor, RunSource::Claude, RunSource::OpenCode, diff --git a/bt-daemon/src/managed_run/muse.rs b/bt-daemon/src/managed_run/muse.rs new file mode 100644 index 00000000..cef246da --- /dev/null +++ b/bt-daemon/src/managed_run/muse.rs @@ -0,0 +1,97 @@ +//! Invocation-local Muse hooks and their route context. +use super::{Injection, ManagedRun, Terminal}; +use crate::agents::Muse; +use crate::args::RunHookCommand; +use crate::wire::SessionRoute; +use anyhow::Context; +use std::ffi::OsString; +use std::path::{Path, PathBuf}; + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +pub(crate) struct MuseManagedContext { + pub route: SessionRoute, + pub managed_run_id: String, + pub socket: PathBuf, +} + +impl ManagedRun for Muse { + fn executable(&self) -> (&'static str, &'static str) { + ("MUSE_BIN", "muse") + } + + fn check(&self, args: &[OsString], _terminal: Terminal) -> anyhow::Result<()> { + crate::setup::muse::check_muse_project_hooks(&std::env::current_dir()?)?; + for (index, arg) in args.iter().enumerate() { + let Some(arg) = arg.to_str() else { continue }; + let workspace = if arg == "--workspace" { + args.get(index + 1).map(PathBuf::from) + } else { + arg.strip_prefix("--workspace=").map(PathBuf::from) + }; + if let Some(workspace) = workspace { + crate::setup::muse::check_muse_project_hooks(&workspace)?; + } + } + let executable = crate::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()); + let output = crate::subprocess::background_command(executable) + .arg("--version") + .output() + .context("check Muse Code version for managed hooks")?; + if !output.status.success() + || crate::setup::muse::muse_version(&output.stdout)? != semver::Version::new(1, 1, 1) + { + anyhow::bail!("Muse managed runs require the verified Muse Code 1.1.1 hook override"); + } + Ok(()) + } + + fn needs_daemon_host(&self) -> bool { + true + } + + fn inject(&self, _hook: &RunHookCommand, _id: &str) -> anyhow::Result { + anyhow::bail!("Muse managed hooks require an invocation route and socket") + } + + fn inject_with_context( + &self, + hook: &RunHookCommand, + id: &str, + route: &SessionRoute, + socket: &Path, + ) -> anyhow::Result { + let directory = tempfile::Builder::new() + .prefix("bt-trace-run-muse-") + .tempdir()?; + crate::paths::ensure_private_dir(directory.path())?; + let context_path = directory.path().join("context.json"); + let context = MuseManagedContext { + route: route.clone(), + managed_run_id: id.into(), + socket: socket.into(), + }; + std::fs::write(&context_path, serde_json::to_vec(&context)?)?; + crate::paths::restrict_file_to_owner(&context_path)?; + let mut hook = hook.clone(); + hook.args.extend([ + "--flush-on-turn-end".into(), + "--managed-context".into(), + context_path.into_os_string(), + ]); + let command = super::managed_hook_shell_command(&hook, "muse", cfg!(windows))?; + let hooks = crate::setup::muse::muse_managed_hook_config( + &crate::paths::muse_config_dir(), + &command, + )?; + let hooks_path = directory.path().join("hooks.json"); + std::fs::write(&hooks_path, serde_json::to_vec(&hooks)?)?; + Ok(Injection { + env: vec![( + "TBH_MANAGED_HOOKS_PATH", + hooks_path.to_string_lossy().into_owned(), + )], + guards: vec![Box::new(directory)], + ..Injection::default() + }) + } +} diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs deleted file mode 100644 index 96dc84ac..00000000 --- a/bt-daemon/src/setup.rs +++ /dev/null @@ -1,2277 +0,0 @@ -//! Persistent installation and configuration for coding-agent tracing plugins. - -use crate::paths; -use crate::trace_command::{EnableArgs, SetupAgent}; -use crate::wire::SessionRoute; -use crate::TraceCommandOutput; -use anyhow::{bail, Context}; -use fs2::FileExt; -use serde_json::{Map, Value}; -use std::io::Write; -use std::path::{Path, PathBuf}; -use std::process::Command as ProcessCommand; - -const CODEX_MARKETPLACE: &str = "braintrust-codex-plugins"; -const CODEX_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-codex-plugin"; -const CODEX_PLUGIN: &str = "trace-codex@braintrust-codex-plugins"; -const CLAUDE_MARKETPLACE: &str = "braintrust-claude-plugin"; -const CLAUDE_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-claude-plugin"; -const CLAUDE_PLUGIN: &str = "trace-claude-code@braintrust-claude-plugin"; -const GROK_PLUGIN: &str = "trace-grok"; -const GROK_PLUGIN_SOURCE: &str = "braintrustdata/braintrust-grok-plugin"; -const OPENCODE_PACKAGE: &str = "@braintrust/trace-opencode"; -const PI_PACKAGE: &str = "@braintrust/pi-extension"; -const OPENCODE_PACKAGE_MANIFEST: &str = - include_str!("../../src/plugins/opencode/content/package.json"); -const ANTIGRAVITY_PLUGIN: &str = "braintrust-antigravity-tracing"; -const MUSE_HOOK_EVENTS: &[&str] = &[ - "SessionStart", - "UserPromptSubmit", - "PreLLMCall", - "PostLLMCall", - "Stop", - "SessionEnd", -]; -const LEGACY_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"]; -const ANTIGRAVITY_PLUGIN_SOURCE: &str = - "https://github.com/braintrustdata/braintrust-antigravity-plugin"; - -fn package_version(manifest: &str) -> anyhow::Result { - let manifest = serde_json::from_str::(manifest)?; - manifest - .get("version") - .and_then(Value::as_str) - .map(str::to_owned) - .ok_or_else(|| anyhow::anyhow!("package manifest has no string version")) -} - -fn npm_major_spec(package: &str, manifest: &str) -> anyhow::Result { - let version = package_version(manifest)?; - let major = version - .split_once('.') - .map(|(major, _)| major) - .ok_or_else(|| anyhow::anyhow!("package version is not semver: {version}"))? - .parse::() - .with_context(|| format!("package version is not semver: {version}"))?; - Ok(format!("{package}@^{major}")) -} - -fn opencode_plugin_spec() -> anyhow::Result { - npm_major_spec(OPENCODE_PACKAGE, OPENCODE_PACKAGE_MANIFEST) -} - -pub(crate) fn pi_plugin_spec() -> String { - format!("npm:{PI_PACKAGE}") -} - -fn version_is_older(installed: &str, expected: &str) -> bool { - let parse = |version: &str| { - let mut parts = version.split(['.', '-', '+']); - Some(( - parts.next()?.parse::().ok()?, - parts.next().unwrap_or("0").parse::().ok()?, - parts.next().unwrap_or("0").parse::().ok()?, - )) - }; - parse(installed) - .zip(parse(expected)) - .is_some_and(|(installed, expected)| installed < expected) -} - -pub(crate) fn update_warning(source: &str) -> Option { - let stale = match source { - "codex" => installed_json_version("codex", &["plugin", "list", "--json"], |value| { - codex_plugin(value).and_then(|plugin| plugin.get("version")).and_then(Value::as_str) - }, &package_version(include_str!("../../src/plugins/codex/content/plugins/trace-codex/.codex-plugin/plugin.json")).ok()?), - "claude" => installed_json_version("claude", &["plugin", "list", "--json"], |value| { - claude_plugin(value).and_then(|plugin| plugin.get("version")).and_then(Value::as_str) - }, &package_version(include_str!("../../src/plugins/claude/content/plugins/trace-claude-code/.claude-plugin/plugin.json")).ok()?), - "opencode" => opencode_update_required(), - "pi" => pi_update_required(), - _ => false, - }; - stale.then(|| format!("tracing plugin is out of date; run `bt trace update {source}`")) -} - -fn installed_json_version( - program: &str, - args: &[&str], - find_version: impl Fn(&Value) -> Option<&str>, - expected: &str, -) -> bool { - let output = ProcessCommand::new(program).args(args).output(); - let Ok(output) = output else { return false }; - if !output.status.success() { - return false; - } - serde_json::from_slice::(&output.stdout) - .ok() - .and_then(|value| find_version(&value).map(str::to_owned)) - .is_some_and(|installed| version_is_older(&installed, expected)) -} - -fn opencode_update_required() -> bool { - let settings_path = paths::agent_settings_path("opencode", None); - let path = settings_path - .parent() - .unwrap_or_else(|| Path::new(".")) - .join("opencode.json"); - let Ok(raw) = std::fs::read(path) else { - return false; - }; - let Ok(config) = serde_json::from_slice::(&raw) else { - return false; - }; - let expected = opencode_plugin_spec().ok(); - config - .get("plugin") - .and_then(Value::as_array) - .and_then(|plugins| { - plugins.iter().find_map(|plugin| { - plugin - .as_str() - .filter(|plugin| plugin.starts_with(OPENCODE_PACKAGE)) - }) - }) - .is_some_and(|plugin| Some(plugin) != expected.as_deref()) -} - -fn pi_update_required() -> bool { - let output = ProcessCommand::new("pi").arg("list").output(); - let Ok(output) = output else { return false }; - if !output.status.success() { - return false; - } - let installed = String::from_utf8_lossy(&output.stdout); - let expected = pi_plugin_spec(); - installed.lines().any(|line| { - let plugin = line.trim(); - plugin.starts_with(&expected) && plugin != expected - }) -} - -trait CommandRunner { - fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result; - fn json_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result; - fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()>; - fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()>; -} - -struct SystemCommandRunner; - -impl CommandRunner for SystemCommandRunner { - fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { - let output = ProcessCommand::new(program) - .args(args) - .output() - .with_context(|| { - format!("failed to run `{program}`; install {program} and ensure it is on PATH") - })?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); - } - serde_json::from_slice(&output.stdout) - .with_context(|| format!("`{program} {}` returned invalid JSON", args.join(" "))) - } - - fn json_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result { - let output = ProcessCommand::new(program) - .args(args) - .env("HOME", home) - .output() - .with_context(|| { - format!("failed to run `{program}`; install {program} and ensure it is on PATH") - })?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); - } - serde_json::from_slice(&output.stdout) - .with_context(|| format!("`{program} {}` returned invalid JSON", args.join(" "))) - } - - fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { - let status = ProcessCommand::new(program) - .args(args) - .status() - .with_context(|| { - format!("failed to run `{program}`; install {program} and ensure it is on PATH") - })?; - if !status.success() { - bail!("`{program} {}` failed with {status}", args.join(" ")); - } - Ok(()) - } - - fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()> { - let output = ProcessCommand::new(program) - .args(args) - .env("HOME", home) - .output() - .with_context(|| { - format!("failed to run `{program}`; install {program} and ensure it is on PATH") - })?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); - } - Ok(()) - } -} - -fn github_repo_matches(source: &str, expected: &str) -> bool { - let source = source.trim().trim_end_matches('/'); - let source = source.strip_suffix(".git").unwrap_or(source); - let source = source - .strip_prefix("https://github.com/") - .or_else(|| source.strip_prefix("git@github.com:")) - .unwrap_or(source); - source == expected -} - -fn codex_marketplace(value: &Value) -> Option<&Value> { - value - .get("marketplaces") - .and_then(Value::as_array)? - .iter() - .find(|item| item.get("name").and_then(Value::as_str) == Some(CODEX_MARKETPLACE)) -} - -fn codex_marketplace_is_published(item: &Value) -> bool { - item.get("marketplaceSource") - .and_then(|source| source.get("source")) - .and_then(Value::as_str) - .is_some_and(|source| github_repo_matches(source, CODEX_MARKETPLACE_SOURCE)) -} - -fn setup_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let marketplaces = runner.json("codex", &["plugin", "marketplace", "list", "--json"])?; - match codex_marketplace(&marketplaces) { - Some(marketplace) if codex_marketplace_is_published(marketplace) => runner.run( - "codex", - &["plugin", "marketplace", "upgrade", CODEX_MARKETPLACE], - )?, - Some(_) => { - runner.run( - "codex", - &["plugin", "marketplace", "remove", CODEX_MARKETPLACE], - )?; - runner.run( - "codex", - &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], - )?; - } - None => runner.run( - "codex", - &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], - )?, - } - - // Adding is idempotent and reconciles the installed cache to the refreshed - // marketplace snapshot. - runner.run("codex", &["plugin", "add", CODEX_PLUGIN]) -} - -fn codex_plugin(value: &Value) -> Option<&Value> { - value - .get("installed") - .and_then(Value::as_array)? - .iter() - .find(|item| item.get("pluginId").and_then(Value::as_str) == Some(CODEX_PLUGIN)) -} - -fn disable_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("codex", &["plugin", "list", "--json"])?; - if codex_plugin(&plugins).is_some() { - runner.run("codex", &["plugin", "remove", CODEX_PLUGIN, "--json"])?; - } - Ok(()) -} - -fn update_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("codex", &["plugin", "list", "--json"])?; - if codex_plugin(&plugins).is_none() { - bail!("Codex tracing plugin is not installed; run `bt trace enable codex`"); - } - let marketplaces = runner.json("codex", &["plugin", "marketplace", "list", "--json"])?; - let marketplace = codex_marketplace(&marketplaces).ok_or_else(|| { - anyhow::anyhow!("Codex tracing marketplace is not installed; run `bt trace enable codex`") - })?; - if !codex_marketplace_is_published(marketplace) { - bail!("Codex tracing marketplace is not the published Braintrust marketplace; run `bt trace enable codex`"); - } - runner.run( - "codex", - &["plugin", "marketplace", "upgrade", CODEX_MARKETPLACE], - )?; - runner.run("codex", &["plugin", "add", CODEX_PLUGIN]) -} - -fn claude_marketplace(value: &Value) -> Option<&Value> { - value - .as_array()? - .iter() - .find(|item| item.get("name").and_then(Value::as_str) == Some(CLAUDE_MARKETPLACE)) -} - -fn claude_marketplace_is_published(item: &Value) -> bool { - item.get("source").and_then(Value::as_str) == Some("github") - && item - .get("repo") - .and_then(Value::as_str) - .is_some_and(|repo| github_repo_matches(repo, CLAUDE_MARKETPLACE_SOURCE)) -} - -fn claude_plugin(value: &Value) -> Option<&Value> { - value - .as_array()? - .iter() - .find(|item| item.get("id").and_then(Value::as_str) == Some(CLAUDE_PLUGIN)) -} - -fn setup_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let marketplaces = runner.json("claude", &["plugin", "marketplace", "list", "--json"])?; - let marketplace_replaced = match claude_marketplace(&marketplaces) { - Some(marketplace) if claude_marketplace_is_published(marketplace) => { - runner.run( - "claude", - &["plugin", "marketplace", "update", CLAUDE_MARKETPLACE], - )?; - false - } - Some(_) => { - runner.run( - "claude", - &["plugin", "marketplace", "remove", CLAUDE_MARKETPLACE], - )?; - runner.run( - "claude", - &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], - )?; - true - } - None => { - runner.run( - "claude", - &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], - )?; - false - } - }; - - // Claude removes a marketplace's installed plugins when that marketplace - // is removed, so replacing a stale source requires a fresh installation. - if marketplace_replaced { - return runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]); - } - - let plugins = runner.json("claude", &["plugin", "list", "--json"])?; - match claude_plugin(&plugins) { - None => runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]), - Some(plugin) => { - runner.run("claude", &["plugin", "update", CLAUDE_PLUGIN])?; - if plugin.get("enabled").and_then(Value::as_bool) == Some(false) { - runner.run("claude", &["plugin", "enable", CLAUDE_PLUGIN])?; - } - Ok(()) - } - } -} - -fn legacy_claude_tracing_env_keys(path: &Path) -> anyhow::Result> { - let raw = match std::fs::read(path) { - Ok(raw) => raw, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), - Err(error) => { - return Err(error) - .with_context(|| format!("failed to inspect Claude settings: {}", path.display())) - } - }; - let value: Value = serde_json::from_slice(&raw) - .with_context(|| format!("Claude settings are not valid JSON: {}", path.display()))?; - let env = value.get("env").and_then(Value::as_object); - Ok(LEGACY_CLAUDE_TRACING_ENV_KEYS - .into_iter() - .filter(|key| env.is_some_and(|env| env.contains_key(*key))) - .collect()) -} - -fn warn_legacy_claude_tracing_env() { - let path = paths::claude_settings_path(); - let keys = match legacy_claude_tracing_env_keys(&path) { - Ok(keys) => keys, - Err(error) => { - eprintln!("warning: could not inspect legacy Claude tracing settings: {error}"); - return; - } - }; - if !keys.is_empty() { - eprintln!( - "warning: obsolete Braintrust tracing settings in {}: {}; remove these keys from env. BRAINTRUST_API_KEY is left unchanged because the Braintrust MCP plugin may use it.", - path.display(), - keys.join(", ") - ); - } -} - -fn disable_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("claude", &["plugin", "list", "--json"])?; - if claude_plugin(&plugins).is_some() { - runner.run("claude", &["plugin", "uninstall", CLAUDE_PLUGIN])?; - } - Ok(()) -} - -fn grok_plugin(value: &Value) -> Option<&Value> { - value - .as_array()? - .iter() - .find(|item| item.get("name").and_then(Value::as_str) == Some(GROK_PLUGIN)) -} - -fn grok_plugin_is_published(item: &Value) -> bool { - item.get("source") - .and_then(Value::as_str) - .is_some_and(|source| github_repo_matches(source, GROK_PLUGIN_SOURCE)) -} - -fn setup_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("grok", &["plugin", "list", "--json"])?; - // `bt trace enable grok` is the user's trust boundary. Grok's `--trust` - // applies to this plugin installation and does not change folder trust. - match grok_plugin(&plugins) { - Some(plugin) if grok_plugin_is_published(plugin) => { - runner.run("grok", &["plugin", "update", GROK_PLUGIN])?; - } - Some(_) => { - runner.run("grok", &["plugin", "uninstall", GROK_PLUGIN, "--confirm"])?; - runner.run( - "grok", - &["plugin", "install", GROK_PLUGIN_SOURCE, "--trust"], - )?; - } - None => runner.run( - "grok", - &["plugin", "install", GROK_PLUGIN_SOURCE, "--trust"], - )?, - } - runner.run("grok", &["plugin", "enable", GROK_PLUGIN]) -} - -fn disable_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("grok", &["plugin", "list", "--json"])?; - if grok_plugin(&plugins).is_some_and(grok_plugin_is_published) { - runner.run("grok", &["plugin", "uninstall", GROK_PLUGIN, "--confirm"])?; - } - Ok(()) -} - -fn update_grok(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("grok", &["plugin", "list", "--json"])?; - let plugin = grok_plugin(&plugins).ok_or_else(|| { - anyhow::anyhow!("Grok tracing plugin is not installed; run `bt trace enable grok`") - })?; - if !grok_plugin_is_published(plugin) { - bail!("Grok tracing plugin is not the published Braintrust plugin; run `bt trace enable grok`"); - } - runner.run("grok", &["plugin", "update", GROK_PLUGIN]) -} - -fn update_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugins = runner.json("claude", &["plugin", "list", "--json"])?; - if claude_plugin(&plugins).is_none() { - bail!("Claude Code tracing plugin is not installed; run `bt trace enable claude`"); - } - let marketplaces = runner.json("claude", &["plugin", "marketplace", "list", "--json"])?; - let marketplace = claude_marketplace(&marketplaces).ok_or_else(|| { - anyhow::anyhow!( - "Claude Code tracing marketplace is not installed; run `bt trace enable claude`" - ) - })?; - if !claude_marketplace_is_published(marketplace) { - bail!("Claude Code tracing marketplace is not the published Braintrust marketplace; run `bt trace enable claude`"); - } - runner.run( - "claude", - &["plugin", "marketplace", "update", CLAUDE_MARKETPLACE], - )?; - runner.run("claude", &["plugin", "update", CLAUDE_PLUGIN]) -} - -fn load_object(path: &Path) -> anyhow::Result> { - match std::fs::read(path) { - Ok(raw) => { - let value: Value = serde_json::from_slice(&raw) - .with_context(|| format!("invalid JSON configuration: {}", path.display()))?; - value.as_object().cloned().ok_or_else(|| { - anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) - }) - } - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Map::new()), - Err(error) => { - Err(error).with_context(|| format!("failed to read configuration: {}", path.display())) - } - } -} - -fn write_object_atomic(path: &Path, object: Map) -> anyhow::Result<()> { - crate::settings::with_settings_lock(path, || write_object_atomic_unlocked(path, object)) -} - -fn write_object_atomic_unlocked(path: &Path, object: Map) -> anyhow::Result<()> { - let parent = path - .parent() - .ok_or_else(|| anyhow::anyhow!("configuration path has no parent: {}", path.display()))?; - std::fs::create_dir_all(parent).with_context(|| { - format!( - "failed to create configuration directory: {}", - parent.display() - ) - })?; - let mut encoded = serde_json::to_string_pretty(&Value::Object(object))?; - encoded.push('\n'); - let mut temporary = tempfile::NamedTempFile::new_in(parent) - .with_context(|| format!("failed to create temporary file in {}", parent.display()))?; - temporary.write_all(encoded.as_bytes()).with_context(|| { - format!( - "failed to write temporary configuration for {}", - path.display() - ) - })?; - temporary - .persist(path) - .map_err(|error| error.error) - .with_context(|| format!("failed to replace configuration: {}", path.display()))?; - Ok(()) -} - -fn setup_opencode_at(path: &Path) -> anyhow::Result<()> { - let mut config = load_object(path)?; - let plugins = config - .entry("plugin") - .or_insert_with(|| Value::Array(Vec::new())) - .as_array_mut() - .ok_or_else(|| { - anyhow::anyhow!( - "OpenCode `plugin` config must be an array: {}", - path.display() - ) - })?; - plugins.retain(|plugin| { - plugin.as_str().is_none_or(|plugin| { - plugin != "@braintrust/trace-opencode" - && !plugin.starts_with("@braintrust/trace-opencode@") - }) - }); - plugins.push(Value::String(opencode_plugin_spec()?)); - write_object_atomic(path, config) -} - -fn setup_opencode() -> anyhow::Result<()> { - let settings_path = paths::agent_settings_path("opencode", None); - let path = settings_path - .parent() - .unwrap_or_else(|| Path::new(".")) - .join("opencode.json"); - setup_opencode_at(&path) -} - -fn update_opencode_at(path: &Path) -> anyhow::Result<()> { - let mut config = match std::fs::read(path) { - Ok(raw) => serde_json::from_slice::(&raw) - .with_context(|| format!("invalid JSON configuration: {}", path.display()))? - .as_object() - .cloned() - .ok_or_else(|| { - anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) - })?, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => { - bail!("OpenCode tracing plugin is not installed; run `bt trace enable opencode`") - } - Err(error) => { - return Err(error) - .with_context(|| format!("failed to read configuration: {}", path.display())) - } - }; - let plugins = config - .get_mut("plugin") - .and_then(Value::as_array_mut) - .ok_or_else(|| { - anyhow::anyhow!( - "OpenCode tracing plugin is not installed; run `bt trace enable opencode`" - ) - })?; - let found = plugins.iter().any(|plugin| { - plugin.as_str().is_some_and(|plugin| { - plugin == OPENCODE_PACKAGE || plugin.starts_with(&format!("{OPENCODE_PACKAGE}@")) - }) - }); - if !found { - bail!("OpenCode tracing plugin is not installed; run `bt trace enable opencode`"); - } - plugins.retain(|plugin| { - plugin.as_str().is_none_or(|plugin| { - plugin != OPENCODE_PACKAGE && !plugin.starts_with(&format!("{OPENCODE_PACKAGE}@")) - }) - }); - plugins.push(Value::String(opencode_plugin_spec()?)); - write_object_atomic(path, config) -} - -fn update_opencode() -> anyhow::Result<()> { - let settings_path = paths::agent_settings_path("opencode", None); - let path = settings_path - .parent() - .unwrap_or_else(|| Path::new(".")) - .join("opencode.json"); - update_opencode_at(&path) -} - -fn remove_opencode_plugin_at(path: &Path) -> anyhow::Result<()> { - let mut config = match std::fs::read(path) { - Ok(raw) => serde_json::from_slice::(&raw) - .with_context(|| format!("invalid JSON configuration: {}", path.display()))? - .as_object() - .cloned() - .ok_or_else(|| { - anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) - })?, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), - Err(error) => { - return Err(error) - .with_context(|| format!("failed to read configuration: {}", path.display())) - } - }; - let Some(plugins) = config.get_mut("plugin") else { - return Ok(()); - }; - let plugins = plugins.as_array_mut().ok_or_else(|| { - anyhow::anyhow!( - "OpenCode `plugin` config must be an array: {}", - path.display() - ) - })?; - let original_len = plugins.len(); - plugins.retain(|plugin| { - plugin.as_str().is_none_or(|plugin| { - plugin != "@braintrust/trace-opencode" - && !plugin.starts_with("@braintrust/trace-opencode@") - }) - }); - if plugins.len() != original_len { - write_object_atomic(path, config)?; - } - Ok(()) -} - -fn disable_opencode() -> anyhow::Result<()> { - let settings_path = paths::agent_settings_path("opencode", None); - let path = settings_path - .parent() - .unwrap_or_else(|| Path::new(".")) - .join("opencode.json"); - remove_opencode_plugin_at(&path) -} - -fn setup_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugin = pi_plugin_spec(); - runner.run("pi", &["install", &plugin]) -} - -fn disable_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugin = pi_plugin_spec(); - runner.run("pi", &["uninstall", &plugin]) -} - -fn update_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - let plugin = pi_plugin_spec(); - runner.run("pi", &["update", &plugin]) -} - -/// 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 { - // Muse treats command-hook failures as non-fatal to the agent turn. - muse_hook_config_for( - MUSE_HOOK_EVENTS, - "bt trace hook --source muse --flush-on-turn-end", - ) -} - -fn muse_hook_config_for(events: &[&str], command: &str) -> Map { - let hooks = 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)), - ]) -} - -/// Build an invocation-local managed file without changing Muse's settings. -/// Preserve handlers administered by someone else, but never run two Muse -/// Braintrust forwarders with different destinations in the same session. -pub(crate) fn muse_managed_hook_config(config_dir: &Path, command: &str) -> anyhow::Result { - let settings = load_object(&config_dir.join("settings.json"))?; - if settings - .get("hooks") - .is_some_and(muse_contains_braintrust_hook) - { - bail!("Muse user hooks already include Braintrust tracing; remove that duplicate before `bt trace run muse`"); - } - let existing = match settings.get("managed_hooks_path").and_then(Value::as_str) { - Some(path) => { - let path = Path::new(path); - let path = if path.is_absolute() { - path.to_path_buf() - } else { - config_dir.join(path) - }; - if !path.exists() || muse_hook_is_owned(&path)? { - Map::from_iter([("schema_version".into(), Value::from(1))]) - } else { - let raw = std::fs::read(&path) - .with_context(|| format!("read Muse managed hooks {}", path.display()))?; - let value: Value = serde_json::from_slice(&raw)?; - if muse_contains_braintrust_hook(&value) { - bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); - } - if muse_has_strict_handler_fields(&value) { - bail!("Muse managed hooks use the strict handler schema; cannot safely merge invocation hooks"); - } - value - .as_object() - .cloned() - .context("Muse managed hooks must be a JSON object")? - } - } - None => Map::from_iter([("schema_version".into(), Value::from(1))]), - }; - let mut result = existing; - let mut hooks = result - .remove("hooks") - .unwrap_or_else(|| Value::Object(Map::new())); - let hooks = hooks - .as_object_mut() - .context("Muse managed hooks must be an object")?; - let generated = muse_hook_config_for(MUSE_HOOK_EVENTS, command); - for (event, handlers) in generated["hooks"].as_object().unwrap() { - let existing = hooks - .entry(event.clone()) - .or_insert_with(|| Value::Array(Vec::new())); - existing - .as_array_mut() - .with_context(|| format!("Muse managed hook event {event} must be an array"))? - .extend(handlers.as_array().unwrap().iter().cloned()); - } - result.insert("hooks".into(), hooks.clone().into()); - Ok(Value::Object(result)) -} - -pub(crate) fn check_muse_project_hooks(start: &Path) -> anyhow::Result<()> { - for dir in start.ancestors() { - let path = dir.join(".muse").join("hooks.json"); - if !path.exists() { - continue; - } - let hooks: Value = serde_json::from_slice(&std::fs::read(&path)?) - .with_context(|| format!("read Muse project hooks {}", path.display()))?; - if muse_contains_braintrust_hook(&hooks) { - bail!("Muse project hooks {} already include Braintrust tracing; cannot safely run a managed trace", path.display()); - } - break; - } - Ok(()) -} - -fn muse_contains_braintrust_hook(value: &Value) -> bool { - match value { - Value::Object(object) => object.iter().any(|(key, value)| { - (key == "command" - && value.as_str().is_some_and(|command| { - command.contains("--source") - && command.contains("muse") - && command.contains("hook") - })) - || muse_contains_braintrust_hook(value) - }), - Value::Array(items) => items.iter().any(muse_contains_braintrust_hook), - _ => false, - } -} - -fn muse_has_strict_handler_fields(value: &Value) -> bool { - match value { - Value::Object(object) => { - object.contains_key("name") - || object.contains_key("timeout_ms") - || object.values().any(muse_has_strict_handler_fields) - } - Value::Array(items) => items.iter().any(muse_has_strict_handler_fields), - _ => false, - } -} - -fn legacy_muse_hook_config() -> Map { - muse_hook_config_for( - LEGACY_MUSE_HOOK_EVENTS, - "bt trace hook --source muse --flush-on-turn-end || true", - ) -} - -pub(crate) fn muse_version(output: &[u8]) -> anyhow::Result { - let text = std::str::from_utf8(output).context("Muse version is not UTF-8")?; - text.split_whitespace() - .find_map(|word| semver::Version::parse(word).ok()) - .ok_or_else(|| { - anyhow::anyhow!( - "Muse Code returned an unrecognized version: {}", - text.trim() - ) - }) -} - -/// Muse serializes changes to settings.json with this native lock file. Hold -/// it across the read, hook update, and settings update so we cannot discard a -/// concurrent change made by Muse itself. -fn with_muse_settings_lock( - config_dir: &Path, - action: impl FnOnce() -> anyhow::Result, -) -> anyhow::Result { - std::fs::create_dir_all(config_dir) - .with_context(|| format!("create Muse config directory {}", config_dir.display()))?; - let lock_path = config_dir.join(".settings.json.lock"); - let lock = std::fs::OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(&lock_path) - .with_context(|| format!("open Muse settings lock {}", lock_path.display()))?; - lock.lock_exclusive() - .with_context(|| format!("lock Muse settings {}", lock_path.display()))?; - let result = action(); - FileExt::unlock(&lock)?; - result -} - -fn muse_hook_is_owned(path: &Path) -> anyhow::Result { - match std::fs::read(path) { - Ok(raw) => { - let value: Value = serde_json::from_slice(&raw) - .with_context(|| format!("invalid Muse hook configuration: {}", path.display()))?; - Ok(value == Value::Object(muse_hook_config()) - || value == Value::Object(legacy_muse_hook_config())) - } - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), - Err(error) => { - Err(error).with_context(|| format!("read Muse hook configuration: {}", path.display())) - } - } -} - -fn check_muse_hook_ownership(path: &Path) -> anyhow::Result<()> { - if path.exists() && !muse_hook_is_owned(path)? { - bail!("Muse hook file {} differs from Braintrust's generated hooks; refusing to overwrite or remove it", path.display()); - } - Ok(()) -} - -fn setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { - with_muse_settings_lock(config_dir, || { - 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") { - if existing.as_str() != Some(owned_path.as_ref()) { - bail!( - "Muse already uses managed_hooks_path {}; refusing to overwrite it", - existing - ); - } - } - check_muse_hook_ownership(&hook_path)?; - settings.entry("schema_version").or_insert(Value::from(1)); - settings.insert( - "managed_hooks_path".into(), - Value::String(owned_path.into_owned()), - ); - write_object_atomic_unlocked(&hook_path, muse_hook_config())?; - write_object_atomic_unlocked(&settings_path, settings) - }) -} - -fn disable_muse_at(config_dir: &Path) -> anyhow::Result<()> { - with_muse_settings_lock(config_dir, || { - let settings_path = config_dir.join("settings.json"); - let hook_path = muse_hook_path(config_dir); - let mut settings = load_object(&settings_path)?; - check_muse_hook_ownership(&hook_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_unlocked(&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 update_muse_at(config_dir: &Path) -> anyhow::Result<()> { - with_muse_settings_lock(config_dir, || { - let settings = load_object(&config_dir.join("settings.json"))?; - let hook_path = muse_hook_path(config_dir); - if settings.get("managed_hooks_path").and_then(Value::as_str) - != Some(hook_path.to_string_lossy().as_ref()) - { - bail!("Muse tracing hooks are not installed; run `bt trace enable muse`"); - } - if !hook_path.exists() { - bail!("Muse tracing hook file is missing; run `bt trace enable muse`"); - } - check_muse_hook_ownership(&hook_path)?; - write_object_atomic_unlocked(&hook_path, muse_hook_config()) - }) -} - -pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { - let mut warnings = Vec::new(); - match ProcessCommand::new("muse").arg("--version").output() { - Ok(output) if output.status.success() => match muse_version(&output.stdout) { - Ok(version) if version < semver::Version::new(1, 1, 1) => warnings - .push("Muse Code 1.1.1 or newer is required for verified command hooks".into()), - Ok(_) => {} - Err(error) => warnings.push(error.to_string()), - }, - _ => warnings.push( - "Muse Code executable is unavailable; install Muse Code before enabling tracing".into(), - ), - } - if !ProcessCommand::new("bt") - .arg("--version") - .output() - .is_ok_and(|output| output.status.success()) - { - warnings.push( - "Braintrust `bt` executable is unavailable on PATH; Muse command hooks cannot forward events".into(), - ); - } - let settings = match load_object(&config_dir.join("settings.json")) { - Ok(settings) => settings, - Err(error) => { - warnings.push(format!("Muse settings could not be read: {error}")); - return warnings; - } - }; - let hook_path = muse_hook_path(config_dir); - if settings.get("managed_hooks_path").and_then(Value::as_str) - != Some(hook_path.to_string_lossy().as_ref()) - { - warnings.push("Muse managed_hooks_path does not point to Braintrust's hook file; run `bt trace enable muse`".into()); - } - match std::fs::read(&hook_path) { - Ok(raw) => match serde_json::from_slice::(&raw) { - Ok(hooks) if hooks == Value::Object(muse_hook_config()) => {} - Ok(hooks) if hooks == Value::Object(legacy_muse_hook_config()) => warnings.push( - "Muse Braintrust hooks use an older command; run `bt trace update muse`".into(), - ), - Ok(_) => warnings.push("Muse Braintrust hook file differs from expected event definitions; restore it or run `bt trace enable muse`".into()), - Err(error) => warnings.push(format!("Muse Braintrust hook file is invalid JSON: {error}")), - }, - Err(error) => warnings.push(format!("Muse Braintrust hook file could not be read: {error}")), - } - warnings -} - -fn setup_muse() -> anyhow::Result<()> { - let output = ProcessCommand::new("muse") - .arg("--version") - .output() - .context("Muse Code is not installed; install Muse Code before enabling tracing")?; - if !output.status.success() { - bail!("Muse Code could not report its version; install Muse Code before enabling tracing"); - } - let installed = muse_version(&output.stdout)?; - if installed < semver::Version::new(1, 1, 1) { - bail!("Muse Code 1.1.1 or newer is required for verified command hooks"); - } - setup_muse_at(&paths::muse_config_dir()) -} -fn disable_muse() -> anyhow::Result<()> { - disable_muse_at(&paths::muse_config_dir()) -} - -fn update_muse() -> anyhow::Result<()> { - update_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!( - "Antigravity configuration directory must end in `.gemini/config`: {}", - config_dir.display() - ); - } - let gemini_dir = config_dir.parent().ok_or_else(|| { - anyhow::anyhow!( - "Antigravity configuration directory has no parent: {}", - config_dir.display() - ) - })?; - if gemini_dir.file_name().and_then(|part| part.to_str()) != Some(".gemini") { - bail!( - "Antigravity configuration directory must end in `.gemini/config`: {}", - config_dir.display() - ); - } - gemini_dir.parent().ok_or_else(|| { - anyhow::anyhow!( - "Antigravity configuration directory has no home parent: {}", - config_dir.display() - ) - }) -} - -fn remove_legacy_antigravity_registration(config_dir: &Path) -> anyhow::Result<()> { - let hooks_path = config_dir.join("hooks.json"); - if !hooks_path.exists() { - return Ok(()); - } - let mut hooks = load_object(&hooks_path)?; - if hooks.remove(ANTIGRAVITY_PLUGIN).is_some() { - write_object_atomic(&hooks_path, hooks)?; - } - Ok(()) -} - -fn setup_antigravity_at(runner: &mut impl CommandRunner, config_dir: &Path) -> anyhow::Result<()> { - runner.run_in_home( - "agy", - &["plugin", "install", ANTIGRAVITY_PLUGIN_SOURCE], - antigravity_home(config_dir)?, - )?; - runner.run_in_home( - "agy", - &["plugin", "enable", ANTIGRAVITY_PLUGIN], - antigravity_home(config_dir)?, - )?; - - remove_legacy_antigravity_registration(config_dir) -} - -fn setup_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - setup_antigravity_at(runner, &paths::antigravity_config_dir()) -} - -fn disable_antigravity_at( - runner: &mut impl CommandRunner, - config_dir: &Path, -) -> anyhow::Result<()> { - // Antigravity's uninstall command is idempotent when the plugin is absent. - runner.run_in_home( - "agy", - &["plugin", "uninstall", ANTIGRAVITY_PLUGIN], - antigravity_home(config_dir)?, - )?; - remove_legacy_antigravity_registration(config_dir) -} - -fn disable_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - disable_antigravity_at(runner, &paths::antigravity_config_dir()) -} - -fn update_antigravity_at(runner: &mut impl CommandRunner, config_dir: &Path) -> anyhow::Result<()> { - let home = antigravity_home(config_dir)?; - let plugins = runner.json_in_home("agy", &["plugin", "list"], home)?; - let installed = plugins - .get("imports") - .and_then(Value::as_array) - .is_some_and(|imports| { - imports.iter().any(|plugin| { - plugin.get("name").and_then(Value::as_str) == Some(ANTIGRAVITY_PLUGIN) - }) - }); - if !installed { - bail!( - "Google Antigravity tracing plugin is not installed; run `bt trace enable antigravity`" - ); - } - runner.run_in_home( - "agy", - &["plugin", "install", ANTIGRAVITY_PLUGIN_SOURCE], - home, - ) -} - -fn update_antigravity(runner: &mut impl CommandRunner) -> anyhow::Result<()> { - update_antigravity_at(runner, &paths::antigravity_config_dir()) -} - -fn enable_tracing_at(path: &Path, mut route: SessionRoute) -> anyhow::Result<()> { - let mut settings = load_object(path)?; - if route.additional_metadata.is_none() { - route.additional_metadata = settings - .get("route") - .and_then(|route| route.get("additional_metadata")) - .or_else(|| settings.get("additional_metadata")) - .filter(|metadata| metadata.is_object()) - .cloned(); - } - if route.tags.is_empty() { - route.tags = settings - .get("route") - .and_then(|route| route.get("tags")) - .or_else(|| settings.get("tags")) - .and_then(Value::as_array) - .map(|tags| { - tags.iter() - .filter_map(Value::as_str) - .map(ToOwned::to_owned) - .collect() - }) - .unwrap_or_default(); - } - if route.span_plugins.is_empty() { - route.span_plugins = settings - .get("route") - .and_then(|route| route.get("span_plugins")) - .and_then(|plugins| serde_json::from_value(plugins.clone()).ok()) - .unwrap_or_default(); - } - settings.insert("trace_to_braintrust".into(), Value::Bool(true)); - settings.insert("route".into(), serde_json::to_value(route)?); - for key in [ - "traceToBraintrust", - "profile", - "org_name", - "api_key", - "api_url", - "app_url", - "project", - "destination", - "additional_metadata", - "tags", - ] { - settings.remove(key); - } - write_object_atomic(path, settings)?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)) - .with_context(|| format!("failed to protect agent settings: {}", path.display()))?; - } - Ok(()) -} - -fn enable_tracing(source: &str, route: SessionRoute) -> anyhow::Result { - let path = paths::agent_settings_path(source, None); - enable_tracing_at(&path, route)?; - Ok(path) -} - -fn remove_tracing_settings(path: &Path) -> anyhow::Result<()> { - match std::fs::remove_file(path) { - Ok(()) => Ok(()), - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), - Err(error) => Err(error) - .with_context(|| format!("failed to remove tracing settings: {}", path.display())), - } -} - -fn finish_disable(adapter_result: anyhow::Result<()>, settings_path: &Path) -> anyhow::Result<()> { - let settings_result = remove_tracing_settings(settings_path); - adapter_result.and(settings_result) -} - -/// Uninstall an agent's tracing adapter and remove its Braintrust-owned settings. -pub fn run_disable(agent: SetupAgent) -> anyhow::Result { - let mut runner = SystemCommandRunner; - let (source, display_name) = agent_details(agent); - let adapter_result = match agent { - SetupAgent::Codex => disable_codex(&mut runner), - SetupAgent::Claude => disable_claude(&mut runner), - 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); - finish_disable(adapter_result, &settings_path)?; - Ok(TraceCommandOutput::disable( - source, - display_name, - settings_path, - )) -} - -/// Update an already installed tracing adapter without writing tracing settings, -/// enabling a disabled plugin, or creating an agent configuration file. -pub fn run_update(agent: SetupAgent) -> anyhow::Result { - let mut runner = SystemCommandRunner; - let (source, display_name) = agent_details(agent); - match agent { - SetupAgent::Codex => update_codex(&mut runner)?, - SetupAgent::Claude => update_claude(&mut runner)?, - SetupAgent::OpenCode => update_opencode()?, - SetupAgent::Pi => update_pi(&mut runner)?, - SetupAgent::Grok => update_grok(&mut runner)?, - SetupAgent::Muse => update_muse()?, - SetupAgent::Antigravity => update_antigravity(&mut runner)?, - } - Ok(TraceCommandOutput::update(source, display_name)) -} - -fn agent_details(agent: SetupAgent) -> (&'static str, &'static str) { - match agent { - SetupAgent::Codex => ("codex", "Codex"), - SetupAgent::Claude => ("claude", "Claude Code"), - SetupAgent::OpenCode => ("opencode", "OpenCode"), - SetupAgent::Pi => ("pi", "Pi"), - SetupAgent::Grok => ("grok", "Grok"), - SetupAgent::Muse => ("muse", "Muse Code"), - SetupAgent::Antigravity => ("antigravity", "Google Antigravity"), - } -} - -/// Install or refresh one agent's published tracing adapter and persist its -/// non-secret route selection. -pub fn run_enable(args: EnableArgs, route: SessionRoute) -> anyhow::Result { - let mut runner = SystemCommandRunner; - let (source, display_name) = match args.agent { - SetupAgent::Codex => { - setup_codex(&mut runner)?; - ("codex", "Codex") - } - SetupAgent::Claude => { - setup_claude(&mut runner)?; - warn_legacy_claude_tracing_env(); - ("claude", "Claude Code") - } - SetupAgent::OpenCode => { - setup_opencode()?; - ("opencode", "OpenCode") - } - SetupAgent::Pi => { - setup_pi(&mut runner)?; - ("pi", "Pi") - } - SetupAgent::Grok => { - setup_grok(&mut runner)?; - ("grok", "Grok") - } - SetupAgent::Muse => { - setup_muse()?; - ("muse", "Muse Code") - } - SetupAgent::Antigravity => { - setup_antigravity(&mut runner)?; - ("antigravity", "Google Antigravity") - } - }; - let settings_path = enable_tracing(source, route)?; - Ok(TraceCommandOutput::setup( - source, - display_name, - settings_path, - )) -} - -/// Backwards-compatible library entry point for hosts that used the former setup name. -pub fn run_setup(args: EnableArgs, route: SessionRoute) -> anyhow::Result { - run_enable(args, route) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::wire::{AuthSelection, TraceDestination}; - use std::collections::VecDeque; - - struct FakeRunner { - responses: VecDeque, - calls: Vec, - home_calls: Vec<(String, PathBuf)>, - } - - impl FakeRunner { - fn new(responses: impl IntoIterator) -> Self { - Self { - responses: responses.into_iter().collect(), - calls: Vec::new(), - home_calls: Vec::new(), - } - } - - fn called(&self, command: &str) -> bool { - self.calls.iter().any(|call| call == command) - } - } - - impl CommandRunner for FakeRunner { - fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { - self.calls.push(format!("{program} {}", args.join(" "))); - self.responses - .pop_front() - .ok_or_else(|| anyhow::anyhow!("missing fake JSON response")) - } - - fn json_in_home( - &mut self, - program: &str, - args: &[&str], - home: &Path, - ) -> anyhow::Result { - self.calls.push(format!("{program} {}", args.join(" "))); - self.home_calls - .push((format!("{program} {}", args.join(" ")), home.to_path_buf())); - self.responses - .pop_front() - .ok_or_else(|| anyhow::anyhow!("missing fake JSON response")) - } - - fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { - self.calls.push(format!("{program} {}", args.join(" "))); - Ok(()) - } - - fn run_in_home(&mut self, program: &str, args: &[&str], home: &Path) -> anyhow::Result<()> { - self.calls.push(format!("{program} {}", args.join(" "))); - self.home_calls - .push((format!("{program} {}", args.join(" ")), home.to_path_buf())); - Ok(()) - } - } - - struct MissingAgyRunner; - - impl CommandRunner for MissingAgyRunner { - fn json(&mut self, _: &str, _: &[&str]) -> anyhow::Result { - unreachable!() - } - - fn json_in_home(&mut self, _: &str, _: &[&str], _: &Path) -> anyhow::Result { - unreachable!() - } - - fn run(&mut self, _: &str, _: &[&str]) -> anyhow::Result<()> { - unreachable!() - } - - fn run_in_home(&mut self, program: &str, _: &[&str], _: &Path) -> anyhow::Result<()> { - anyhow::bail!("failed to run `{program}`; install {program} and ensure it is on PATH") - } - } - - #[test] - fn codex_installs_from_the_published_marketplace_when_missing() { - let mut runner = FakeRunner::new([serde_json::json!({"marketplaces": []})]); - - setup_codex(&mut runner).unwrap(); - - assert!( - runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") - ); - assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); - } - - #[test] - fn codex_refreshes_the_published_marketplace_and_plugin() { - let mut runner = FakeRunner::new([serde_json::json!({ - "marketplaces": [{ - "name": CODEX_MARKETPLACE, - "marketplaceSource": { - "sourceType": "github", - "source": CODEX_MARKETPLACE_SOURCE - } - }] - })]); - - setup_codex(&mut runner).unwrap(); - - assert!(runner.called("codex plugin marketplace upgrade braintrust-codex-plugins")); - assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); - assert!(!runner.called("codex plugin marketplace remove braintrust-codex-plugins")); - } - - #[test] - fn codex_replaces_a_same_name_local_marketplace() { - let mut runner = FakeRunner::new([serde_json::json!({ - "marketplaces": [{ - "name": CODEX_MARKETPLACE, - "marketplaceSource": {"sourceType": "local", "source": "/tmp/stale"} - }] - })]); - - setup_codex(&mut runner).unwrap(); - - assert!(runner.called("codex plugin marketplace remove braintrust-codex-plugins")); - assert!( - runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") - ); - assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); - } - - #[test] - fn claude_installs_from_the_published_marketplace_when_missing() { - let mut runner = FakeRunner::new([serde_json::json!([]), serde_json::json!([])]); - - setup_claude(&mut runner).unwrap(); - - assert!( - runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") - ); - assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); - } - - #[test] - fn claude_refreshes_the_published_marketplace_and_plugin() { - let mut runner = FakeRunner::new([ - serde_json::json!([{ - "name": CLAUDE_MARKETPLACE, - "source": "github", - "repo": CLAUDE_MARKETPLACE_SOURCE - }]), - serde_json::json!([{ - "id": CLAUDE_PLUGIN, - "version": "1.4.4", - "enabled": true - }]), - ]); - - setup_claude(&mut runner).unwrap(); - - assert!(runner.called("claude plugin marketplace update braintrust-claude-plugin")); - assert!(runner.called("claude plugin update trace-claude-code@braintrust-claude-plugin")); - assert!(!runner.called("claude plugin marketplace remove braintrust-claude-plugin")); - } - - #[test] - fn claude_replaces_a_same_name_local_marketplace() { - let mut runner = FakeRunner::new([serde_json::json!([{ - "name": CLAUDE_MARKETPLACE, - "source": "directory", - "path": "/tmp/stale" - }])]); - - setup_claude(&mut runner).unwrap(); - - assert!(runner.called("claude plugin marketplace remove braintrust-claude-plugin")); - assert!( - runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") - ); - assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); - } - - #[test] - fn claude_updates_then_enables_a_disabled_plugin() { - let mut runner = FakeRunner::new([ - serde_json::json!([{ - "name": CLAUDE_MARKETPLACE, - "source": "github", - "repo": CLAUDE_MARKETPLACE_SOURCE - }]), - serde_json::json!([{"id": CLAUDE_PLUGIN, "enabled": false}]), - ]); - - setup_claude(&mut runner).unwrap(); - - let update = runner - .calls - .iter() - .position(|call| { - call == "claude plugin update trace-claude-code@braintrust-claude-plugin" - }) - .unwrap(); - let enable = runner - .calls - .iter() - .position(|call| { - call == "claude plugin enable trace-claude-code@braintrust-claude-plugin" - }) - .unwrap(); - assert!(update < enable); - } - - #[test] - fn claude_legacy_env_diagnostic_is_read_only_and_preserves_api_key() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("settings.json"); - let contents = r#"{ - "env": { - "BRAINTRUST_CC_PROJECT": "old-project", - "BRAINTRUST_CC_DEBUG": "1", - "BRAINTRUST_API_KEY": "still-used-by-mcp", - "OTHER": "value" - } - }"#; - std::fs::write(&path, contents).unwrap(); - - assert_eq!( - legacy_claude_tracing_env_keys(&path).unwrap(), - ["BRAINTRUST_CC_PROJECT", "BRAINTRUST_CC_DEBUG"] - ); - assert_eq!(std::fs::read_to_string(path).unwrap(), contents); - } - - #[test] - fn grok_installs_published_plugin_then_enables_it_when_missing() { - let mut runner = FakeRunner::new([serde_json::json!([ - {"name": "other-plugin", "source": "somebody/other-plugin"} - ])]); - - setup_grok(&mut runner).unwrap(); - - assert_eq!( - runner.calls, - [ - "grok plugin list --json", - "grok plugin install braintrustdata/braintrust-grok-plugin --trust", - "grok plugin enable trace-grok", - ] - ); - } - - #[test] - fn grok_repeated_enable_updates_and_enables_the_published_plugin() { - let mut runner = FakeRunner::new([serde_json::json!([ - { - "name": GROK_PLUGIN, - "source": "https://github.com/braintrustdata/braintrust-grok-plugin.git", - "status": "installed" - }, - {"name": "other-plugin", "source": "somebody/other-plugin"} - ])]); - - setup_grok(&mut runner).unwrap(); - - assert_eq!( - runner.calls, - [ - "grok plugin list --json", - "grok plugin update trace-grok", - "grok plugin enable trace-grok", - ] - ); - } - - #[test] - fn grok_update_requires_the_published_plugin_without_enabling_it() { - let mut runner = FakeRunner::new([serde_json::json!([{ - "name": GROK_PLUGIN, - "source": GROK_PLUGIN_SOURCE, - "status": "installed" - }])]); - - update_grok(&mut runner).unwrap(); - - assert_eq!( - runner.calls, - ["grok plugin list --json", "grok plugin update trace-grok"] - ); - } - - #[test] - fn grok_enable_reconciles_only_a_conflicting_same_name_plugin() { - let mut runner = FakeRunner::new([serde_json::json!([ - { - "name": GROK_PLUGIN, - "source": "/tmp/local-trace-grok", - "status": "installed" - }, - {"name": "other-plugin", "source": "somebody/other-plugin"} - ])]); - - setup_grok(&mut runner).unwrap(); - - assert_eq!( - runner.calls, - [ - "grok plugin list --json", - "grok plugin uninstall trace-grok --confirm", - "grok plugin install braintrustdata/braintrust-grok-plugin --trust", - "grok plugin enable trace-grok", - ] - ); - } - - #[test] - fn grok_disable_removes_only_the_published_braintrust_plugin() { - let mut runner = FakeRunner::new([serde_json::json!([ - { - "name": GROK_PLUGIN, - "source": GROK_PLUGIN_SOURCE, - "status": "installed" - }, - {"name": "other-plugin", "source": "somebody/other-plugin"} - ])]); - - disable_grok(&mut runner).unwrap(); - - assert_eq!( - runner.calls, - [ - "grok plugin list --json", - "grok plugin uninstall trace-grok --confirm", - ] - ); - - let mut local = FakeRunner::new([serde_json::json!([{ - "name": GROK_PLUGIN, - "source": "/tmp/local-trace-grok", - "status": "installed" - }])]); - disable_grok(&mut local).unwrap(); - assert_eq!(local.calls, ["grok plugin list --json"]); - } - - #[test] - fn opencode_reconciles_the_published_plugin_and_preserves_config() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("opencode.json"); - std::fs::write( - &path, - r#"{"plugin":["other","@braintrust/trace-opencode@0.9.0"],"model":"test/model"}"#, - ) - .unwrap(); - - setup_opencode_at(&path).unwrap(); - - let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); - assert_eq!(config["model"], "test/model"); - assert_eq!( - config["plugin"], - serde_json::json!(["other", opencode_plugin_spec().unwrap()]) - ); - } - - #[test] - fn pi_installs_the_latest_published_extension() { - let mut runner = FakeRunner::new([]); - - setup_pi(&mut runner).unwrap(); - - assert!(runner.called("pi install npm:@braintrust/pi-extension")); - } - - #[test] - fn pi_updates_the_npm_extension_without_installing_it() { - let mut runner = FakeRunner::new([]); - - update_pi(&mut runner).unwrap(); - - assert!(runner.called("pi update npm:@braintrust/pi-extension")); - assert!(!runner - .calls - .iter() - .any(|call| call.starts_with("pi install "))); - } - - #[test] - fn opencode_update_requires_an_existing_plugin_and_preserves_other_config() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("opencode.json"); - std::fs::write( - &path, - r#"{"plugin":["other","@braintrust/trace-opencode@^1"],"model":"test/model"}"#, - ) - .unwrap(); - - update_opencode_at(&path).unwrap(); - - let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); - assert_eq!(config["model"], "test/model"); - assert_eq!( - config["plugin"], - serde_json::json!(["other", opencode_plugin_spec().unwrap()]) - ); - - let missing = temp.path().join("missing.json"); - assert!(update_opencode_at(&missing).is_err()); - assert!(!missing.exists()); - } - - #[test] - fn codex_update_refuses_to_install_a_missing_plugin() { - let mut runner = FakeRunner::new([serde_json::json!({"installed": []})]); - - assert!(update_codex(&mut runner).is_err()); - assert!(!runner - .calls - .iter() - .any(|call| call.starts_with("codex plugin add "))); - } - - #[test] - fn antigravity_installs_published_plugin_and_removes_legacy_registration() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - std::fs::create_dir_all(&config_dir).unwrap(); - let hooks_path = config_dir.join("hooks.json"); - std::fs::write( - &hooks_path, - serde_json::to_vec(&serde_json::json!({ - ANTIGRAVITY_PLUGIN: {"Stop": []}, - "other-plugin": {"Stop": [{"type": "command", "command": "other"}]} - })) - .unwrap(), - ) - .unwrap(); - let mut runner = FakeRunner::new([]); - - setup_antigravity_at(&mut runner, &config_dir).unwrap(); - - assert!(runner.called(&format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"))); - assert!(runner.called(&format!("agy plugin enable {ANTIGRAVITY_PLUGIN}"))); - let hooks: Value = serde_json::from_slice(&std::fs::read(&hooks_path).unwrap()).unwrap(); - assert_eq!(hooks["other-plugin"]["Stop"][0]["command"], "other"); - assert!(hooks.get(ANTIGRAVITY_PLUGIN).is_none()); - } - - #[test] - fn antigravity_update_checks_and_updates_the_same_overridden_home() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - std::fs::create_dir_all(&config_dir).unwrap(); - let expected_home = temp.path().to_path_buf(); - let mut runner = FakeRunner::new([serde_json::json!({ - "imports": [{"name": ANTIGRAVITY_PLUGIN}], - })]); - - update_antigravity_at(&mut runner, &config_dir).unwrap(); - - assert_eq!( - runner.home_calls, - vec![ - ("agy plugin list".into(), expected_home.clone()), - ( - format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"), - expected_home, - ), - ] - ); - } - - #[test] - fn antigravity_setup_is_idempotent() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - std::fs::create_dir_all(&config_dir).unwrap(); - std::fs::write( - config_dir.join("hooks.json"), - serde_json::to_vec(&serde_json::json!({ - ANTIGRAVITY_PLUGIN: {"Stop": []}, - "other-plugin": {"Stop": []} - })) - .unwrap(), - ) - .unwrap(); - let mut runner = FakeRunner::new([]); - - setup_antigravity_at(&mut runner, &config_dir).unwrap(); - let first = std::fs::read(config_dir.join("hooks.json")).unwrap(); - setup_antigravity_at(&mut runner, &config_dir).unwrap(); - let second = std::fs::read(config_dir.join("hooks.json")).unwrap(); - - assert_eq!(first, second); - assert_eq!( - runner - .calls - .iter() - .filter(|call| { - call.as_str() == format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}") - }) - .count(), - 2 - ); - } - - #[test] - fn antigravity_setup_relies_on_native_plugin_hooks() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - let mut runner = FakeRunner::new([]); - - setup_antigravity_at(&mut runner, &config_dir).unwrap(); - - assert!(!config_dir.join("hooks.json").exists()); - assert!(runner.called(&format!("agy plugin install {ANTIGRAVITY_PLUGIN_SOURCE}"))); - assert!(runner.called(&format!("agy plugin enable {ANTIGRAVITY_PLUGIN}"))); - } - - #[test] - fn antigravity_setup_reports_a_missing_cli_without_changing_hooks() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - std::fs::create_dir_all(&config_dir).unwrap(); - let hooks_path = config_dir.join("hooks.json"); - let original = br#"{"other-plugin":{"Stop":[]}}"#; - std::fs::write(&hooks_path, original).unwrap(); - - let error = setup_antigravity_at(&mut MissingAgyRunner, &config_dir).unwrap_err(); - - assert_eq!( - error.to_string(), - "failed to run `agy`; install agy and ensure it is on PATH" - ); - assert_eq!(std::fs::read(hooks_path).unwrap(), original); - } - - #[test] - fn antigravity_disable_removes_only_the_managed_registration() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join(".gemini/config"); - std::fs::create_dir_all(&config_dir).unwrap(); - let hooks_path = config_dir.join("hooks.json"); - std::fs::write( - &hooks_path, - serde_json::to_vec(&serde_json::json!({ - ANTIGRAVITY_PLUGIN: {"Stop": []}, - "other-plugin": {"Stop": [{"type": "command", "command": "other"}]} - })) - .unwrap(), - ) - .unwrap(); - let mut runner = FakeRunner::new([]); - - disable_antigravity_at(&mut runner, &config_dir).unwrap(); - - assert!(runner.called("agy plugin uninstall braintrust-antigravity-tracing")); - let hooks: Value = serde_json::from_slice(&std::fs::read(&hooks_path).unwrap()).unwrap(); - assert!(hooks.get(ANTIGRAVITY_PLUGIN).is_none()); - assert_eq!(hooks["other-plugin"]["Stop"][0]["command"], "other"); - } - - #[test] - fn tracing_settings_preserve_unrelated_fields_and_remove_legacy_keys() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("braintrust.json"); - std::fs::write( - &path, - r#"{ - "traceToBraintrust": false, - "profile": "stale-profile", - "org_name": "stale-org", - "api_key": "stale-secret", - "api_url": "https://stale-api.example", - "app_url": "https://stale-app.example", - "project": "old", - "destination": "old-destination", - "additional_metadata": {"migrated": true}, - "auth": {"type": "legacy"} - }"#, - ) - .unwrap(); - let route = SessionRoute { - auth: AuthSelection { - source: crate::wire::AuthSource::SavedProfile, - profile_id: None, - profile: Some("work".into()), - org_name: Some("Braintrust SDKs".into()), - }, - destination: Some(TraceDestination::ProjectLogs { - project_id: None, - project_name: Some("coding-agents".into()), - }), - ..SessionRoute::default() - }; - - enable_tracing_at(&path, route).unwrap(); - - let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!(settings["trace_to_braintrust"], true); - assert_eq!(settings["route"]["auth"]["profile"], "work"); - assert_eq!( - settings["route"]["destination"]["project_name"], - "coding-agents" - ); - assert_eq!(settings["auth"]["type"], "legacy"); - assert_eq!(settings["route"]["additional_metadata"]["migrated"], true); - for key in [ - "traceToBraintrust", - "profile", - "org_name", - "api_key", - "api_url", - "app_url", - "project", - "destination", - "additional_metadata", - ] { - assert!(settings.get(key).is_none(), "legacy key {key} remained"); - } - } - - #[test] - fn tracing_settings_preserve_metadata_until_setup_explicitly_replaces_it() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("braintrust.json"); - std::fs::write( - &path, - r#"{"route":{"additional_metadata":{"ci":true},"tags":["saved"]}}"#, - ) - .unwrap(); - - let route = SessionRoute::default(); - enable_tracing_at(&path, route).unwrap(); - let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!( - settings["route"]["additional_metadata"], - serde_json::json!({"ci": true}) - ); - assert_eq!(settings["route"]["tags"], serde_json::json!(["saved"])); - - let route = SessionRoute { - additional_metadata: Some(serde_json::json!({"run_id": "new"})), - tags: vec!["replacement".to_string()], - ..SessionRoute::default() - }; - enable_tracing_at(&path, route).unwrap(); - let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!( - settings["route"]["additional_metadata"], - serde_json::json!({"run_id": "new"}) - ); - assert_eq!( - settings["route"]["tags"], - serde_json::json!(["replacement"]) - ); - } - - #[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()); - assert_eq!(hooks["hooks"].as_object().unwrap().len(), 6); - assert!(hooks["hooks"].get("PreToolUse").is_none()); - assert_eq!( - hooks["hooks"]["PreLLMCall"][0]["hooks"][0]["command"], - "bt trace hook --source muse --flush-on-turn-end" - ); - assert!(config_dir.join(".settings.json.lock").exists()); - - 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_managed_run_replaces_only_owned_braintrust_hooks() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - setup_muse_at(&config_dir).unwrap(); - let saved = std::fs::read(muse_hook_path(&config_dir)).unwrap(); - let generated = - muse_managed_hook_config(&config_dir, "'bt' 'hook' '--source' 'muse'").unwrap(); - assert_eq!( - generated["hooks"].as_object().unwrap().len(), - MUSE_HOOK_EVENTS.len() - ); - assert_eq!( - generated["hooks"]["SessionStart"][0]["hooks"][0]["command"], - "'bt' 'hook' '--source' 'muse'" - ); - assert_eq!(std::fs::read(muse_hook_path(&config_dir)).unwrap(), saved); - } - - #[test] - fn muse_managed_run_keeps_foreign_handlers_and_rejects_duplicate_user_hooks() { - 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.json"}"#, - ) - .unwrap(); - let original = r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"true"}]}]}}"#; - std::fs::write(config_dir.join("other.json"), original).unwrap(); - let generated = muse_managed_hook_config(&config_dir, "bt hook --source muse").unwrap(); - assert_eq!(generated["hooks"]["Stop"].as_array().unwrap().len(), 2); - assert_eq!(generated["hooks"]["Stop"][0]["hooks"][0]["command"], "true"); - assert_eq!( - std::fs::read_to_string(config_dir.join("other.json")).unwrap(), - original - ); - - std::fs::write( - config_dir.join("settings.json"), - r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"bt trace hook --source muse"}]}]}}"#, - ) - .unwrap(); - assert!( - muse_managed_hook_config(&config_dir, "bt hook --source muse") - .unwrap_err() - .to_string() - .contains("user hooks already include Braintrust") - ); - - std::fs::write( - config_dir.join("settings.json"), - r#"{"schema_version":1,"managed_hooks_path":"other.json"}"#, - ) - .unwrap(); - std::fs::write( - config_dir.join("other.json"), - r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","name":"foreign","timeout_ms":1000,"command":"true"}]}]}}"#, - ) - .unwrap(); - assert!( - muse_managed_hook_config(&config_dir, "bt hook --source muse") - .unwrap_err() - .to_string() - .contains("strict handler schema") - ); - } - - #[test] - fn muse_managed_run_rejects_duplicate_project_hooks() { - let temp = tempfile::tempdir().unwrap(); - let project = temp.path().join("project"); - let nested = project.join("nested"); - std::fs::create_dir_all(project.join(".muse")).unwrap(); - std::fs::create_dir_all(&nested).unwrap(); - std::fs::write( - project.join(".muse/hooks.json"), - r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"'bt' 'trace' 'hook' '--source' 'muse'"}]}]}}"#, - ) - .unwrap(); - assert!(check_muse_project_hooks(&nested) - .unwrap_err() - .to_string() - .contains("project hooks")); - } - - #[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 muse_setup_and_disable_preserve_foreign_hook_file() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - std::fs::create_dir_all(&config_dir).unwrap(); - let hook_path = muse_hook_path(&config_dir); - let custom = r#"{"schema_version":1,"hooks":{"Stop":[]}}"#; - std::fs::write(&hook_path, custom).unwrap(); - - assert!(setup_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - assert!(disable_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - assert_eq!(std::fs::read_to_string(&hook_path).unwrap(), custom); - } - - #[test] - fn muse_update_requires_an_owned_install_and_upgrades_legacy_command() { - let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - assert!(update_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("not installed")); - - setup_muse_at(&config_dir).unwrap(); - let hook_path = muse_hook_path(&config_dir); - write_object_atomic(&hook_path, legacy_muse_hook_config()).unwrap(); - assert!(muse_doctor_warnings(&config_dir) - .iter() - .any(|warning| warning.contains("older command"))); - update_muse_at(&config_dir).unwrap(); - let installed = load_object(&hook_path).unwrap(); - assert_eq!(installed, muse_hook_config()); - - std::fs::write(&hook_path, r#"{"hooks":{}}"#).unwrap(); - assert!(update_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - } - - #[test] - fn muse_version_requires_a_parseable_native_version() { - assert_eq!( - muse_version(b"Muse Code 1.1.1 (1.1.1-R2514.1)\n").unwrap(), - semver::Version::new(1, 1, 1) - ); - assert!(muse_version(b"unexpected output").is_err()); - } - - #[test] - fn disabling_removes_only_the_braintrust_settings_file() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("braintrust.json"); - std::fs::write( - &path, - r#"{"traceToBraintrust":true,"route":{"destination":{"project_name":"coding-agents"}},"other":true}"#, - ) - .unwrap(); - - remove_tracing_settings(&path).unwrap(); - assert!(!path.exists()); - } - - #[test] - fn disabling_removes_settings_even_when_adapter_cleanup_fails() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("braintrust.json"); - std::fs::write(&path, "{}").unwrap(); - - let error = finish_disable(Err(anyhow::anyhow!("adapter unavailable")), &path) - .unwrap_err() - .to_string(); - - assert_eq!(error, "adapter unavailable"); - assert!(!path.exists()); - } - - #[test] - fn disabling_installed_plugins_uses_each_agents_uninstall_command() { - let mut codex = FakeRunner::new([serde_json::json!({ - "installed": [{"pluginId": CODEX_PLUGIN}] - })]); - disable_codex(&mut codex).unwrap(); - assert!(codex.called("codex plugin remove trace-codex@braintrust-codex-plugins --json")); - - let mut claude = FakeRunner::new([serde_json::json!([{"id": CLAUDE_PLUGIN}])]); - disable_claude(&mut claude).unwrap(); - assert!(claude.called("claude plugin uninstall trace-claude-code@braintrust-claude-plugin")); - - let mut pi = FakeRunner::new([]); - disable_pi(&mut pi).unwrap(); - assert!(pi.called("pi uninstall npm:@braintrust/pi-extension")); - } - - #[test] - fn disabling_opencode_removes_only_the_managed_plugin() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("opencode.json"); - std::fs::write( - &path, - serde_json::to_vec(&serde_json::json!({ - "plugin": ["other", opencode_plugin_spec().unwrap()], - "model": "test/model", - })) - .unwrap(), - ) - .unwrap(); - - remove_opencode_plugin_at(&path).unwrap(); - - let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); - assert_eq!(config["plugin"], serde_json::json!(["other"])); - assert_eq!(config["model"], "test/model"); - } - - #[test] - fn tracing_settings_preserve_plugins_until_setup_explicitly_replaces_them() { - let temp = tempfile::tempdir().unwrap(); - let path = temp.path().join("braintrust.json"); - std::fs::write(&path, r#"{"route":{"span_plugins":["old.mjs"]}}"#).unwrap(); - - enable_tracing_at(&path, SessionRoute::default()).unwrap(); - let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!( - settings["route"]["span_plugins"], - serde_json::json!(["old.mjs"]) - ); - - enable_tracing_at( - &path, - SessionRoute { - span_plugins: vec![PathBuf::from("first.mjs"), PathBuf::from("second.mjs")], - ..SessionRoute::default() - }, - ) - .unwrap(); - let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!( - settings["route"]["span_plugins"], - serde_json::json!(["first.mjs", "second.mjs"]) - ); - } -} diff --git a/bt-daemon/src/setup/common.rs b/bt-daemon/src/setup/common.rs index b50f0461..a943e032 100644 --- a/bt-daemon/src/setup/common.rs +++ b/bt-daemon/src/setup/common.rs @@ -291,7 +291,7 @@ pub(super) fn write_object_atomic_with( crate::settings::with_settings_lock(path, || write_object_atomic_unlocked(path, object, access)) } -fn write_object_atomic_unlocked( +pub(super) fn write_object_atomic_unlocked( path: &Path, object: Map, access: FileAccess, diff --git a/bt-daemon/src/setup/mod.rs b/bt-daemon/src/setup/mod.rs index 61aaa3c0..cb869dfd 100644 --- a/bt-daemon/src/setup/mod.rs +++ b/bt-daemon/src/setup/mod.rs @@ -10,6 +10,7 @@ mod codex; mod common; pub(crate) mod cursor; mod grok; +pub(crate) mod muse; mod opencode; pub(crate) mod pi; diff --git a/bt-daemon/src/setup/muse.rs b/bt-daemon/src/setup/muse.rs new file mode 100644 index 00000000..7d32ecdc --- /dev/null +++ b/bt-daemon/src/setup/muse.rs @@ -0,0 +1,603 @@ +//! Muse Code managed command hooks. +use super::common::{load_object, write_object_atomic_unlocked, CommandRunner, FileAccess}; +use super::Setup; +use crate::agents::Muse; +use crate::paths; +use anyhow::{bail, Context}; +use fs2::FileExt; +use serde_json::{Map, Value}; +use std::path::{Path, PathBuf}; + +const MUSE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreLLMCall", + "PostLLMCall", + "Stop", + "SessionEnd", +]; +const LEGACY_MUSE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreLLMCall", + "PostLLMCall", + "PreCompact", + "PostCompact", + "SubagentStart", + "SubagentStop", + "Stop", + "SessionEnd", +]; + +fn muse_hook_path(config_dir: &Path) -> PathBuf { + config_dir.join("braintrust-hooks.json") +} + +fn muse_hook_config() -> Map { + // Muse treats command-hook failures as non-fatal to the agent turn. + muse_hook_config_for( + MUSE_HOOK_EVENTS, + "bt trace hook --source muse --flush-on-turn-end", + ) +} + +fn muse_hook_config_for(events: &[&str], command: &str) -> Map { + let hooks = 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)), + ]) +} + +/// Build an invocation-local managed file without changing Muse's settings. +/// Preserve handlers administered by someone else, but never run two Muse +/// Braintrust forwarders with different destinations in the same session. +pub(crate) fn muse_managed_hook_config(config_dir: &Path, command: &str) -> anyhow::Result { + let settings = load_object(&config_dir.join("settings.json"))?; + if settings + .get("hooks") + .is_some_and(muse_contains_braintrust_hook) + { + bail!("Muse user hooks already include Braintrust tracing; remove that duplicate before `bt trace run muse`"); + } + let existing = match settings.get("managed_hooks_path").and_then(Value::as_str) { + Some(path) => { + let path = Path::new(path); + let path = if path.is_absolute() { + path.to_path_buf() + } else { + config_dir.join(path) + }; + if !path.exists() || muse_hook_is_owned(&path)? { + Map::from_iter([("schema_version".into(), Value::from(1))]) + } else { + let raw = std::fs::read(&path) + .with_context(|| format!("read Muse managed hooks {}", path.display()))?; + let value: Value = serde_json::from_slice(&raw)?; + if muse_contains_braintrust_hook(&value) { + bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); + } + if muse_has_strict_handler_fields(&value) { + bail!("Muse managed hooks use the strict handler schema; cannot safely merge invocation hooks"); + } + value + .as_object() + .cloned() + .context("Muse managed hooks must be a JSON object")? + } + } + None => Map::from_iter([("schema_version".into(), Value::from(1))]), + }; + let mut result = existing; + let mut hooks = result + .remove("hooks") + .unwrap_or_else(|| Value::Object(Map::new())); + let hooks = hooks + .as_object_mut() + .context("Muse managed hooks must be an object")?; + let generated = muse_hook_config_for(MUSE_HOOK_EVENTS, command); + for (event, handlers) in generated["hooks"] + .as_object() + .context("generated Muse hooks must be an object")? + { + let existing = hooks + .entry(event.clone()) + .or_insert_with(|| Value::Array(Vec::new())); + existing + .as_array_mut() + .with_context(|| format!("Muse managed hook event {event} must be an array"))? + .extend( + handlers + .as_array() + .context("generated Muse handlers must be an array")? + .iter() + .cloned(), + ); + } + result.insert("hooks".into(), hooks.clone().into()); + Ok(Value::Object(result)) +} + +pub(crate) fn check_muse_project_hooks(start: &Path) -> anyhow::Result<()> { + for dir in start.ancestors() { + let path = dir.join(".muse").join("hooks.json"); + if !path.exists() { + continue; + } + let hooks: Value = serde_json::from_slice(&std::fs::read(&path)?) + .with_context(|| format!("read Muse project hooks {}", path.display()))?; + if muse_contains_braintrust_hook(&hooks) { + bail!("Muse project hooks {} already include Braintrust tracing; cannot safely run a managed trace", path.display()); + } + break; + } + Ok(()) +} + +fn muse_contains_braintrust_hook(value: &Value) -> bool { + match value { + Value::Object(object) => object.iter().any(|(key, value)| { + (key == "command" + && value.as_str().is_some_and(|command| { + command.contains("--source") + && command.contains("muse") + && command.contains("hook") + })) + || muse_contains_braintrust_hook(value) + }), + Value::Array(items) => items.iter().any(muse_contains_braintrust_hook), + _ => false, + } +} + +fn muse_has_strict_handler_fields(value: &Value) -> bool { + match value { + Value::Object(object) => { + object.contains_key("name") + || object.contains_key("timeout_ms") + || object.values().any(muse_has_strict_handler_fields) + } + Value::Array(items) => items.iter().any(muse_has_strict_handler_fields), + _ => false, + } +} + +fn legacy_muse_hook_config() -> Map { + muse_hook_config_for( + LEGACY_MUSE_HOOK_EVENTS, + "bt trace hook --source muse --flush-on-turn-end || true", + ) +} + +pub(crate) fn muse_version(output: &[u8]) -> anyhow::Result { + let text = std::str::from_utf8(output).context("Muse version is not UTF-8")?; + text.split_whitespace() + .find_map(|word| semver::Version::parse(word).ok()) + .ok_or_else(|| { + anyhow::anyhow!( + "Muse Code returned an unrecognized version: {}", + text.trim() + ) + }) +} + +/// Muse serializes changes to settings.json with this native lock file. Hold +/// it across the read, hook update, and settings update so we cannot discard a +/// concurrent change made by Muse itself. +fn with_muse_settings_lock( + config_dir: &Path, + action: impl FnOnce() -> anyhow::Result, +) -> anyhow::Result { + std::fs::create_dir_all(config_dir) + .with_context(|| format!("create Muse config directory {}", config_dir.display()))?; + let lock_path = config_dir.join(".settings.json.lock"); + let lock = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&lock_path) + .with_context(|| format!("open Muse settings lock {}", lock_path.display()))?; + lock.lock_exclusive() + .with_context(|| format!("lock Muse settings {}", lock_path.display()))?; + let result = action(); + FileExt::unlock(&lock)?; + result +} + +fn muse_hook_is_owned(path: &Path) -> anyhow::Result { + match std::fs::read(path) { + Ok(raw) => { + let value: Value = serde_json::from_slice(&raw) + .with_context(|| format!("invalid Muse hook configuration: {}", path.display()))?; + Ok(value == Value::Object(muse_hook_config()) + || value == Value::Object(legacy_muse_hook_config())) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => { + Err(error).with_context(|| format!("read Muse hook configuration: {}", path.display())) + } + } +} + +fn check_muse_hook_ownership(path: &Path) -> anyhow::Result<()> { + if path.exists() && !muse_hook_is_owned(path)? { + bail!("Muse hook file {} differs from Braintrust's generated hooks; refusing to overwrite or remove it", path.display()); + } + Ok(()) +} + +fn setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + 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") { + if existing.as_str() != Some(owned_path.as_ref()) { + bail!( + "Muse already uses managed_hooks_path {}; refusing to overwrite it", + existing + ); + } + } + check_muse_hook_ownership(&hook_path)?; + settings.entry("schema_version").or_insert(Value::from(1)); + settings.insert( + "managed_hooks_path".into(), + Value::String(owned_path.into_owned()), + ); + write_muse_object(&hook_path, muse_hook_config())?; + write_muse_object(&settings_path, settings) + }) +} + +fn disable_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + let settings_path = config_dir.join("settings.json"); + let hook_path = muse_hook_path(config_dir); + let mut settings = load_object(&settings_path)?; + check_muse_hook_ownership(&hook_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_muse_object(&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 update_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + let settings = load_object(&config_dir.join("settings.json"))?; + let hook_path = muse_hook_path(config_dir); + if settings.get("managed_hooks_path").and_then(Value::as_str) + != Some(hook_path.to_string_lossy().as_ref()) + { + bail!("Muse tracing hooks are not installed; run `bt trace enable muse`"); + } + if !hook_path.exists() { + bail!("Muse tracing hook file is missing; run `bt trace enable muse`"); + } + check_muse_hook_ownership(&hook_path)?; + write_muse_object(&hook_path, muse_hook_config()) + }) +} + +pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { + let mut warnings = Vec::new(); + match crate::subprocess::background_command("muse") + .arg("--version") + .output() + { + Ok(output) if output.status.success() => match muse_version(&output.stdout) { + Ok(version) if version < semver::Version::new(1, 1, 1) => warnings + .push("Muse Code 1.1.1 or newer is required for verified command hooks".into()), + Ok(_) => {} + Err(error) => warnings.push(error.to_string()), + }, + _ => warnings.push( + "Muse Code executable is unavailable; install Muse Code before enabling tracing".into(), + ), + } + if !crate::subprocess::background_command("bt") + .arg("--version") + .output() + .is_ok_and(|output| output.status.success()) + { + warnings.push( + "Braintrust `bt` executable is unavailable on PATH; Muse command hooks cannot forward events".into(), + ); + } + let settings = match load_object(&config_dir.join("settings.json")) { + Ok(settings) => settings, + Err(error) => { + warnings.push(format!("Muse settings could not be read: {error}")); + return warnings; + } + }; + let hook_path = muse_hook_path(config_dir); + if settings.get("managed_hooks_path").and_then(Value::as_str) + != Some(hook_path.to_string_lossy().as_ref()) + { + warnings.push("Muse managed_hooks_path does not point to Braintrust's hook file; run `bt trace enable muse`".into()); + } + match std::fs::read(&hook_path) { + Ok(raw) => match serde_json::from_slice::(&raw) { + Ok(hooks) if hooks == Value::Object(muse_hook_config()) => {} + Ok(hooks) if hooks == Value::Object(legacy_muse_hook_config()) => warnings.push( + "Muse Braintrust hooks use an older command; run `bt trace update muse`".into(), + ), + Ok(_) => warnings.push("Muse Braintrust hook file differs from expected event definitions; restore it or run `bt trace enable muse`".into()), + Err(error) => warnings.push(format!("Muse Braintrust hook file is invalid JSON: {error}")), + }, + Err(error) => warnings.push(format!("Muse Braintrust hook file could not be read: {error}")), + } + warnings +} + +fn setup_muse() -> anyhow::Result<()> { + let output = crate::subprocess::background_command("muse") + .arg("--version") + .output() + .context("Muse Code is not installed; install Muse Code before enabling tracing")?; + if !output.status.success() { + bail!("Muse Code could not report its version; install Muse Code before enabling tracing"); + } + let installed = muse_version(&output.stdout)?; + if installed < semver::Version::new(1, 1, 1) { + bail!("Muse Code 1.1.1 or newer is required for verified command hooks"); + } + setup_muse_at(&paths::muse_config_dir()) +} +fn disable_muse() -> anyhow::Result<()> { + disable_muse_at(&paths::muse_config_dir()) +} + +fn update_muse() -> anyhow::Result<()> { + update_muse_at(&paths::muse_config_dir()) +} + +fn write_muse_object(path: &Path, object: Map) -> anyhow::Result<()> { + write_object_atomic_unlocked(path, object, FileAccess::Inherited) +} + +impl Setup for Muse { + fn enable(&self, _runner: &mut dyn CommandRunner) -> anyhow::Result<()> { + setup_muse() + } + fn disable(&self, _runner: &mut dyn CommandRunner) -> anyhow::Result<()> { + disable_muse() + } + fn update(&self, _runner: &mut dyn CommandRunner) -> anyhow::Result<()> { + update_muse() + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[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()); + assert_eq!(hooks["hooks"].as_object().unwrap().len(), 6); + assert!(hooks["hooks"].get("PreToolUse").is_none()); + assert_eq!( + hooks["hooks"]["PreLLMCall"][0]["hooks"][0]["command"], + "bt trace hook --source muse --flush-on-turn-end" + ); + assert!(config_dir.join(".settings.json.lock").exists()); + + 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_managed_run_replaces_only_owned_braintrust_hooks() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + setup_muse_at(&config_dir).unwrap(); + let saved = std::fs::read(muse_hook_path(&config_dir)).unwrap(); + let generated = + muse_managed_hook_config(&config_dir, "'bt' 'hook' '--source' 'muse'").unwrap(); + assert_eq!( + generated["hooks"].as_object().unwrap().len(), + MUSE_HOOK_EVENTS.len() + ); + assert_eq!( + generated["hooks"]["SessionStart"][0]["hooks"][0]["command"], + "'bt' 'hook' '--source' 'muse'" + ); + assert_eq!(std::fs::read(muse_hook_path(&config_dir)).unwrap(), saved); + } + + #[test] + fn muse_managed_run_keeps_foreign_handlers_and_rejects_duplicate_user_hooks() { + 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.json"}"#, + ) + .unwrap(); + let original = r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"true"}]}]}}"#; + std::fs::write(config_dir.join("other.json"), original).unwrap(); + let generated = muse_managed_hook_config(&config_dir, "bt hook --source muse").unwrap(); + assert_eq!(generated["hooks"]["Stop"].as_array().unwrap().len(), 2); + assert_eq!(generated["hooks"]["Stop"][0]["hooks"][0]["command"], "true"); + assert_eq!( + std::fs::read_to_string(config_dir.join("other.json")).unwrap(), + original + ); + + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"bt trace hook --source muse"}]}]}}"#, + ) + .unwrap(); + assert!( + muse_managed_hook_config(&config_dir, "bt hook --source muse") + .unwrap_err() + .to_string() + .contains("user hooks already include Braintrust") + ); + + std::fs::write( + config_dir.join("settings.json"), + r#"{"schema_version":1,"managed_hooks_path":"other.json"}"#, + ) + .unwrap(); + std::fs::write( + config_dir.join("other.json"), + r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","name":"foreign","timeout_ms":1000,"command":"true"}]}]}}"#, + ) + .unwrap(); + assert!( + muse_managed_hook_config(&config_dir, "bt hook --source muse") + .unwrap_err() + .to_string() + .contains("strict handler schema") + ); + } + + #[test] + fn muse_managed_run_rejects_duplicate_project_hooks() { + let temp = tempfile::tempdir().unwrap(); + let project = temp.path().join("project"); + let nested = project.join("nested"); + std::fs::create_dir_all(project.join(".muse")).unwrap(); + std::fs::create_dir_all(&nested).unwrap(); + std::fs::write( + project.join(".muse/hooks.json"), + r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"'bt' 'trace' 'hook' '--source' 'muse'"}]}]}}"#, + ) + .unwrap(); + assert!(check_muse_project_hooks(&nested) + .unwrap_err() + .to_string() + .contains("project hooks")); + } + + #[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 muse_setup_and_disable_preserve_foreign_hook_file() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + std::fs::create_dir_all(&config_dir).unwrap(); + let hook_path = muse_hook_path(&config_dir); + let custom = r#"{"schema_version":1,"hooks":{"Stop":[]}}"#; + std::fs::write(&hook_path, custom).unwrap(); + + assert!(setup_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + assert!(disable_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + assert_eq!(std::fs::read_to_string(&hook_path).unwrap(), custom); + } + + #[test] + fn muse_update_requires_an_owned_install_and_upgrades_legacy_command() { + let temp = tempfile::tempdir().unwrap(); + let config_dir = temp.path().join("muse"); + assert!(update_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("not installed")); + + setup_muse_at(&config_dir).unwrap(); + let hook_path = muse_hook_path(&config_dir); + write_muse_object(&hook_path, legacy_muse_hook_config()).unwrap(); + assert!(muse_doctor_warnings(&config_dir) + .iter() + .any(|warning| warning.contains("older command"))); + update_muse_at(&config_dir).unwrap(); + let installed = load_object(&hook_path).unwrap(); + assert_eq!(installed, muse_hook_config()); + + std::fs::write(&hook_path, r#"{"hooks":{}}"#).unwrap(); + assert!(update_muse_at(&config_dir) + .unwrap_err() + .to_string() + .contains("refusing to overwrite")); + } + + #[test] + fn muse_version_requires_a_parseable_native_version() { + assert_eq!( + muse_version(b"Muse Code 1.1.1 (1.1.1-R2514.1)\n").unwrap(), + semver::Version::new(1, 1, 1) + ); + assert!(muse_version(b"unexpected output").is_err()); + } +} diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs index 18399fcf..4837b7d9 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -4,6 +4,7 @@ //! integrations they control. Hosts such as `bt` provide global auth flags and //! dispatch these commands without duplicating agent-specific CLI knowledge. +use crate::agents::{self, Agent, Identity}; use crate::{HookArgs, ImportArgs, RunArgs, ServeArgs, StatusArgs}; use clap::{Args, Subcommand, ValueEnum}; use std::path::PathBuf; @@ -29,7 +30,7 @@ pub enum TraceCommand { /// Run the tracing daemon (foreground). #[command(hide = true)] Daemon(ServeArgs), - /// Forward one coding-agent hook event (read from stdin) to the daemon. + /// Capture one coding-agent hook event from stdin; exits nonzero if capture fails. #[command(hide = true)] Hook(HookArgs), /// Print daemon/session status. @@ -63,6 +64,7 @@ pub struct DoctorArgs { #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum DoctorAgent { Codex, + #[value(alias = "muse-code")] Muse, #[value(name = "claude", alias = "claude-code")] Claude, @@ -71,30 +73,20 @@ pub enum DoctorAgent { Pi, Antigravity, Grok, + Cursor, } impl DoctorAgent { - pub(crate) fn source(self) -> &'static str { + pub(crate) fn identity(self) -> &'static Identity { match self { - Self::Codex => "codex", - Self::Muse => "muse", - Self::Claude => "claude", - Self::OpenCode => "opencode", - Self::Pi => "pi", - Self::Antigravity => "antigravity", - Self::Grok => "grok", - } - } - - 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", - Self::Antigravity => "Google Antigravity", - Self::Grok => "Grok", + Self::Codex => agents::Codex.identity(), + Self::Muse => agents::Muse.identity(), + Self::Claude => agents::Claude.identity(), + Self::OpenCode => agents::OpenCode.identity(), + Self::Pi => agents::Pi.identity(), + Self::Antigravity => agents::Antigravity.identity(), + Self::Grok => agents::Grok.identity(), + Self::Cursor => agents::Cursor.identity(), } } } @@ -139,7 +131,8 @@ pub struct UpdateArgs { pub enum SetupAgent { /// Install the published Codex tracing plugin. Codex, - /// Configure Muse Code tracing hooks. + /// Install the Muse Code tracing hooks. + #[command(alias = "muse-code")] Muse, /// Install the published Claude Code tracing plugin. Claude, @@ -150,10 +143,27 @@ pub enum SetupAgent { Pi, /// Install the published Grok tracing plugin. Grok, + /// Install the published Cursor tracing plugin. + Cursor, /// Install the Google Antigravity tracing hooks. Antigravity, } +impl SetupAgent { + pub(crate) fn identity(self) -> &'static Identity { + match self { + Self::Codex => agents::Codex.identity(), + Self::Muse => agents::Muse.identity(), + Self::Claude => agents::Claude.identity(), + Self::OpenCode => agents::OpenCode.identity(), + Self::Pi => agents::Pi.identity(), + Self::Grok => agents::Grok.identity(), + Self::Cursor => agents::Cursor.identity(), + Self::Antigravity => agents::Antigravity.identity(), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -242,6 +252,16 @@ mod tests { .. }) if value == r#"{"import":true}"# )); + + let cursor_import = Cli::try_parse_from(["bt", "import", "cursor", "session-id"]).unwrap(); + assert!(matches!( + cursor_import.trace.command, + TraceCommand::Import(ImportArgs { + source: crate::ImportSource::Cursor, + ref session_ids, + .. + }) if session_ids == &["session-id"] + )); } #[test] @@ -275,11 +295,27 @@ mod tests { )); } + #[test] + fn run_accepts_cursor_and_forwards_agent_arguments() { + let parsed = Cli::try_parse_from([ + "bt", "run", "cursor", "--tag", "ci", "--", "--print", "hello", + ]) + .unwrap(); + assert!(matches!( + parsed.trace.command, + TraceCommand::Run(RunArgs { + source: crate::RunSource::Cursor, + ref agent_args, + ref tags, + .. + }) if agent_args == &["--print", "hello"] && tags == &["ci"] + )); + } + #[test] 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"), @@ -290,14 +326,15 @@ mod tests { "Google Antigravity", ), ("grok", DoctorAgent::Grok, "grok", "Grok"), + ("cursor", DoctorAgent::Cursor, "cursor", "Cursor"), ] { let cli = Cli::try_parse_from(["bt", "doctor", agent]).unwrap(); assert!(matches!( cli.trace.command, TraceCommand::Doctor(DoctorArgs { agent }) if agent == expected )); - assert_eq!(expected.source(), source); - assert_eq!(expected.display_name(), display_name); + assert_eq!(expected.identity().id, source); + assert_eq!(expected.identity().display_name, display_name); } } @@ -326,28 +363,7 @@ mod tests { } #[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() { + fn grok_uses_shared_enable_and_disable_commands() { for command in ["enable", "setup"] { let parsed = Cli::try_parse_from(["bt", command, "grok"]).unwrap(); assert!(matches!( @@ -366,17 +382,6 @@ mod tests { agent: SetupAgent::Grok }) )); - - let parsed = Cli::try_parse_from(["bt", "doctor", "grok"]).unwrap(); - assert!(matches!( - parsed.trace.command, - TraceCommand::Doctor(DoctorArgs { - agent: DoctorAgent::Grok - }) - )); - - assert_eq!(DoctorAgent::Grok.source(), "grok"); - assert_eq!(DoctorAgent::Grok.display_name(), "Grok"); } #[test] @@ -387,7 +392,7 @@ mod tests { "opencode", "pi", "grok", - "muse", + "cursor", "antigravity", ] { assert!(matches!( @@ -400,18 +405,6 @@ mod tests { } } - #[test] - fn run_accepts_muse() { - let parsed = Cli::try_parse_from(["bt", "run", "muse", "exec", "hello"]).unwrap(); - assert!(matches!( - parsed.trace.command, - TraceCommand::Run(RunArgs { - source: crate::RunSource::Muse, - .. - }) - )); - } - #[test] fn public_commands_preserve_repeated_plugin_order() { for args in [ diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index e637b193..bc616d02 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -10,9 +10,9 @@ use crate::trace_command::{DoctorAgent, DoctorArgs, TraceCommand}; use crate::wire::{AuthSelection, AuthSource, SessionConfig, SessionRoute, TraceDestination}; use crate::{ braintrust_serve_options, paths, run_disable, run_enable, run_import, run_serve, run_status, - run_traced, shutdown_daemon, AuthDiagnostic, AuthLease, AuthProvider, AuthResolveReason, - BraintrustSinkConfig, DaemonDiagnostic, DaemonStatus, DoctorCommandOutput, HostInfo, - OutputFormat, Registry, RunHookCommand, ServeOptions, StatusArgs, TraceArgs, + run_traced_with_host, shutdown_daemon, AuthDiagnostic, AuthLease, AuthProvider, + AuthResolveReason, BraintrustSinkConfig, DaemonDiagnostic, DaemonStatus, DoctorCommandOutput, + HostInfo, OutputFormat, Registry, RunHookCommand, ServeOptions, StatusArgs, TraceArgs, TraceCommandOutput, }; use async_trait::async_trait; @@ -533,7 +533,9 @@ async fn doctor_output_at( warnings.push(warning); } if args.agent == DoctorAgent::Muse { - warnings.extend(crate::setup::muse_doctor_warnings(&paths::muse_config_dir())); + warnings.extend(crate::setup::muse::muse_doctor_warnings( + &paths::muse_config_dir(), + )); } let (route, route_source) = match settings.route { @@ -678,6 +680,14 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul if crate::hook::suppress_hook(&hook_args) { return Ok(()); } + if hook_args.managed_context.is_some() { + return crate::hook::run_hook_with_route( + hook_args, + SessionRoute::default(), + host_info(&host), + ) + .await; + } let settings = crate::settings::AgentSettings::load_for_hook(&hook_args.source)?; run_hook_command(hook_args, &host, settings).await } @@ -736,7 +746,8 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul apply_additional_metadata(&mut route, run_args.additional_metadata.as_deref())?; apply_tags(&mut route, &run_args.tags)?; let hook_command = child_command(&host.command, "hook"); - let status = run_traced(run_args, hook_command, route).await?; + let status = + run_traced_with_host(run_args, hook_command, route, host_info(&host)).await?; if status.success() { Ok(()) } else { @@ -1092,6 +1103,7 @@ mod tests { flush_timeout_ms: 10_000, additional_metadata: None, managed_run_hook: false, + managed_context: None, } } diff --git a/bt-daemon/src/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs index 08577b67..6d7bf849 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -1,3 +1,4 @@ +use crate::agents::{registrar, Agent}; use crate::wire::Envelope; use crate::ImportSource; use anyhow::{bail, Context}; @@ -9,98 +10,172 @@ use std::path::{Path, PathBuf}; mod antigravity; mod claude; mod codex; +mod cursor; pub(crate) mod muse; +mod pipeline; + +pub use pipeline::{import_transcript, import_transcripts, run_import}; + +/// Importing one agent's native transcripts. Implemented in each agent's +/// module here and registered in [`crate::agents`]. +pub(crate) trait TranscriptImport: Agent { + /// Every session transcript the agent has saved under `home`. + fn discover(&self, home: &Path) -> anyhow::Result>; + /// The saved transcript for one session under `home`. + fn find(&self, home: &Path, session_id: &str) -> anyhow::Result; + /// State for following one growing transcript. + fn tail(&self) -> Box; +} + +/// Where an agent saves its transcripts and how they are named. Each agent +/// builds one privately; this holds the search shared by all of them. +pub(super) struct TranscriptLayout { + pub display_name: &'static str, + /// Directories searched recursively for `.jsonl` transcripts. + pub roots: Vec, + /// The session a transcript belongs to, if the path is one. + pub session_id: fn(&Path) -> Option, + /// Whether a path is named for a session id. + pub filename_matches: fn(&Path, &str) -> bool, +} + +/// Follows one growing transcript, turning the records read so far into +/// synthetic hook events and emitting only those not emitted before. +pub(crate) trait TailSession: Send { + /// Synthetic hook events for every record read so far. + fn envelopes( + &mut self, + path: &Path, + records: &IncrementalRecords, + ) -> anyhow::Result>; + + /// The events from `events` to emit now. `finalize` ends the session. + fn poll( + &mut self, + events: Vec, + len: u64, + finalize: bool, + ) -> anyhow::Result>; + + /// Whether a rewritten transcript continues with the same translator + /// instead of starting a new one. If so, the fresh tail that replaces this + /// one receives [`Self::rewritten`]. + fn keeps_translator_on_rewrite(&self) -> bool { + false + } + + /// Called on a fresh tail after a rewrite, with the records from before + /// and after it, when [`Self::keeps_translator_on_rewrite`] is true. + fn rewritten(&mut self, _previous: &[Value], _current: &[Value]) {} + + /// Whether an attached import may stop at a final record that is still + /// being written when it shuts down. + fn tolerates_incomplete_final_record(&self) -> bool { + false + } +} + +fn importer(source: ImportSource) -> &'static dyn TranscriptImport { + let name = source.identity().id; + registrar() + .import + .get(name) + .unwrap_or_else(|| panic!("no transcript import is registered for {name}")) +} pub(crate) fn resolve_transcripts( session_ids: &[String], all: bool, source: ImportSource, ) -> anyhow::Result> { + let home = std::env::var_os("HOME") + .or_else(|| std::env::var_os("USERPROFILE")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(".")); + let importer = importer(source); match (all, session_ids.is_empty()) { - (true, true) => discover_transcripts(source), + (true, true) => importer.discover(&home), (false, false) => session_ids .iter() - .map(|session_id| resolve_transcript(session_id, source)) + .map(|session_id| importer.find(&home, session_id)) .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::Muse => Vec::new(), - ImportSource::Claude => claude::roots(&home), - ImportSource::Antigravity => antigravity::roots(&home), +impl TranscriptLayout { + /// One transcript per session found under the roots, in path order. + pub fn discover(&self) -> anyhow::Result> { + let mut sessions = BTreeMap::>::new(); + for path in self.candidates() { + if let Some(session_id) = (self.session_id)(&path) { + sessions.entry(session_id).or_default().push(path); + } + } + if sessions.is_empty() { + bail!( + "no {} transcripts found; searched {}", + self.display_name, + self.locations() + ); + } + sessions + .into_iter() + .map(|(session_id, paths)| self.single(&session_id, paths)) + .collect() } -} -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); + /// The one transcript named for `session_id`. + pub fn find(&self, session_id: &str) -> anyhow::Result { + validate_session_id(session_id)?; + let matches: Vec<_> = self + .candidates() + .into_iter() + .filter(|path| (self.filename_matches)(path, session_id)) + .collect(); + if matches.is_empty() { + bail!( + "no {} transcript found for session {session_id}; searched {}", + self.display_name, + self.locations() + ); + } + self.single(session_id, matches) } - 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) - ); + fn candidates(&self) -> Vec { + let mut candidates = Vec::new(); + for root in &self.roots { + find_jsonl_files(root, &mut candidates); + } + candidates.sort(); + candidates.dedup(); + candidates } - 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 + fn single(&self, session_id: &str, paths: Vec) -> anyhow::Result { + match <[PathBuf; 1]>::try_from(paths) { + Ok([path]) => Ok(path), + Err(paths) => bail!( + "multiple {} transcripts found for session {session_id}: {}", + self.display_name, + paths .iter() .map(|path| path.display().to_string()) .collect::>() - .join(", "); - bail!( - "multiple {} transcripts found for session {session_id}: {locations}", - source_name(source) - ); - } + .join(", ") + ), } } - Ok(resolved) + + fn locations(&self) -> String { + self.roots + .iter() + .map(|root| root.display().to_string()) + .collect::>() + .join(", ") + } } fn find_jsonl_files(directory: &Path, matches: &mut Vec) { @@ -123,62 +198,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), - } -} - -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::Muse => false, - ImportSource::Claude => claude::filename_matches(path, session_id), - ImportSource::Antigravity => antigravity::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) - ) - } - } -} - -pub(crate) fn validate_session_id(session_id: &str) -> anyhow::Result<()> { +fn validate_session_id(session_id: &str) -> anyhow::Result<()> { if session_id.is_empty() || !session_id .bytes() @@ -189,47 +209,20 @@ pub(crate) fn validate_session_id(session_id: &str) -> anyhow::Result<()> { Ok(()) } -fn source_name(source: ImportSource) -> &'static str { - match source { - ImportSource::Codex => "Codex", - ImportSource::Muse => "Muse Code", - ImportSource::Claude => "Claude Code", - ImportSource::Antigravity => "Google Antigravity", - } -} - #[cfg(test)] pub(crate) fn transcript_envelopes( path: &Path, source: ImportSource, ) -> anyhow::Result> { let mut records = IncrementalRecords::default(); - records.refresh(path, true)?; - envelopes_from_records(path, source, &records) -} - -fn envelopes_from_records( - path: &Path, - source: ImportSource, - records: &IncrementalRecords, -) -> anyhow::Result> { - match source { - ImportSource::Codex => codex::envelopes(path, &records.values), - ImportSource::Muse => muse::envelopes(path), - ImportSource::Claude => claude::envelopes( - path, - &records.values, - &records.end_offsets, - records.read_offset, - ), - ImportSource::Antigravity => { - antigravity::envelopes(path, &records.values, &records.end_offsets) - } - } + records.refresh_with_final_partial(path, true, false)?; + importer(source).tail().envelopes(path, &records) } +/// The records of a transcript read so far, re-read incrementally as it +/// grows. #[derive(Default)] -struct IncrementalRecords { +pub(crate) struct IncrementalRecords { values: Vec, end_offsets: Vec, read_offset: u64, @@ -246,7 +239,12 @@ enum Refresh { } impl IncrementalRecords { - fn refresh(&mut self, path: &Path, finalize: bool) -> anyhow::Result { + fn refresh_with_final_partial( + &mut self, + path: &Path, + finalize: bool, + allow_incomplete_final_record: bool, + ) -> anyhow::Result { let metadata = std::fs::metadata(path) .with_context(|| format!("read transcript metadata {}", path.display()))?; let len = metadata.len(); @@ -298,7 +296,13 @@ impl IncrementalRecords { values.push(value); end_offsets.push(offset); } - Err(_) if !finalize && offset == len && !line.ends_with('\n') => break, + Err(_) + if (allow_incomplete_final_record || !finalize) + && offset == len + && !line.ends_with('\n') => + { + break + } Err(error) => { return Err(error).with_context(|| { format!("parse transcript {} line {}", path.display(), line_count) @@ -350,43 +354,44 @@ fn read_anchor(path: &Path, through: u64, max_len: usize) -> anyhow::Result, records: IncrementalRecords, retry_envelopes: bool, translator_reset: bool, -} - -enum TailState { - Codex(codex::Tail), - Muse, - Claude(claude::Tail), - Antigravity(antigravity::Tail), + allow_incomplete_final_record: bool, } impl TranscriptTail { pub(crate) fn new(path: PathBuf, source: ImportSource) -> Self { + let importer = importer(source); Self { path, - source, - state: Self::new_state(source), + importer, + session: importer.tail(), records: IncrementalRecords::default(), retry_envelopes: false, translator_reset: false, + allow_incomplete_final_record: false, } } - 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()), - } + pub(crate) fn allow_incomplete_final_record_on_shutdown(&mut self) { + self.allow_incomplete_final_record = true; } pub(crate) fn poll(&mut self, finalize: bool) -> anyhow::Result> { - let refresh = match self.records.refresh(&self.path, finalize) { + let keeps_translator = self.session.keeps_translator_on_rewrite(); + let previous_records = if keeps_translator { + self.records.values.clone() + } else { + Vec::new() + }; + let refresh = match self.records.refresh_with_final_partial( + &self.path, + finalize, + self.session.tolerates_incomplete_final_record() && self.allow_incomplete_final_record, + ) { Ok(refresh) => refresh, Err(_) if !finalize => return Ok(Vec::new()), Err(error) => return Err(error), @@ -395,10 +400,15 @@ impl TranscriptTail { return Ok(Vec::new()); } if refresh == Refresh::Reset { - self.state = Self::new_state(self.source); - self.translator_reset = true; + self.session = self.importer.tail(); + if keeps_translator { + self.session + .rewritten(&previous_records, &self.records.values); + } else { + self.translator_reset = true; + } } - let events = match envelopes_from_records(&self.path, self.source, &self.records) { + let events = match self.session.envelopes(&self.path, &self.records) { Ok(events) => events, Err(_) if !finalize => { self.retry_envelopes = true; @@ -410,12 +420,7 @@ impl TranscriptTail { let len = std::fs::metadata(&self.path) .with_context(|| format!("read transcript metadata {}", self.path.display()))? .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), - } + self.session.poll(events, len, finalize) } pub(crate) fn take_translator_reset(&mut self) -> bool { @@ -526,7 +531,7 @@ mod tests { std::fs::write(&transcript, "{}\n").unwrap(); assert_eq!( - resolve_transcript_in("session-123", ImportSource::Codex, &[root]).unwrap(), + codex::layout(vec![root]).find("session-123").unwrap(), transcript ); } @@ -542,7 +547,7 @@ mod tests { std::fs::write(&transcript, "{}\n").unwrap(); assert_eq!( - resolve_transcript_in("session-123", ImportSource::Claude, &[root]).unwrap(), + claude::layout(vec![root]).find("session-123").unwrap(), transcript ); } @@ -563,19 +568,252 @@ mod tests { std::fs::write(unrelated, "{}\n").unwrap(); assert_eq!( - resolve_transcript_in("conversation-123", ImportSource::Antigravity, &[root]).unwrap(), + antigravity::layout(vec![root]) + .find("conversation-123") + .unwrap(), transcript ); } #[test] - fn rejects_unsafe_session_ids() { - let error = resolve_transcript_in( - "../session", - ImportSource::Codex, - &[PathBuf::from("unused")], + fn finds_cursor_transcripts_by_nested_session_identity_and_ignores_symlinks() { + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("projects"); + let transcript = root + .join("workspace") + .join("agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + std::fs::write(&transcript, r#"{"role":"user","message":{"content":"hi"}}"#).unwrap(); + let subagent = root + .join("workspace") + .join("agent-transcripts/agent-session-123/agent-session-123.jsonl"); + std::fs::create_dir_all(subagent.parent().unwrap()).unwrap(); + std::fs::write(&subagent, "{}\n").unwrap(); + + assert_eq!( + cursor::layout(vec![root.clone()]) + .find("session-123") + .unwrap(), + transcript + ); + assert_eq!( + cursor::layout(vec![root]).discover().unwrap(), + vec![transcript] + ); + } + + #[test] + fn cursor_import_emits_only_supported_transcript_enrichment_and_boundaries() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp + .path() + .join("workspace/agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + let contents = concat!( + "{\"role\":\"user\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"question\"}]}}\n", + "{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"tool_use\",\"name\":\"Read\",\"input\":{\"path\":\"a.txt\"}}]}}\n", + "{\"role\":\"assistant\",\"message\":{\"content\":[{\"type\":\"text\",\"text\":\"answer\"}]}}\n", + "{\"type\":\"turn_ended\",\"status\":\"success\"}\n" + ); + std::fs::write(&transcript, contents).unwrap(); + let events = transcript_envelopes(&transcript, ImportSource::Cursor).unwrap(); + + assert_eq!( + events + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["_bt_importStart", "_bt_importCheckpoint", "_bt_importStop"] + ); + assert!(events.iter().all(|event| event.source == "cursor")); + assert_eq!( + events[1].payload["_bt_transcript_mirror"]["through"], + contents.len() as u64 + ); + assert_eq!( + events[1].payload["historical_fidelity"], + "user_and_assistant_text_and_terminal_status" + ); + assert_eq!(events[2].payload["status"], "success"); + assert_eq!(events[1].payload["start_time_estimated"], true); + } + + #[test] + fn cursor_import_preserves_terminal_error_message() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp + .path() + .join("workspace/agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + std::fs::write( + &transcript, + "{\"role\":\"user\",\"message\":{\"content\":\"question\"}}\n{\"type\":\"turn_ended\",\"status\":\"error\",\"error\":\"WritableIterable is closed\"}\n", + ) + .unwrap(); + + let events = transcript_envelopes(&transcript, ImportSource::Cursor).unwrap(); + let stop = events + .iter() + .find(|event| event.event == "_bt_importStop") + .unwrap(); + assert_eq!(stop.payload["status"], "error"); + assert_eq!(stop.payload["error_message"], "WritableIterable is closed"); + } + + #[test] + fn cursor_attach_tails_growth_without_repeating_boundaries_or_checkpoints() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp + .path() + .join("workspace/agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + let first_record = b"{\"role\":\"user\",\"message\":{\"content\":\"one\"}}\n"; + std::fs::write(&transcript, first_record).unwrap(); + let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); + let first = tail.poll(false).unwrap(); + assert_eq!( + first + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["_bt_importStart", "_bt_importCheckpoint"] + ); + let snapshot = first[1].payload["_bt_transcript_mirror"]["mirror"] + .as_str() + .unwrap() + .to_owned(); + assert_ne!(PathBuf::from(&snapshot), transcript); + assert_eq!( + first[1].payload["_bt_transcript_mirror"]["through"], + first_record.len() as u64 + ); + assert!(tail.poll(false).unwrap().is_empty()); + + use std::io::Write; + std::fs::OpenOptions::new() + .append(true) + .open(&transcript) + .unwrap() + .write_all(b"{\"role\":\"assistant\",\"message\":{\"content\":\"two\"}}\n") + .unwrap(); + let grown = tail.poll(false).unwrap(); + assert_eq!( + grown + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["_bt_importCheckpoint"] + ); + assert_eq!( + grown[0].payload["_bt_transcript_mirror"]["mirror"], + snapshot + ); + assert_eq!( + std::fs::read_to_string(&snapshot).unwrap().lines().count(), + 2 + ); + assert_eq!( + tail.poll(true) + .unwrap() + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["_bt_importStop"] + ); + std::fs::write(&transcript, b"replacement transcript\n").unwrap(); + let snapshot_contents = std::fs::read_to_string(&snapshot).unwrap(); + assert!(snapshot_contents.contains("\"content\":\"one\"")); + assert!(snapshot_contents.contains("\"content\":\"two\"")); + } + + #[test] + fn cursor_transcript_replacement_rotates_snapshot_without_resetting_translator() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp + .path() + .join("workspace/agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + std::fs::write( + &transcript, + "{\"role\":\"user\",\"message\":{\"content\":\"the original longer prompt\"}}\n", + ) + .unwrap(); + let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); + let first = tail.poll(false).unwrap(); + let first_snapshot = first[1].payload["_bt_transcript_mirror"]["mirror"] + .as_str() + .unwrap() + .to_owned(); + + std::fs::write( + &transcript, + "{\"role\":\"user\",\"message\":{\"content\":\"latest\"}}\n", ) - .unwrap_err(); + .unwrap(); + let replacement = tail.poll(false).unwrap(); + assert_eq!(replacement[0].event, "_bt_importStart"); + assert_eq!(replacement[1].event, "_bt_importCheckpoint"); + let replacement_snapshot = replacement[1].payload["_bt_transcript_mirror"]["mirror"] + .as_str() + .unwrap(); + assert_ne!(first_snapshot, replacement_snapshot); + assert!(!tail.take_translator_reset()); + } + + #[test] + fn cursor_attach_finalizes_after_an_incomplete_trailing_record() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp + .path() + .join("workspace/agent-transcripts/session-123/session-123.jsonl"); + std::fs::create_dir_all(transcript.parent().unwrap()).unwrap(); + std::fs::write( + &transcript, + "{\"role\":\"user\",\"message\":{\"content\":\"complete\"}}\n", + ) + .unwrap(); + + let mut tail = TranscriptTail::new(transcript.clone(), ImportSource::Cursor); + tail.allow_incomplete_final_record_on_shutdown(); + assert_eq!( + tail.poll(false).unwrap().last().unwrap().event, + "_bt_importCheckpoint" + ); + use std::io::Write; + std::fs::OpenOptions::new() + .append(true) + .open(&transcript) + .unwrap() + .write_all(b"{\"role\":\"assistant\",\"message\":") + .unwrap(); + + let final_events = tail.poll(true).unwrap(); + assert_eq!( + final_events + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + vec!["_bt_importCheckpoint", "_bt_importStop"] + ); + } + + #[test] + fn completed_import_still_rejects_an_incomplete_trailing_record() { + let temp = tempfile::tempdir().unwrap(); + let transcript = temp.path().join("transcript.jsonl"); + std::fs::write(&transcript, "{\"type\":\"complete\"}\n{\"partial\"").unwrap(); + + let error = IncrementalRecords::default() + .refresh_with_final_partial(&transcript, true, false) + .unwrap_err(); + assert!(error.to_string().contains("parse transcript")); + } + + #[test] + fn rejects_unsafe_session_ids() { + let error = codex::layout(vec![PathBuf::from("unused")]) + .find("../session") + .unwrap_err(); assert!(error.to_string().contains("invalid session id")); } @@ -587,18 +825,16 @@ mod tests { 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(); + let missing = claude::layout(vec![first.clone(), second.clone()]) + .find("missing") + .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(); + let ambiguous = claude::layout(vec![first, second]) + .find("duplicate") + .unwrap_err(); assert!(ambiguous .to_string() .contains("multiple Claude Code transcripts")); @@ -631,7 +867,7 @@ mod tests { std::fs::write(root.join("not-a-transcript.jsonl"), "{}\n").unwrap(); assert_eq!( - discover_transcripts_in(ImportSource::Codex, &[root]).unwrap(), + codex::layout(vec![root]).discover().unwrap(), vec![first, second] ); } @@ -657,13 +893,15 @@ mod tests { std::fs::write(project.join("notes.jsonl"), "{}\n").unwrap(); assert_eq!( - discover_transcripts_in(ImportSource::Claude, &[root]).unwrap(), + claude::layout(vec![root]).discover().unwrap(), vec![transcript] ); } #[test] - fn claude_import_emits_related_subagent_lifecycle_inside_parent_turn() { + fn claude_import_keeps_delayed_child_and_native_handback_under_launching_turn() { + use crate::translate::{Registry, SessionCtx, SpanOp, SpanType}; + 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"); @@ -671,8 +909,12 @@ mod tests { 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:03Z","message":{"content":"unrelated question"}}), 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"}]}}), + json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:07Z","origin":{"kind":"peer","handback":true,"senderTaskId":"child-a","from":"child-a","name":"reviewer"},"message":{"content":"Another Claude session sent a message:\nreviewed\nThis agent is working inside the same session."}}), + json!({"type":"assistant","sessionId":"session-a","timestamp":"2026-01-01T00:00:08Z","message":{"id":"handback-response","content":[{"type":"text","text":"accepted review"}]}}), + json!({"type":"user","sessionId":"session-a","timestamp":"2026-01-01T00:00:09Z","message":{"content":"next human question"}}), ]; std::fs::write( &transcript, @@ -690,27 +932,72 @@ mod tests { .unwrap(); let events = transcript_envelopes(&transcript, ImportSource::Claude).unwrap(); - let names = events + let context = SessionCtx { + session_id: "session-a".into(), + config: None, + }; + let mut translator = Registry::default_agents().create("claude-code", "session-a"); + let mut ops = Vec::new(); + for event in &events { + ops.extend(translator.handle(event, &context).unwrap()); + while let Some(pending) = translator.drain_pending(&context).unwrap() { + ops.extend(pending); + } + } + ops.extend(translator.finalize(&context).unwrap()); + while let Some(pending) = translator.drain_pending(&context).unwrap() { + ops.extend(pending); + } + let inserts = ops + .iter() + .filter_map(|op| match op { + SpanOp::Insert(row) => Some(row), + SpanOp::Merge(_) => None, + }) + .collect::>(); + let turns = inserts .iter() - .map(|event| event.event.as_str()) + .filter(|row| row.name.starts_with("Turn ")) .collect::>(); assert_eq!( - names, - vec![ - "SessionStart", - "UserPromptSubmit", - "SubagentStart", - "SubagentStop", - "Stop", - "SessionEnd" - ] + turns + .iter() + .map(|row| row.name.as_str()) + .collect::>(), + ["Turn 1", "Turn 2", "Turn 3"] ); - assert_eq!(events[2].payload["agent_id"], json!("child-a")); - assert_eq!(events[2].payload["agent_type"], json!("reviewer")); + let launch = turns[0]; + let child = inserts + .iter() + .find(|row| { + row.span_type == SpanType::Task + && row + .metadata + .as_ref() + .and_then(|value| value.get("agent_id")) + == Some(&json!("child-a")) + }) + .unwrap(); + assert_eq!(child.parent_span_ids, std::slice::from_ref(&launch.span_id)); + let continuation = inserts + .iter() + .find(|row| { + row.metadata + .as_ref() + .and_then(|value| value.get("turn_trigger")) + == Some(&json!("agent_message")) + }) + .unwrap(); assert_eq!( - Path::new(events[3].payload["agent_transcript_path"].as_str().unwrap()), - subagent + continuation.parent_span_ids, + std::slice::from_ref(&launch.span_id) ); + assert!(ops.iter().any(|op| match op { + SpanOp::Insert(row) | SpanOp::Merge(row) => { + row.span_id == continuation.span_id + && row.output.as_ref() == Some(&json!("accepted review")) + } + })); } #[test] @@ -994,7 +1281,12 @@ mod tests { let path = temp.path().join("session.jsonl"); std::fs::write(&path, "{\"type\":\"session_meta\"").unwrap(); let mut records = IncrementalRecords::default(); - assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Unchanged); + assert_eq!( + records + .refresh_with_final_partial(&path, false, false) + .unwrap(), + Refresh::Unchanged + ); assert!(records.values.is_empty()); std::fs::write( @@ -1002,11 +1294,21 @@ mod tests { "{\"type\":\"session_meta\",\"payload\":{\"id\":\"session\"}}\n", ) .unwrap(); - assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Appended); + assert_eq!( + records + .refresh_with_final_partial(&path, false, false) + .unwrap(), + Refresh::Appended + ); assert_eq!(records.values.len(), 1); std::fs::write(&path, "{\"type\":\"event_msg\"}\n").unwrap(); - assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Reset); + assert_eq!( + records + .refresh_with_final_partial(&path, false, false) + .unwrap(), + Refresh::Reset + ); assert_eq!(records.values, vec![json!({"type":"event_msg"})]); std::fs::write( @@ -1014,7 +1316,12 @@ mod tests { "{\"type\":\"replacement_with_a_longer_prefix\"}\n{\"type\":\"second\"}\n", ) .unwrap(); - assert_eq!(records.refresh(&path, false).unwrap(), Refresh::Reset); + assert_eq!( + records + .refresh_with_final_partial(&path, false, false) + .unwrap(), + Refresh::Reset + ); assert_eq!( records.values, vec![ diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index d0f0476e..5c91246d 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -257,7 +257,7 @@ where .event .as_ref() .ok_or_else(|| anyhow!("Muse run record has no event"))?; - let kind = kind.expect("recognized run event has kind"); + let kind = kind.context("recognized Muse run event has no kind")?; let run_id = payload .run_id .as_deref() @@ -349,9 +349,15 @@ where { self.start_model(&run_id, ts, None, Value::Null)?; } - let run = self.runs.get_mut(&run_id).expect("model start creates run"); + let run = self + .runs + .get_mut(&run_id) + .context("Muse model start did not create a run")?; run.last_assistant = Some(text.clone()); - run.active_model.as_mut().expect("model exists").text = Some(text); + run.active_model + .as_mut() + .context("Muse run has no active model")? + .text = Some(text); } "terminal" => { let terminal = event @@ -410,7 +416,10 @@ where } for run_id in self.runs.keys().cloned().collect::>() { self.finish_model(&run_id, end_ts, Some("Run ended without a terminal event"))?; - let run = self.runs.remove(&run_id).expect("run exists"); + let run = self + .runs + .remove(&run_id) + .context("Muse run disappeared while finalizing")?; self.emit( "Stop", end_ts, @@ -441,6 +450,7 @@ where } } +#[cfg(test)] pub(crate) fn envelopes(path: &Path) -> Result> { let header = read_header(path)?; let session = single_session(&header, path)?; @@ -452,6 +462,7 @@ pub(crate) fn session_id(path: &Path) -> Result { Ok(single_session(&header, path)?.session_id.clone()) } +#[cfg(test)] pub(crate) fn envelopes_for_session(path: &Path, expected_id: &str) -> Result> { let mut events = Vec::new(); for_each_envelope_for_session(path, expected_id, |event| { @@ -461,6 +472,7 @@ pub(crate) fn envelopes_for_session(path: &Path, expected_id: &str) -> Result( path: &Path, expected_id: &str, diff --git a/bt-daemon/src/transcript_import/muse_runtime.rs b/bt-daemon/src/transcript_import/muse_runtime.rs new file mode 100644 index 00000000..f460a544 --- /dev/null +++ b/bt-daemon/src/transcript_import/muse_runtime.rs @@ -0,0 +1,563 @@ +//! Muse exports are streamed through the shared import processor. +use super::*; +use anyhow::Context; +use serde::Deserialize; +use std::ffi::OsString; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::{ + AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines, +}; + +const MUSE_MSP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const MUSE_MSP_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const MUSE_MSP_MAX_PAGES: usize = 1_000; +const MUSE_MSP_DIAGNOSTIC_BYTES: usize = 64 * 1024; + +pub(super) async fn run( + args: ImportArgs, + opts: ServeOptions, + config: Option, + ledger: Option, +) -> anyhow::Result> { + if args.attach { + return attach_muse_session(&args.session_ids[0], opts, config, ledger).await; + } + let ids = if args.all { + muse_msp_session_ids().await? + } else { + args.session_ids + }; + if ids.is_empty() { + anyhow::bail!("no durable Muse sessions found to import"); + } + import_muse_sessions_with_ledger(&ids, opts, config, ledger).await +} + +async fn export_muse_session(session_id: &str, path: &std::path::Path) -> anyhow::Result<()> { + transcript_import::validate_session_id(session_id)?; + let executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| OsString::from("muse")); + let mut command: tokio::process::Command = + crate::subprocess::background_command(&executable).into(); + let mut child = command + .kill_on_drop(true) + .args(["export", "--session", session_id, "--out"]) + .arg(path) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| { + format!( + "run {} export for Muse session {session_id}", + executable.to_string_lossy() + ) + })?; + let stderr = child + .stderr + .take() + .context("Muse export stderr was not piped")?; + let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); + let status = tokio::time::timeout(Duration::from_secs(300), child.wait()) + .await + .with_context(|| format!("Muse export for session {session_id} timed out"))? + .with_context(|| format!("wait for Muse export of session {session_id}"))?; + let stderr = diagnostics + .await + .context("collect Muse export diagnostics")??; + if !status.success() { + anyhow::bail!( + "Muse export for session {session_id} failed: {}", + stderr.trim() + ); + } + Ok(()) +} + +/// 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> { + let executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| OsString::from("muse")); + let mut command: tokio::process::Command = + crate::subprocess::background_command(&executable).into(); + let mut child = command + .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().context("Muse MSP stdin was not piped")?; + let stdout = child + .stdout + .take() + .context("Muse MSP stdout was not piped")?; + let stderr = child + .stderr + .take() + .context("Muse MSP stderr was not piped")?; + let diagnostics = tokio::spawn(muse_drain_stderr(stderr)); + let mut reader = BufReader::new(stdout).lines(); + let listing = async { + 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 requires the initialized notification before other requests. + muse_msp_notify(&mut stdin, "initialized").await?; + let mut cursor = None; + let mut seen_cursors = std::collections::HashSet::new(); + let mut sessions = Vec::new(); + for page_index in 0..MUSE_MSP_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?; + match muse_msp_append_page(result, page_index + 1, &mut sessions, &mut seen_cursors)? { + None => { + sessions.sort(); + sessions.dedup(); + return Ok(sessions); + } + Some(next) => cursor = Some(next), + } + } + anyhow::bail!("Muse session/list exceeded {MUSE_MSP_MAX_PAGES} pages") + } + .await; + drop(stdin); + drop(reader); + let status = match tokio::time::timeout(MUSE_MSP_SHUTDOWN_TIMEOUT, child.wait()).await { + Ok(status) => status.context("wait for Muse MSP host")?, + Err(_) => { + child + .start_kill() + .context("stop unresponsive Muse MSP host")?; + child + .wait() + .await + .context("reap unresponsive Muse MSP host")?; + anyhow::bail!("Muse MSP host did not exit after stdin closed") + } + }; + let stderr = diagnostics + .await + .context("collect Muse MSP diagnostics")??; + let sessions = listing?; + if !status.success() { + anyhow::bail!("Muse MSP session listing failed: {}", stderr.trim()); + } + Ok(sessions) +} + +#[derive(Deserialize)] +struct MuseMspSessionPage { + sessions: Vec, + #[serde(rename = "nextCursor")] + next_cursor: Option, +} + +#[derive(Deserialize)] +struct MuseMspSessionEntry { + #[serde(rename = "sessionId", alias = "session_id")] + session_id: String, +} + +fn muse_msp_append_page( + result: serde_json::Value, + page_index: usize, + sessions: &mut Vec, + seen_cursors: &mut std::collections::HashSet, +) -> anyhow::Result> { + let page: MuseMspSessionPage = serde_json::from_value(result) + .with_context(|| format!("decode Muse session/list page {page_index}"))?; + for session in page.sessions { + transcript_import::validate_session_id(&session.session_id).with_context(|| { + format!("invalid Muse session id on session/list page {page_index}") + })?; + sessions.push(session.session_id); + } + if let Some(cursor) = &page.next_cursor { + if !seen_cursors.insert(cursor.clone()) { + anyhow::bail!("Muse session/list repeated cursor on page {page_index}") + } + } + Ok(page.next_cursor) +} + +async fn muse_drain_stderr(mut stderr: tokio::process::ChildStderr) -> anyhow::Result { + let mut output = Vec::new(); + let mut buffer = [0u8; 4096]; + loop { + let count = stderr.read(&mut buffer).await?; + if count == 0 { + break; + } + let remaining = MUSE_MSP_DIAGNOSTIC_BYTES.saturating_sub(output.len()); + output.extend_from_slice(&buffer[..count.min(remaining)]); + } + Ok(String::from_utf8_lossy(&output).into_owned()) +} + +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'); + tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { + stdin.write_all(frame.as_bytes()).await?; + stdin.flush().await?; + Ok(()) + }) + .await + .map_err(|_| anyhow::anyhow!("Muse MSP timed out sending {method}"))? +} + +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'); + tokio::time::timeout(MUSE_MSP_REQUEST_TIMEOUT, async { + stdin.write_all(request.as_bytes()).await?; + stdin.flush().await?; + 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("jsonrpc").and_then(serde_json::Value::as_str) != Some("2.0") { + anyhow::bail!( + "Muse MSP sent a frame without JSON-RPC 2.0 while waiting for {method}" + ) + } + let Some(response_id) = frame.get("id") else { + // Notifications can be interleaved with responses. + continue; + }; + if response_id.as_u64() != Some(id) { + anyhow::bail!("Muse MSP returned an unexpected response id for {method}") + } + 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}"))? +} + +pub(super) async fn import_muse_sessions_with_ledger( + session_ids: &[String], + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + let mut processor = ImportProcessor::new(opts, config, ledger_dir)?; + let mut summaries = Vec::new(); + for session_id in session_ids { + let directory = tempfile::Builder::new() + .prefix("bt-muse-import-") + .tempdir()?; + let export = directory.path().join("session.json"); + export_muse_session(session_id, &export).await?; + let ended = process_muse_export(&mut processor, &export, session_id).await?; + let summary = if ended { + processor.finish_session(session_id).await? + } else { + processor.detach_session(session_id).await? + }; + if let Some(summary) = summary { + summaries.push(summary); + } + } + summaries.extend(processor.finish().await?); + Ok(summaries) +} + +pub(super) async fn process_muse_export( + processor: &mut ImportProcessor, + export: &std::path::Path, + session_id: &str, +) -> anyhow::Result { + process_muse_snapshot(processor, export, session_id, true, 0) + .await + .map(|(_, ended)| ended) +} + +/// The file is fully parsed on every poll. Only envelopes beyond the previous +/// snapshot are sent to the long-lived translator and sink, so open model and +/// turn state is retained without replaying old starts into it. +pub(super) async fn process_muse_snapshot( + processor: &mut ImportProcessor, + export: &std::path::Path, + session_id: &str, + attach_snapshot: bool, + skip_envelopes: usize, +) -> anyhow::Result<(usize, bool)> { + let (sender, mut receiver) = tokio::sync::mpsc::channel::>(4); + let export_for_parser = export.to_path_buf(); + let expected_id = session_id.to_owned(); + let parser = tokio::task::spawn_blocking(move || { + let mut batch = Vec::with_capacity(128); + let mut envelope_count = 0usize; + let ended = transcript_import::muse::for_each_snapshot_envelope_for_session( + &export_for_parser, + &expected_id, + attach_snapshot, + |entry| { + envelope_count += 1; + if envelope_count <= skip_envelopes { + return Ok(()); + } + batch.push(entry); + if batch.len() == 128 { + sender + .blocking_send(std::mem::take(&mut batch)) + .map_err(|_| { + anyhow::anyhow!("Muse import stopped while parsing {expected_id}") + })?; + } + Ok(()) + }, + )?; + if envelope_count < skip_envelopes { + anyhow::bail!( + "Muse export for {expected_id} shrank from {skip_envelopes} to {envelope_count} trace events" + ); + } + if !batch.is_empty() { + sender + .blocking_send(batch) + .map_err(|_| anyhow::anyhow!("Muse import stopped while parsing {expected_id}"))?; + } + Ok::<(usize, bool), anyhow::Error>((envelope_count, ended)) + }); + let mut processing_error = None; + while let Some(batch) = receiver.recv().await { + if let Err(error) = processor.process(batch).await { + processing_error = Some(error); + break; + } + } + drop(receiver); + let parse_result = parser.await.context("Muse export parser task failed")?; + if let Some(error) = processing_error { + return Err(error); + } + parse_result +} + +pub(super) async fn attach_muse_session( + session_id: &str, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + transcript_import::validate_session_id(session_id)?; + let directory = tempfile::Builder::new() + .prefix("bt-muse-attach-") + .tempdir()?; + let export = directory.path().join("session.json"); + attach_muse_export(&export, session_id, true, opts, config, ledger_dir).await +} + +pub(super) async fn attach_muse_export( + export: &std::path::Path, + session_id: &str, + reexport: bool, + opts: ServeOptions, + config: Option, + ledger_dir: Option, +) -> anyhow::Result> { + let mut processor = ImportProcessor::new(opts, config, ledger_dir)?; + let mut emitted = 0usize; + let shutdown = tokio::signal::ctrl_c(); + tokio::pin!(shutdown); + loop { + if reexport { + export_muse_session(session_id, export).await?; + } + let (count, ended) = + process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; + if count != emitted { + processor.flush_open().await?; + } + emitted = count; + if ended { + return processor.finish().await; + } + tokio::select! { + result = &mut shutdown => { + result?; + break; + } + _ = tokio::time::sleep(Duration::from_secs(1)) => {} + } + } + // Capture one final complete snapshot before interrupting any still-open + // spans, matching the finalization behavior of other attach sources. + if reexport { + export_muse_session(session_id, export).await?; + } + process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; + processor.finish().await +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[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 muse_msp_pages_reject_partial_or_cyclic_session_lists() { + let mut sessions = Vec::new(); + let mut cursors = std::collections::HashSet::new(); + assert_eq!( + muse_msp_append_page( + json!({"sessions":[{"sessionId":"session-a"}],"nextCursor":"next"}), + 1, + &mut sessions, + &mut cursors, + ) + .unwrap(), + Some("next".to_owned()) + ); + assert_eq!(sessions, ["session-a"]); + assert!(muse_msp_append_page( + json!({"sessions":[],"nextCursor":"next"}), + 2, + &mut sessions, + &mut cursors, + ) + .unwrap_err() + .to_string() + .contains("repeated cursor")); + assert!(muse_msp_append_page( + json!({"sessions":[{"status":"notLoaded"}]}), + 2, + &mut sessions, + &mut cursors, + ) + .unwrap_err() + .to_string() + .contains("decode Muse session/list page 2")); + } + + #[tokio::test] + async fn muse_msp_request_accepts_notifications_but_rejects_wrong_response_ids() { + async fn exchange(response: &str) -> anyhow::Result { + let (client, server) = tokio::io::duplex(4096); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, mut server_write) = tokio::io::split(server); + let response = response.to_owned(); + let server = tokio::spawn(async move { + let request = BufReader::new(server_read) + .lines() + .next_line() + .await + .unwrap() + .unwrap(); + assert_eq!( + serde_json::from_str::(&request).unwrap()["id"], + 7 + ); + server_write.write_all(response.as_bytes()).await.unwrap(); + }); + let result = muse_msp_request( + &mut client_write, + &mut BufReader::new(client_read).lines(), + 7, + "session/list", + json!({}), + ) + .await; + server.await.unwrap(); + result + } + + let result = exchange( + "{\"jsonrpc\":\"2.0\",\"method\":\"progress\"}\n{\"jsonrpc\":\"2.0\",\"id\":7,\"result\":{\"sessions\":[]}}\n", + ) + .await + .unwrap(); + assert_eq!(result, json!({"sessions": []})); + assert!(exchange("{\"jsonrpc\":\"2.0\",\"id\":8,\"result\":{}}\n") + .await + .unwrap_err() + .to_string() + .contains("unexpected response id")); + } +} diff --git a/bt-daemon/src/transcript_import/pipeline.rs b/bt-daemon/src/transcript_import/pipeline.rs index 7ccfdf12..b7725bcb 100644 --- a/bt-daemon/src/transcript_import/pipeline.rs +++ b/bt-daemon/src/transcript_import/pipeline.rs @@ -1,6 +1,9 @@ //! Transcript import: run native transcript records through the translators //! and a sink without a daemon, recording deliveries in the ledger. +#[path = "muse_runtime.rs"] +mod muse_runtime; + use anyhow::Context; use braintrust_sdk_rust::SpanComponents; use std::path::PathBuf; @@ -27,6 +30,9 @@ 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 { + return muse_runtime::run(args, 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 { @@ -193,6 +199,28 @@ async fn import_transcript_with_ledger( attach: bool, ledger_dir: Option, ) -> anyhow::Result> { + if source == ImportSource::Muse { + let session_id = transcript_import::muse::session_id(file)?; + if attach { + return muse_runtime::attach_muse_export( + file, + &session_id, + false, + opts, + config, + ledger_dir, + ) + .await; + } + let mut processor = ImportProcessor::new(opts, config, ledger_dir)?; + let ended = muse_runtime::process_muse_export(&mut processor, file, &session_id).await?; + let summary = if ended { + processor.finish_session(&session_id).await? + } else { + processor.detach_session(&session_id).await? + }; + return Ok(summary.into_iter().collect()); + } let mut tail = transcript_import::TranscriptTail::new(file.to_path_buf(), source); if attach { tail.allow_incomplete_final_record_on_shutdown(); @@ -241,6 +269,22 @@ async fn import_transcripts_with_ledger( config: Option, ledger_dir: Option, ) -> anyhow::Result> { + if source == ImportSource::Muse { + let mut processor = ImportProcessor::new(opts, config, ledger_dir)?; + let mut summaries = Vec::new(); + for file in files { + let session_id = transcript_import::muse::session_id(file)?; + let ended = + muse_runtime::process_muse_export(&mut processor, file, &session_id).await?; + let summary = if ended { + processor.finish_session(&session_id).await? + } else { + processor.detach_session(&session_id).await? + }; + summaries.extend(summary); + } + return Ok(summaries); + } let mut processor = ImportProcessor::new(opts, config, ledger_dir)?; let mut summaries = Vec::new(); for file in files { @@ -473,7 +517,7 @@ impl ImportProcessor { let ops = live.translator.flush(&live.ctx)?; Self::emit_translator_batches(&mut live, ops).await?; live.sink.flush().await?; - summaries.push(Self::summary(sid, live)); + summaries.push(Self::summary(sid, live, true)); } summaries.sort_by(|left, right| left.session_id.cmp(&right.session_id)); Ok(summaries) @@ -486,27 +530,251 @@ impl ImportProcessor { let ops = live.translator.flush(&live.ctx)?; Self::emit_translator_batches(&mut live, ops).await?; live.sink.flush().await?; - Ok(Some(Self::summary(session_id.to_string(), live))) + Ok(Some(Self::summary(session_id.to_string(), live, true))) + } + + async fn flush_open(&mut self) -> anyhow::Result<()> { + for live in self.sessions.values_mut() { + live.sink.flush().await?; + live.pending_ops = 0; + } + Ok(()) } - fn summary(session_id: String, live: ImportLive) -> ImportSummary { + async fn detach_session(&mut self, session_id: &str) -> anyhow::Result> { + let Some(mut live) = self.sessions.remove(session_id) else { + return Ok(None); + }; + live.sink.flush().await?; + Ok(Some(Self::summary(session_id.to_string(), live, false))) + } + + fn summary(session_id: String, live: ImportLive, finalized: bool) -> ImportSummary { ImportSummary { session_id, destination: live.destination, root_span_id: live.root_span_id, span_count: live.span_ids.len(), - finalized: true, + finalized, } } } #[cfg(test)] mod tests { + use super::muse_runtime::*; use super::*; use crate::debug_serve_options; + use crate::{ids, DebugSinkFactory, Registry}; use braintrust_sdk_rust::SpanObjectType; use clap::Parser; use serde_json::json; + use std::sync::Arc; + + #[tokio::test] + async fn muse_attach_snapshots_reuse_one_trace_and_keep_open_spans() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("export.json"); + let output = temp.path().join("spans"); + let mut processor = ImportProcessor::new( + ServeOptions { + version: "test".into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(DebugSinkFactory { + dir: output.clone(), + }), + auth_provider: None, + }, + None, + None, + ) + .unwrap(); + let mut events = vec![ + json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}), + json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","request_record_id":"request-1"}}}}), + ]; + let write = |events: &[serde_json::Value]| { + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-attach"}], + "events":events, + })) + .unwrap(), + ) + .unwrap(); + }; + write(&events); + let (count, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, 0) + .await + .unwrap(); + assert_eq!(count, 3); + assert!(!ended); + processor.flush_open().await.unwrap(); + let (same_count, ended) = + process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) + .await + .unwrap(); + assert_eq!(same_count, count); + assert!(!ended); + events.extend([ + json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_completed","usage":{"input_tokens":2}}}}}), + json!({"recorded_at":1_003_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"done"}}}}), + json!({"recorded_at":1_004_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), + json!({"recorded_at":1_005_000,"envelope":{"payload":{"kind":"session_end"}}}), + ]); + write(&events); + let (_, ended) = process_muse_snapshot(&mut processor, &export, "muse-attach", true, count) + .await + .unwrap(); + assert!(ended); + processor.finish().await.unwrap(); + let rows = std::fs::read_to_string(output.join("muse-attach.ndjson")) + .unwrap() + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + for identity in ["root", "turn:run-1", "llm:request-1"] { + let id = ids::span_id(&ids::session_namespace("muse", "muse-attach"), identity); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(|value| value.as_str()) + == Some(id.as_str())) + .count(), + 1, + "duplicate insert for {identity}" + ); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Merge/span_id") + .and_then(|value| value.as_str()) + == Some(id.as_str())) + .count(), + 1, + "missing or duplicate completion for {identity}" + ); + } + } + + #[tokio::test] + async fn repeated_muse_import_does_not_emit_a_second_trace() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("muse-export.json"); + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-ledger"}], + "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":"terminal","terminal":"completed"}}}}, + {"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}} + ] + })) + .unwrap(), + ) + .unwrap(); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let import = || { + import_transcript_with_ledger( + &export, + ImportSource::Muse, + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + false, + Some(ledger.clone()), + ) + }; + import().await.unwrap(); + let rows = output.join("spans/muse-ledger.ndjson"); + let first = std::fs::read_to_string(&rows).unwrap(); + assert!(first.contains("Muse Code")); + import().await.unwrap(); + assert_eq!(std::fs::read_to_string(&rows).unwrap(), first); + } + + #[tokio::test] + async fn later_muse_import_completes_an_active_snapshot_on_the_same_root() { + let temp = tempfile::tempdir().unwrap(); + let export = temp.path().join("muse-export.json"); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let started = json!({"recorded_at":1_000_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}); + let write = |events: Vec| { + std::fs::write( + &export, + serde_json::to_vec(&json!({ + "export_schema_version":1, + "sessions":[{"session_id":"muse-active-import"}], + "events":events, + })) + .unwrap(), + ) + .unwrap(); + }; + let import = || { + import_transcript_with_ledger( + &export, + ImportSource::Muse, + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + false, + Some(ledger.clone()), + ) + }; + write(vec![started.clone()]); + let summaries = import().await.unwrap(); + assert!(!summaries[0].finalized); + let rows_path = output.join("spans/muse-active-import.ndjson"); + let first = std::fs::read_to_string(&rows_path).unwrap(); + assert!(!first.contains("\"Merge\"")); + + write(vec![ + started, + json!({"recorded_at":1_001_000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}), + json!({"recorded_at":1_002_000,"envelope":{"payload":{"kind":"session_end"}}}), + ]); + let summaries = import().await.unwrap(); + assert!(summaries[0].finalized); + let rows = std::fs::read_to_string(&rows_path).unwrap(); + let parsed = rows + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + let root_id = ids::span_id( + &ids::session_namespace("muse", "muse-active-import"), + "root", + ); + assert_eq!( + parsed + .iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(|value| value.as_str()) + == Some(root_id.as_str())) + .count(), + 2 + ); + assert_eq!( + parsed + .iter() + .filter(|row| row + .pointer("/Merge/span_id") + .and_then(|value| value.as_str()) + == Some(root_id.as_str())) + .count(), + 1 + ); + assert!(rows.contains("\"end_ms\":1002")); + let completed = rows.clone(); + import().await.unwrap(); + assert_eq!(std::fs::read_to_string(&rows_path).unwrap(), completed); + } #[derive(Debug, Parser)] struct ImportCli { diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index 754a91b2..d39e9a28 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -10,24 +10,23 @@ mod antigravity; mod claude; mod codex; +pub(crate) mod cursor; mod debug; +mod decode; mod git; mod grok; mod muse; mod opencode; mod pi; +pub(crate) use pi::request_config as pi_request_config; mod recent; mod tool; +mod turn_lineage; +pub use turn_lineage::TURN_SPAN_ID_KEY; -pub use antigravity::AntigravityTranslatorFactory; -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; +use debug::DebugTranslatorFactory; +use crate::agents::{registrar, Agent}; use crate::wire::{Envelope, SessionConfig}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -54,6 +53,14 @@ pub enum SpanType { Wait, } +/// Immutable versions captured when a coding-agent span is first created. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct OriginSnapshot { + pub plugin_version: Option, + pub bt_version: String, + pub source_version: Option, +} + /// A resolved span row, ready for a sink to insert or merge. Field set is the /// subset every current plugin uses; extend as translators need more. #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -84,9 +91,21 @@ pub struct SpanRow { #[serde(skip)] #[doc(hidden)] pub late_merge_key: Option, + /// Span id of the user turn that owns this row, recorded by the translator + /// from its native turn tracking. [`turn_lineage`] publishes it as + /// `metadata.turn_span_id`. + #[serde(skip)] + #[doc(hidden)] + pub turn_span_id: Option, /// Labels for filtering in Braintrust (e.g. `compaction`, `permission-request`). #[serde(default, skip_serializing_if = "Option::is_none")] pub tags: Option>, + /// Translator-owned marker for turn creation, not other task spans. + #[serde(skip)] + pub is_turn: bool, + /// Delivery-owned creation snapshot, present on each coding-agent span. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin: Option, } /// A span operation. `Insert` creates (or replaces) a row; `Merge` updates an @@ -160,15 +179,27 @@ pub trait AgentTranslator: Send { } } -/// Builds translator instances for a given `source`. -pub trait TranslatorFactory: Send + Sync { - fn source(&self) -> &str; +/// Builds translator instances for an agent's hook source. Each agent type +/// implements this in its translator module and registers itself in +/// [`crate::agents`]; the source is the agent's [`Identity::source`]. +/// +/// [`Identity::source`]: crate::agents::Identity::source +pub(crate) trait TranslatorFactory: Agent { fn create(&self, session_id: &str) -> Box; } /// Maps canonical and supported alias source strings to translator factories. pub struct Registry { - factories: HashMap>, + factories: HashMap>, +} + +/// Map an agent name or alias to the source its hooks report. Unknown +/// sources pass through. +pub(crate) fn canonical_source_name(source: &str) -> &str { + registrar() + .translate + .get(source) + .map_or(source, |agent| agent.identity().source) } impl Registry { @@ -177,20 +208,16 @@ impl Registry { let mut r = Registry { factories: HashMap::new(), }; - r.register(Box::new(DebugTranslatorFactory)); - let git = Arc::new(git::GitMetadataCache::default()); - r.register(Box::new(AntigravityTranslatorFactory::new(git.clone()))); - 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.insert(Arc::new(DebugTranslatorFactory)); + for agent in registrar().translate.shared() { + r.insert(agent.clone()); + } r } - pub fn register(&mut self, factory: Box) { - self.factories.insert(factory.source().to_string(), factory); + fn insert(&mut self, factory: Arc) { + self.factories + .insert(factory.identity().source.to_string(), factory); } /// Known sources, for the `initialize` capabilities list. @@ -202,16 +229,11 @@ impl Registry { /// Resolve daemon source aliases to one stable identity. pub fn canonical_source<'a>(&'a self, source: &'a str) -> Option<&'a str> { - let canonical = match source { - "claude" | "claude-code" => "claude-code", - "open-code" | "opencode" => "opencode", - "antigravity" => "antigravity", - "codex" => "codex", - "grok" => "grok", - "muse" | "muse-code" => "muse", - "pi" => "pi", - "debug" => "debug", - _ => return None, + // Only agent sources and the debug translator are accepted. + let canonical = match registrar().translate.get(source) { + Some(agent) => agent.identity().source, + None if source == "debug" => "debug", + None => return None, }; self.factories.contains_key(canonical).then_some(canonical) } @@ -231,17 +253,21 @@ impl Registry { )) } - pub(crate) fn create_checked_with_session_key( + pub(crate) fn create_checked_with_session_namespace( &self, source: &str, - session_key: &str, + session_namespace: &str, ) -> anyhow::Result> { let canonical = self .canonical_source(source) .ok_or_else(|| anyhow::anyhow!("unsupported coding-agent source {source:?}"))?; let factory = self.factories.get(canonical); - let factory = factory.expect("canonical source must have a factory"); - Ok(factory.create(session_key)) + let factory = factory.ok_or_else(|| { + anyhow::anyhow!("no translator factory for coding-agent source {canonical:?}") + })?; + Ok(Box::new(turn_lineage::TurnLineage::new( + factory.create(session_namespace), + ))) } /// Create a known translator. Production ingress uses @@ -252,3 +278,45 @@ impl Registry { .unwrap_or_else(|error| panic!("{error}")) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_agent_source_has_a_translator() { + assert_eq!( + Registry::default_agents().sources(), + [ + "antigravity", + "claude-code", + "codex", + "cursor", + "debug", + "grok", + "muse", + "opencode", + "pi" + ] + ); + } + + #[test] + fn source_aliases_resolve_to_hook_sources() { + let registry = Registry::default_agents(); + for (alias, source) in [ + ("claude", "claude-code"), + ("claude-code", "claude-code"), + ("open-code", "opencode"), + ("opencode", "opencode"), + ("codex", "codex"), + ("muse-code", "muse"), + ("debug", "debug"), + ] { + assert_eq!(registry.canonical_source(alias), Some(source)); + } + assert_eq!(registry.canonical_source("agy"), None); + assert_eq!(canonical_source_name("claude"), "claude-code"); + assert_eq!(canonical_source_name("unknown"), "unknown"); + } +} diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs index 97284aec..a958b6b9 100644 --- a/bt-daemon/src/translate/muse.rs +++ b/bt-daemon/src/translate/muse.rs @@ -16,20 +16,9 @@ use serde_json::{json, Map, Value}; use std::collections::{HashMap, HashSet}; 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" - } +impl TranslatorFactory for crate::agents::Muse { fn create(&self, session_id: &str) -> Box { - Box::new(MuseTranslator::new(session_id, self.git.clone())) + Box::new(MuseTranslator::new(session_id, GitMetadataCache::shared())) } } @@ -37,6 +26,7 @@ impl TranslatorFactory for MuseTranslatorFactory { struct OpenOp { id: String, parents: Vec, + turn_span_id: Option, } #[derive(Default, Deserialize)] @@ -183,6 +173,7 @@ impl MuseTranslator { span_id: open.id, root_span_id: root_span_id.into(), parent_span_ids: open.parents, + turn_span_id: open.turn_span_id, end_ms: Some(ts), output, metrics, @@ -202,6 +193,7 @@ impl MuseTranslator { span_id: open.id, root_span_id: self.effective_root_id.clone(), parent_span_ids: open.parents, + turn_span_id: open.turn_span_id, end_ms: Some(ts), error: error.map(str::to_owned), ..Default::default() @@ -227,6 +219,7 @@ impl MuseTranslator { OpenOp { id: row.span_id.clone(), parents: row.parent_span_ids.clone(), + turn_span_id: row.turn_span_id.clone(), }, ); vec![SpanOp::Insert(row)] @@ -291,6 +284,11 @@ impl AgentTranslator for MuseTranslator { span_type: SpanType::Task, start_ms: Some(e.ts_ms), input: p.get("prompt").cloned(), + is_turn: true, + turn_span_id: Some(ids::span_id( + &self.session_id, + &format!("turn:{id}"), + )), ..Default::default() }, )); @@ -321,7 +319,8 @@ impl AgentTranslator for MuseTranslator { SpanRow { span_id: ids::span_id(&self.session_id, &format!("llm:{id}")), root_span_id: self.effective_root_id.clone(), - parent_span_ids: vec![parent], + parent_span_ids: vec![parent.clone()], + turn_span_id: (parent != self.root_id).then_some(parent), name: p .get("model") .and_then(Value::as_str) @@ -516,6 +515,24 @@ mod tests { |op| matches!(op, SpanOp::Insert(row) if row.name == "Turn" && row.end_ms.is_none()) )); let llm_start = translator.handle(&events[2], &ctx).unwrap(); + let turn_id = turn_start + .iter() + .find_map(|op| match op { + SpanOp::Insert(row) if row.is_turn => Some(row.span_id.clone()), + _ => None, + }) + .expect("native user turn marker"); + for op in turn_start.iter().chain(llm_start.iter()) { + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op; + if row.span_id != root_id { + assert_eq!( + row.metadata + .as_ref() + .and_then(|metadata| metadata.get(super::super::TURN_SPAN_ID_KEY)), + Some(&json!(turn_id)) + ); + } + } let llm_id = llm_start .iter() .find_map(|op| match op { diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index f49e321d..ec7265fd 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -25,7 +25,7 @@ use bt_daemon::{ Registry, ServeArgs, ServeOptions, Sink, SinkFactory, SpanOp, StatusArgs, }; #[cfg(all(feature = "cli", unix))] -use bt_daemon::{run_traced, RunArgs, RunHookCommand, RunSource, SpanType}; +use bt_daemon::{run_traced_with_host as run_traced, RunArgs, RunHookCommand, RunSource, SpanType}; use std::collections::{HashMap, HashSet}; use std::ffi::OsString; use std::path::{Path, PathBuf}; @@ -2980,9 +2980,6 @@ async fn muse_managed_run_starts_its_isolated_daemon_before_sanitized_hooks() { assert!(status.success()); } -#[cfg(all(feature = "cli", unix))] -type CreatedDebugSinks = Arc)>>>; - #[cfg(all(feature = "cli", unix))] struct FlushTrackingDebugSinkFactory { inner: DebugSinkFactory, From 6dc2a781a2d05e152cb5b82cd47bfef7b5150e12 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:49:08 +0800 Subject: [PATCH 17/17] Complete Muse tracing coverage and replay correctness --- .github/workflows/ci.yml | 37 ++ bt-daemon/README.md | 67 +- bt-daemon/src/delivery_ledger.rs | 155 ++++- bt-daemon/src/hook.rs | 71 ++ bt-daemon/src/managed_run/mod.rs | 12 +- bt-daemon/src/managed_run/muse.rs | 118 +++- bt-daemon/src/pipeline/journal.rs | 27 + bt-daemon/src/setup/muse.rs | 613 +++++++++--------- bt-daemon/src/span_processor.rs | 1 + bt-daemon/src/transcript_import/muse.rs | 335 +++++++++- .../src/transcript_import/muse_runtime.rs | 23 +- bt-daemon/src/transcript_import/pipeline.rs | 226 ++++++- bt-daemon/src/translate/debug.rs | 2 + bt-daemon/src/translate/mod.rs | 11 + bt-daemon/src/translate/muse.rs | 334 +++++++++- bt-daemon/src/translate/muse_events.rs | 474 ++++++++++++++ bt-daemon/tests/braintrust_sink.rs | 1 + .../tests/fixtures/muse-export-1.4.11.json | 347 ++++++++++ .../muse-export-tool-failure-1.4.11.json | 348 ++++++++++ .../tests/fixtures/muse-hooks-1.4.10.jsonl | 11 + bt-daemon/tests/muse_attach.rs | 301 +++++++++ bt-daemon/tests/muse_integration.rs | 188 ++++++ bt-daemon/tests/pipeline.rs | 68 +- bt-daemon/tests/support/agents/mod.rs | 3 + bt-daemon/tests/support/agents/muse.rs | 120 ++++ bt-daemon/tests/support/inference/openai.rs | 51 +- 26 files changed, 3506 insertions(+), 438 deletions(-) create mode 100644 bt-daemon/src/translate/muse_events.rs create mode 100644 bt-daemon/tests/fixtures/muse-export-1.4.11.json create mode 100644 bt-daemon/tests/fixtures/muse-export-tool-failure-1.4.11.json create mode 100644 bt-daemon/tests/fixtures/muse-hooks-1.4.10.jsonl create mode 100644 bt-daemon/tests/muse_attach.rs create mode 100644 bt-daemon/tests/muse_integration.rs create mode 100644 bt-daemon/tests/support/agents/muse.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cec29b0c..bc741387 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -404,6 +404,42 @@ jobs: if: runner.os != 'macOS' run: cargo clippy --manifest-path bt-daemon/Cargo.toml --all-targets --all-features --locked -- -D warnings + muse-integration: + name: Muse native integration (Linux) + env: + TBH_CREDENTIAL_BACKEND: file + TBH_DISABLE_TELEMETRY: "1" + runs-on: ubuntu-24.04 + timeout-minutes: 20 + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 + - name: Install current Muse using its official installer + shell: bash + env: + MUSE_INSTALL_DIR: ${{ runner.temp }}/muse + MUSE_NO_MODIFY_PATH: "1" + run: | + curl --fail --silent --show-error --location https://dev.meta.ai/install.sh -o "$RUNNER_TEMP/install-muse.sh" + bash "$RUNNER_TEMP/install-muse.sh" + "$MUSE_INSTALL_DIR/muse" --version + echo "MUSE_BIN=$MUSE_INSTALL_DIR/muse" >> "$GITHUB_ENV" + - name: Install Rust + run: | + rustup toolchain install stable --profile minimal + rustup default stable + - name: Cache Rust build + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + with: + workspaces: bt-daemon -> target + shared-key: muse-integration + save-if: ${{ github.event_name != 'pull_request' }} + - name: Trace native tools, retries, and failures through mock ingest + run: cargo test --manifest-path bt-daemon/Cargo.toml --all-features --locked --test muse_integration -- --ignored --nocapture + - name: Verify native active attach, detach, and resumed completion + run: cargo test --manifest-path bt-daemon/Cargo.toml --all-features --locked --test muse_attach -- --ignored --nocapture + - name: Verify native managed-hook coexistence and sanitized environments + run: cargo test --manifest-path bt-daemon/Cargo.toml --all-features --locked --test pipeline muse_managed_run -- --ignored --nocapture --test-threads=1 + checks-passed: name: checks-passed needs: @@ -414,6 +450,7 @@ jobs: - pi-compatibility - resolve-coding-agent-versions - daemon + - muse-integration runs-on: ubuntu-24.04 timeout-minutes: 5 if: always() diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 6ac1a269..80ee85d6 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -285,32 +285,49 @@ 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. Exports are parsed incrementally, one session at a time. -Muse Code 1.1.1 setup captures the six verified lifecycle and model hooks: -`SessionStart`, `UserPromptSubmit`, `PreLLMCall`, `PostLLMCall`, `Stop`, and -`SessionEnd`. Tool, permission, subagent, and compaction details are not -available from those hooks and are not inferred by the live translator. +Muse Code 1.3+ setup captures lifecycle, model, tool success/failure, permission +requests, subagent lifecycle, compaction, batch completion, and notifications. +Installation adds owned user hooks while preserving unrelated user hooks and +enterprise managed-hook policy. Disablement removes only Braintrust's hooks. + `import muse --attach` follows an active session by polling full -exports. Each snapshot is parsed with the same reader and translator as -historical import; only new envelopes are delivered, and an active snapshot -does not invent `Stop` or `SessionEnd` events. The attach exits after a native -session end, or finalizes open spans on interruption. This reads the full -export each poll, so polling cost grows with long sessions. MSP offers -cursor-paged `view/page`, but its view events omit the model request/response -records needed by this translator; `view/subscribe` also requires a session -loaded on the same host. -One-shot imports of active Muse sessions also leave spans open, allowing a -later import to add the native completion to the same deterministic trace. -`run muse` uses a process-local managed-hook override in the verified Muse Code -1.1.1 release. It retains compatible unrelated managed hooks, replaces Braintrust's saved -managed hooks only for that invocation, and leaves ordinary Muse sessions and -settings unchanged. The injected hook reads a private invocation context file -because Muse clears its hook environment. A dedicated daemon is started before -Muse so environment-based Braintrust authentication never depends on that -cleared environment. Known duplicate Braintrust user or project hooks stop the -managed run with a diagnostic. A foreign managed file using Muse's strict -handler schema is also rejected rather than silently losing hooks. Muse plugins that independently trace the same -session should be disabled before using `run muse`; new Muse versions must be -verified before the undocumented hook-path override is enabled for them. +exports with the same parser and translator used for historical import. Only +new envelopes are delivered. Active snapshots leave spans open; Ctrl-C detaches +without fabricating a terminal result or launching a final export. A later attach +or import applies authoritative completion to the same span IDs. Polling cost +grows with session length; MSP view events omit the model request/response records +needed by the translator. + +Live model hooks and schema-v1 exports use the native turn, model step, and retry +attempt for shared span identity; the export's record UUID is retained as metadata. +Legacy exports without a model step retain their record UUID, so cross-ingress +model deduplication cannot be guaranteed for those records. Native response +previews remain previews: tracing cannot recover content Muse does not export. +Permission hooks report requests, not approval decisions; compaction hooks lack +native operation IDs. Captured built-in child hooks identify their child session +but omit the parent session, so their work remains in a separate trace rather +than guessing ancestry across concurrent sessions. Missing native facts are not invented. + +`run muse` uses invocation-local `-c hooks=…` configuration and a private context +file because Muse sanitizes the hook environment. It replaces owned user hooks +for that invocation, preserves foreign user hooks and enterprise managed hooks, +and leaves saved settings unchanged. An isolated daemon starts before Muse so +authentication does not depend on the sanitized hook environment. Duplicate +Braintrust project hooks produce a diagnostic. Disable independently tracing +Muse plugins before running a managed session. + +Muse native verification runs on Linux in CI and locally on macOS. The official +installer does not provide a Windows binary; Rust fixture tests still run on +Windows. Native tests use Muse's file credential backend +(`TBH_CREDENTIAL_BACKEND=file`) and disable telemetry so temporary profiles do +not trigger macOS Keychain creation dialogs. Muse hooks ship through the embedding `bt` binary; no separate marketplace +package is needed. After this PR merges, dispatch the consumer's existing +`Update bt-daemon dependency` workflow (`update-bt-daemon.yml` in +`braintrustdata/bt`) to prepare its Git pin and lockfile PR. Merge that consumer +update and build/install `bt`, then verify `bt trace setup +muse`, `bt trace run muse exec …`, `bt trace import muse `, status, +and disablement against an isolated Muse profile. Release publishing remains +owned by the consumer repository. Add `--attach` to keep following an active Codex, Claude, or Antigravity transcript until Ctrl-C. `run [ARGS...]` launches the selected agent with diff --git a/bt-daemon/src/delivery_ledger.rs b/bt-daemon/src/delivery_ledger.rs index 8464e008..ebd8aa9e 100644 --- a/bt-daemon/src/delivery_ledger.rs +++ b/bt-daemon/src/delivery_ledger.rs @@ -8,7 +8,7 @@ //! destination still receives the full trace. use crate::sink::Sink; -use crate::translate::{OriginSnapshot, SpanOp, SpanRow}; +use crate::translate::{LateMergeVersion, OriginSnapshot, SpanOp, SpanRow}; use crate::wire::SessionConfig; use sha2::{Digest, Sha256}; use std::collections::{HashMap, HashSet}; @@ -31,6 +31,8 @@ struct LedgerFile { late_merge_ids: HashSet, #[serde(default, skip_serializing_if = "HashMap::is_empty")] span_origins: HashMap, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + late_merge_versions: HashMap, } // Closing this file releases the OS lock, including on cancellation or errors. @@ -45,6 +47,8 @@ struct DeliveryLedger { pending: HashSet, known_late_merges: HashSet, pending_late_merges: HashSet, + known_versions: HashMap, + pending_versions: HashMap, } impl DeliveryLedger { @@ -80,6 +84,8 @@ impl DeliveryLedger { pending: HashSet::new(), known_late_merges: persisted.late_merge_ids, pending_late_merges: HashSet::new(), + known_versions: persisted.late_merge_versions, + pending_versions: HashMap::new(), }; let lock = ledger.lock().await?; ledger @@ -143,6 +149,9 @@ impl DeliveryLedger { let update: LedgerFile = serde_json::from_slice(&line)?; self.known.extend(update.completed_span_ids); self.known_late_merges.extend(update.late_merge_ids); + for (span, version) in update.late_merge_versions { + remember_version(&mut self.known_versions, &span, &version); + } for (span, origin) in update.span_origins { if pending_origins.remove(&span).is_some() { // Another writer persisted first while this writer's @@ -170,10 +179,25 @@ impl DeliveryLedger { } fn filter(&self, ops: &[SpanOp]) -> Vec { - ops.iter() - .filter(|op| self.should_emit(op)) - .cloned() - .collect() + let mut versions = self.known_versions.clone(); + for (span, version) in &self.pending_versions { + remember_version(&mut versions, span, version); + } + let mut filtered = Vec::new(); + for op in ops { + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op; + if !self.should_emit(op) { + continue; + } + if let Some(version) = &row.late_merge_version { + if !version_allowed(versions.get(&row.span_id), version) { + continue; + } + remember_version(&mut versions, &row.span_id, version); + } + filtered.push(op.clone()); + } + filtered } fn record_emitted(&mut self, ops: &[SpanOp]) { @@ -181,6 +205,9 @@ impl DeliveryLedger { let row = match op { SpanOp::Insert(row) | SpanOp::Merge(row) => row, }; + if let Some(version) = &row.late_merge_version { + remember_version(&mut self.pending_versions, &row.span_id, version); + } if row.end_ms.is_some() { self.pending.insert(row.span_id.clone()); } @@ -194,7 +221,10 @@ impl DeliveryLedger { self.pending.retain(|id| !self.known.contains(id)); self.pending_late_merges .retain(|id| !self.known_late_merges.contains(id)); - if self.pending.is_empty() && self.pending_late_merges.is_empty() { + if self.pending.is_empty() + && self.pending_late_merges.is_empty() + && self.pending_versions.is_empty() + { return Ok(()); } // Only new acknowledgements are serialized. Retain pending state until @@ -203,6 +233,7 @@ impl DeliveryLedger { &LedgerFile { completed_span_ids: self.pending.clone(), late_merge_ids: self.pending_late_merges.clone(), + late_merge_versions: self.pending_versions.clone(), ..Default::default() }, lock, @@ -211,6 +242,9 @@ impl DeliveryLedger { self.known.extend(self.pending.drain()); self.known_late_merges .extend(self.pending_late_merges.drain()); + for (span, version) in self.pending_versions.drain() { + remember_version(&mut self.known_versions, &span, &version); + } Ok(()) } @@ -244,6 +278,26 @@ impl DeliveryLedger { } } +fn version_allowed(previous: Option<&LateMergeVersion>, next: &LateMergeVersion) -> bool { + previous.is_none_or(|previous| { + next.ts_ms >= previous.ts_ms + || (next.native_id.is_some() && next.native_id == previous.native_id) + }) +} + +fn remember_version( + versions: &mut HashMap, + span: &str, + next: &LateMergeVersion, +) { + if versions + .get(span) + .is_none_or(|previous| next.ts_ms >= previous.ts_ms) + { + versions.insert(span.to_owned(), next.clone()); + } +} + fn late_merge_id(row: &SpanRow, key: &str) -> String { format!("{}\u{1f}{key}", row.span_id) } @@ -374,7 +428,7 @@ impl Sink for LedgerSink { true } }); - filtered.retain(|op| ledger.should_emit(op)); + filtered = ledger.filter(&filtered); Some(lock) } else { None @@ -774,6 +828,7 @@ mod tests { completed_span_ids: ledger.known, late_merge_ids: ledger.known_late_merges, span_origins: origins, + late_merge_versions: ledger.known_versions, }) .unwrap(), ) @@ -964,6 +1019,92 @@ mod tests { assert!(repeated_output.lock().is_empty()); } + #[tokio::test] + async fn aggregate_watermarks_reject_older_operations_but_allow_latest_native_enrichment() { + let temp = tempfile::tempdir().unwrap(); + let cfg = config("project-a"); + let versioned = |key: &str, ts_ms, native: &str| { + let mut op = late_merge("root", key); + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = &mut op; + row.late_merge_version = Some(LateMergeVersion { + ts_ms, + native_id: Some(native.into()), + }); + op + }; + let mut first = LedgerSink::new( + Box::new(RecordingSink::default()), + Some(temp.path()), + "muse", + "session", + Some(&cfg), + "test", + ) + .await + .unwrap(); + // Filtering must consider earlier rows in this same uncommitted batch. + assert_eq!( + first + .emit(&[ + versioned("t2-preview", 200, "t2"), + versioned("t1-old", 100, "t1") + ]) + .await + .unwrap(), + 1 + ); + first.flush().await.unwrap(); + drop(first); + let mut recovered = LedgerSink::new( + Box::new(RecordingSink::default()), + Some(temp.path()), + "muse", + "session", + Some(&cfg), + "test", + ) + .await + .unwrap(); + // Native exports precede hook capture timestamps, but belong to the + // same latest operation and can replace its truncated preview. + assert_eq!( + recovered + .emit(&[versioned("t2-full", 150, "t2")]) + .await + .unwrap(), + 1 + ); + assert_eq!( + recovered + .emit(&[versioned("t1-full", 190, "t1")]) + .await + .unwrap(), + 0 + ); + assert_eq!( + recovered + .emit(&[versioned("t2-full", 150, "t2")]) + .await + .unwrap(), + 0 + ); + assert_eq!( + recovered + .emit(&[versioned("t3-full", 201, "t3")]) + .await + .unwrap(), + 1 + ); + assert_eq!( + recovered + .emit(&[versioned("t2-different", 200, "t2")]) + .await + .unwrap(), + 0 + ); + recovered.flush().await.unwrap(); + } + #[tokio::test] async fn a_completed_span_receives_each_distinct_late_merge_once() { let temp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/hook.rs b/bt-daemon/src/hook.rs index 53e58856..9f5feb35 100644 --- a/bt-daemon/src/hook.rs +++ b/bt-daemon/src/hook.rs @@ -199,6 +199,9 @@ async fn capture_hook( anyhow::bail!("Muse managed hook context is incomplete"); } route = context.route.clone(); + if args.source_version.is_none() { + args.source_version = context.source_version.clone(); + } Some(context) } else { None @@ -247,9 +250,32 @@ async fn capture_hook( } fn resolve_dynamic_hook_versions(args: &mut HookArgs, payload: &serde_json::Value) { + resolve_dynamic_hook_versions_with_muse_probe(args, payload, || { + crate::setup::muse::installed_muse_version() + .ok() + .map(|version| version.to_string()) + }); +} + +fn resolve_dynamic_hook_versions_with_muse_probe( + args: &mut HookArgs, + payload: &serde_json::Value, + muse_probe: impl FnOnce() -> Option, +) { if args.source_version.is_none() { args.source_version = source_version_from_env(&args.source) .or_else(|| source_version_from_payload(&args.source, payload)); + let event = args.event.as_deref().or_else(|| { + payload + .get(&args.event_field) + .and_then(serde_json::Value::as_str) + }); + if args.source_version.is_none() + && crate::translate::canonical_source_name(&args.source) == "muse" + && matches!(event, Some("SessionStart" | "session_start")) + { + args.source_version = muse_probe(); + } } if args.plugin_version.is_none() { args.plugin_version = plugin_version_from_plugin_root(&args.source); @@ -622,6 +648,51 @@ mod tests { assert_ne!(initialize["client"]["plugin_version"], "1.0.13"); } + #[test] + fn muse_version_probe_is_session_start_only_and_fail_open() { + let mut args = HookCli::try_parse_from(["test", "--source", "muse"]) + .unwrap() + .args; + resolve_dynamic_hook_versions_with_muse_probe( + &mut args, + &json!({"hook_event_name":"UserPromptSubmit"}), + || panic!("must not spawn a version probe on each event"), + ); + assert!(args.source_version.is_none()); + resolve_dynamic_hook_versions_with_muse_probe( + &mut args, + &json!({"hook_event_name":"SessionStart"}), + || None, + ); + assert!(args.source_version.is_none()); + resolve_dynamic_hook_versions_with_muse_probe( + &mut args, + &json!({"hook_event_name":"SessionStart"}), + || Some("1.4.11".into()), + ); + let envelope = build_hook_envelope( + &args, + SessionRoute::default(), + json!({}), + "session".into(), + "SessionStart".into(), + ); + assert_eq!(envelope.source_version.as_deref(), Some("1.4.11")); + resolve_dynamic_hook_versions_with_muse_probe( + &mut args, + &json!({"hook_event_name":"SessionStart"}), + || panic!("explicit context version must be preserved"), + ); + let mut other = HookCli::try_parse_from(["test", "--source", "cursor"]) + .unwrap() + .args; + resolve_dynamic_hook_versions_with_muse_probe( + &mut other, + &json!({"hook_event_name":"SessionStart"}), + || panic!("Muse probe must not run for other agents"), + ); + } + #[test] fn hook_versions_are_discovered_without_hook_configuration_literals() { let payload = serde_json::json!({"cli_version": "2.3.4"}); diff --git a/bt-daemon/src/managed_run/mod.rs b/bt-daemon/src/managed_run/mod.rs index a35ac5dd..1ee4ff79 100644 --- a/bt-daemon/src/managed_run/mod.rs +++ b/bt-daemon/src/managed_run/mod.rs @@ -41,6 +41,15 @@ pub(crate) trait ManagedRun: Agent { managed_run_id: &str, ) -> anyhow::Result; + /// Combine invocation options with user arguments in the agent's native order. + fn invocation_args( + &self, + injection_args: Vec, + agent_args: Vec, + ) -> Vec { + injection_args.into_iter().chain(agent_args).collect() + } + /// Agents that sanitize hook environments need an already-running daemon /// and a private context file instead of relying on inherited settings. fn needs_daemon_host(&self) -> bool { @@ -160,8 +169,7 @@ async fn run_traced_inner( let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?; let mut command: tokio::process::Command = subprocess::interactive_command(&executable).into(); command - .args(injection.args) - .args(args.agent_args) + .args(agent.invocation_args(injection.args, args.agent_args)) .env("_BT_TRACE_MANAGED_RUN", "1") .env(MANAGED_RUN_ID_ENV, &managed_run_id) .env(settings::INVOCATION_SETTINGS_ENV, invocation_settings); diff --git a/bt-daemon/src/managed_run/muse.rs b/bt-daemon/src/managed_run/muse.rs index cef246da..e740f5fa 100644 --- a/bt-daemon/src/managed_run/muse.rs +++ b/bt-daemon/src/managed_run/muse.rs @@ -3,7 +3,6 @@ use super::{Injection, ManagedRun, Terminal}; use crate::agents::Muse; use crate::args::RunHookCommand; use crate::wire::SessionRoute; -use anyhow::Context; use std::ffi::OsString; use std::path::{Path, PathBuf}; @@ -12,6 +11,8 @@ pub(crate) struct MuseManagedContext { pub route: SessionRoute, pub managed_run_id: String, pub socket: PathBuf, + #[serde(default)] + pub source_version: Option, } impl ManagedRun for Muse { @@ -32,19 +33,36 @@ impl ManagedRun for Muse { crate::setup::muse::check_muse_project_hooks(&workspace)?; } } - let executable = crate::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()); - let output = crate::subprocess::background_command(executable) - .arg("--version") - .output() - .context("check Muse Code version for managed hooks")?; - if !output.status.success() - || crate::setup::muse::muse_version(&output.stdout)? != semver::Version::new(1, 1, 1) - { - anyhow::bail!("Muse managed runs require the verified Muse Code 1.1.1 hook override"); + if crate::setup::muse::installed_muse_version()? < semver::Version::new(1, 3, 0) { + anyhow::bail!( + "Muse Code 1.3.0 or newer is required for invocation-local hook overrides" + ); } Ok(()) } + fn invocation_args( + &self, + injection_args: Vec, + mut agent_args: Vec, + ) -> Vec { + // Muse dispatches subcommands before parsing options. Placing --config + // before `exec` accidentally launches its interactive parser instead. + if agent_args + .first() + .and_then(|arg| arg.to_str()) + .is_some_and(|arg| matches!(arg, "exec" | "resume")) + { + let subcommand = agent_args.remove(0); + std::iter::once(subcommand) + .chain(injection_args) + .chain(agent_args) + .collect() + } else { + injection_args.into_iter().chain(agent_args).collect() + } + } + fn needs_daemon_host(&self) -> bool { true } @@ -69,6 +87,7 @@ impl ManagedRun for Muse { route: route.clone(), managed_run_id: id.into(), socket: socket.into(), + source_version: Some(crate::setup::muse::installed_muse_version()?.to_string()), }; std::fs::write(&context_path, serde_json::to_vec(&context)?)?; crate::paths::restrict_file_to_owner(&context_path)?; @@ -83,15 +102,82 @@ impl ManagedRun for Muse { &crate::paths::muse_config_dir(), &command, )?; - let hooks_path = directory.path().join("hooks.json"); - std::fs::write(&hooks_path, serde_json::to_vec(&hooks)?)?; + let mut args = vec![ + OsString::from("--config"), + OsString::from(format!("hooks={}", serde_json::to_string(&hooks)?)), + ]; + // Old persistent installs used a dedicated managed file. Its handlers + // must be suppressed alongside the replaced user hooks for this run. + if crate::setup::muse::muse_has_owned_managed_hooks(&crate::paths::muse_config_dir())? { + args.extend([ + OsString::from("--config"), + OsString::from("managed_hooks_path=\"\""), + ]); + } Ok(Injection { - env: vec![( - "TBH_MANAGED_HOOKS_PATH", - hooks_path.to_string_lossy().into_owned(), - )], + args, guards: vec![Box::new(directory)], ..Injection::default() }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn managed_context_carries_runtime_version_and_reads_old_contexts() { + let context = MuseManagedContext { + route: SessionRoute::default(), + managed_run_id: "run".into(), + socket: PathBuf::from("daemon.sock"), + source_version: Some("1.4.11".into()), + }; + let mut value = serde_json::to_value(&context).unwrap(); + assert_eq!( + serde_json::from_value::(value.clone()) + .unwrap() + .source_version + .as_deref(), + Some("1.4.11") + ); + value.as_object_mut().unwrap().remove("source_version"); + assert!(serde_json::from_value::(value) + .unwrap() + .source_version + .is_none()); + } + + #[test] + fn invocation_options_follow_native_subcommands() { + for subcommand in ["exec", "resume"] { + let args = Muse.invocation_args( + vec!["--config".into(), "hooks={}".into()], + vec![ + subcommand.into(), + "--workspace".into(), + "path with spaces".into(), + "prompt".into(), + ], + ); + assert_eq!( + args, + [ + subcommand, + "--config", + "hooks={}", + "--workspace", + "path with spaces", + "prompt" + ] + .map(OsString::from) + ); + } + let args = Muse.invocation_args( + vec!["--config".into(), "hooks={}".into()], + vec!["prompt".into()], + ); + assert_eq!(args, ["--config", "hooks={}", "prompt"].map(OsString::from)); + } +} diff --git a/bt-daemon/src/pipeline/journal.rs b/bt-daemon/src/pipeline/journal.rs index 16e03446..b67fab1c 100644 --- a/bt-daemon/src/pipeline/journal.rs +++ b/bt-daemon/src/pipeline/journal.rs @@ -44,6 +44,8 @@ struct StoredOp { version: u32, op: SpanOp, late_merge_key: Option, + #[serde(default)] + late_merge_version: Option, turn_span_id: Option, is_turn: bool, } @@ -55,6 +57,7 @@ impl StoredOp { version: 1, op: op.clone(), late_merge_key: row.late_merge_key.clone(), + late_merge_version: row.late_merge_version.clone(), turn_span_id: row.turn_span_id.clone(), is_turn: row.is_turn, })?) @@ -68,6 +71,7 @@ impl StoredOp { ); let (SpanOp::Insert(row) | SpanOp::Merge(row)) = &mut stored.op; row.late_merge_key = stored.late_merge_key; + row.late_merge_version = stored.late_merge_version; row.turn_span_id = stored.turn_span_id; row.is_turn = stored.is_turn; Ok(stored.op) @@ -923,6 +927,29 @@ mod tests { ..SpanRow::default() }) } + #[test] + fn codec_preserves_aggregate_watermarks_and_reads_older_records() { + let mut op = op("root"); + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = &mut op; + let version = crate::translate::LateMergeVersion { + ts_ms: 123, + native_id: Some("turn".into()), + }; + row.late_merge_version = Some(version.clone()); + let encoded = StoredOp::encode(&op).unwrap(); + let decoded = StoredOp::decode(&encoded).unwrap(); + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = decoded; + assert_eq!(row.late_merge_version, Some(version)); + let mut legacy: serde_json::Value = serde_json::from_slice(&encoded).unwrap(); + legacy.as_object_mut().unwrap().remove("late_merge_version"); + let decoded = StoredOp::decode(&serde_json::to_vec(&legacy).unwrap()).unwrap(); + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = decoded; + assert!(row.late_merge_version.is_none()); + assert!(serde_json::to_value(&op).unwrap()["Insert"] + .get("late_merge_version") + .is_none()); + } + fn route(name: &str) -> SessionRoute { SessionRoute { tags: vec![name.into()], diff --git a/bt-daemon/src/setup/muse.rs b/bt-daemon/src/setup/muse.rs index 7d32ecdc..83ee60d3 100644 --- a/bt-daemon/src/setup/muse.rs +++ b/bt-daemon/src/setup/muse.rs @@ -11,11 +11,31 @@ use std::path::{Path, PathBuf}; const MUSE_HOOK_EVENTS: &[&str] = &[ "SessionStart", "UserPromptSubmit", + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PostToolUseFailure", + "PostToolBatch", "PreLLMCall", "PostLLMCall", + "PreCompact", + "PostCompact", + "SubagentStart", + "SubagentStop", "Stop", + "StopFailure", "SessionEnd", + "Notification", ]; +const PREVIOUS_MUSE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreLLMCall", + "PostLLMCall", + "Stop", + "SessionEnd", +]; +const MUSE_HOOK_COMMAND: &str = "bt trace hook --source muse --flush-on-turn-end"; const LEGACY_MUSE_HOOK_EVENTS: &[&str] = &[ "SessionStart", "UserPromptSubmit", @@ -38,10 +58,7 @@ fn muse_hook_path(config_dir: &Path) -> PathBuf { fn muse_hook_config() -> Map { // Muse treats command-hook failures as non-fatal to the agent turn. - muse_hook_config_for( - MUSE_HOOK_EVENTS, - "bt trace hook --source muse --flush-on-turn-end", - ) + muse_hook_config_for(MUSE_HOOK_EVENTS, MUSE_HOOK_COMMAND) } fn muse_hook_config_for(events: &[&str], command: &str) -> Map { @@ -51,7 +68,6 @@ fn muse_hook_config_for(events: &[&str], command: &str) -> Map { ( (*event).to_string(), serde_json::json!([{ - "matcher": "*", "hooks": [{"type": "command", "command": command}], }]), ) @@ -63,73 +79,130 @@ fn muse_hook_config_for(events: &[&str], command: &str) -> Map { ]) } -/// Build an invocation-local managed file without changing Muse's settings. -/// Preserve handlers administered by someone else, but never run two Muse -/// Braintrust forwarders with different destinations in the same session. +/// Build a per-run user hook overlay. Keeping the managed source separate preserves +/// enterprise handlers, including its stricter file-wide validation schema. pub(crate) fn muse_managed_hook_config(config_dir: &Path, command: &str) -> anyhow::Result { let settings = load_object(&config_dir.join("settings.json"))?; - if settings + check_foreign_managed_hooks(config_dir, &settings)?; + let mut hooks = settings .get("hooks") - .is_some_and(muse_contains_braintrust_hook) - { - bail!("Muse user hooks already include Braintrust tracing; remove that duplicate before `bt trace run muse`"); + .cloned() + .unwrap_or_else(|| serde_json::json!({})); + remove_owned_hooks(&mut hooks)?; + if muse_contains_braintrust_hook(&hooks) { + bail!("Muse user hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); } - let existing = match settings.get("managed_hooks_path").and_then(Value::as_str) { - Some(path) => { + append_hooks(&mut hooks, command)?; + Ok(hooks) +} + +fn managed_hook_path(config_dir: &Path, settings: &Map) -> Option { + settings + .get("managed_hooks_path") + .and_then(Value::as_str) + .map(|path| { let path = Path::new(path); - let path = if path.is_absolute() { - path.to_path_buf() + if path.is_absolute() { + path.to_owned() } else { config_dir.join(path) - }; - if !path.exists() || muse_hook_is_owned(&path)? { - Map::from_iter([("schema_version".into(), Value::from(1))]) + } + }) +} + +pub(crate) fn muse_has_owned_managed_hooks(config_dir: &Path) -> anyhow::Result { + let settings = load_object(&config_dir.join("settings.json"))?; + managed_hook_path(config_dir, &settings) + .map(|path| { + if path == muse_hook_path(config_dir) { + muse_hook_is_owned(&path) } else { - let raw = std::fs::read(&path) - .with_context(|| format!("read Muse managed hooks {}", path.display()))?; - let value: Value = serde_json::from_slice(&raw)?; - if muse_contains_braintrust_hook(&value) { - bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); - } - if muse_has_strict_handler_fields(&value) { - bail!("Muse managed hooks use the strict handler schema; cannot safely merge invocation hooks"); - } - value - .as_object() - .cloned() - .context("Muse managed hooks must be a JSON object")? + Ok(false) + } + }) + .transpose() + .map(|owned| owned.unwrap_or(false)) +} + +fn check_foreign_managed_hooks( + config_dir: &Path, + settings: &Map, +) -> anyhow::Result<()> { + if let Some(path) = managed_hook_path(config_dir, settings) { + if path.exists() && (path != muse_hook_path(config_dir) || !muse_hook_is_owned(&path)?) { + let hooks: Value = serde_json::from_slice(&std::fs::read(path)?)?; + if muse_contains_braintrust_hook(&hooks) { + bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); } } - None => Map::from_iter([("schema_version".into(), Value::from(1))]), - }; - let mut result = existing; - let mut hooks = result - .remove("hooks") - .unwrap_or_else(|| Value::Object(Map::new())); - let hooks = hooks + } + Ok(()) +} + +fn owned_command(value: &Value) -> bool { + value.get("type").and_then(Value::as_str) == Some("command") + && matches!( + value.get("command").and_then(Value::as_str), + Some(MUSE_HOOK_COMMAND) + | Some("bt trace hook --source muse --flush-on-turn-end || true") + ) + && value.as_object().is_some_and(|object| object.len() == 2) +} + +fn remove_owned_hooks(hooks: &mut Value) -> anyhow::Result<()> { + let events = hooks .as_object_mut() - .context("Muse managed hooks must be an object")?; + .context("Muse user hooks must be an object")?; + let mut emptied_events = Vec::new(); + for (event, groups) in events.iter_mut() { + let groups = groups + .as_array_mut() + .context("Muse hook event must be an array")?; + let original_len = groups.len(); + for index in (0..groups.len()).rev() { + let handlers = groups[index] + .get_mut("hooks") + .and_then(Value::as_array_mut) + .context("Muse hook group must contain a hooks array")?; + let original_handlers = handlers.len(); + handlers.retain(|handler| !owned_command(handler)); + if original_handlers != handlers.len() && handlers.is_empty() { + groups.remove(index); + } + } + if original_len != groups.len() && groups.is_empty() { + emptied_events.push(event.clone()); + } + } + for event in emptied_events { + events.remove(&event); + } + Ok(()) +} + +fn append_hooks(hooks: &mut Value, command: &str) -> anyhow::Result<()> { + let events = hooks + .as_object_mut() + .context("Muse hooks must be an object")?; let generated = muse_hook_config_for(MUSE_HOOK_EVENTS, command); - for (event, handlers) in generated["hooks"] + for (event, groups) in generated["hooks"] .as_object() .context("generated Muse hooks must be an object")? { - let existing = hooks + events .entry(event.clone()) - .or_insert_with(|| Value::Array(Vec::new())); - existing + .or_insert_with(|| serde_json::json!([])) .as_array_mut() - .with_context(|| format!("Muse managed hook event {event} must be an array"))? + .context("Muse hook event must be an array")? .extend( - handlers + groups .as_array() - .context("generated Muse handlers must be an array")? + .context("generated Muse groups must be an array")? .iter() .cloned(), ); } - result.insert("hooks".into(), hooks.clone().into()); - Ok(Value::Object(result)) + Ok(()) } pub(crate) fn check_muse_project_hooks(start: &Path) -> anyhow::Result<()> { @@ -164,18 +237,6 @@ fn muse_contains_braintrust_hook(value: &Value) -> bool { } } -fn muse_has_strict_handler_fields(value: &Value) -> bool { - match value { - Value::Object(object) => { - object.contains_key("name") - || object.contains_key("timeout_ms") - || object.values().any(muse_has_strict_handler_fields) - } - Value::Array(items) => items.iter().any(muse_has_strict_handler_fields), - _ => false, - } -} - fn legacy_muse_hook_config() -> Map { muse_hook_config_for( LEGACY_MUSE_HOOK_EVENTS, @@ -195,6 +256,20 @@ pub(crate) fn muse_version(output: &[u8]) -> anyhow::Result { }) } +/// Resolve the running executable rather than stamping an install-time version: +/// Muse's launcher can upgrade independently of Braintrust's hook definitions. +pub(crate) fn installed_muse_version() -> anyhow::Result { + let executable = crate::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()); + let output = crate::subprocess::background_command(executable) + .arg("--version") + .output() + .context("check Muse Code version")?; + if !output.status.success() { + bail!("Muse Code could not report its version"); + } + muse_version(&output.stdout) +} + /// Muse serializes changes to settings.json with this native lock file. Hold /// it across the read, hook update, and settings update so we cannot discard a /// concurrent change made by Muse itself. @@ -224,7 +299,20 @@ fn muse_hook_is_owned(path: &Path) -> anyhow::Result { Ok(raw) => { let value: Value = serde_json::from_slice(&raw) .with_context(|| format!("invalid Muse hook configuration: {}", path.display()))?; + let mut previous = muse_hook_config_for(PREVIOUS_MUSE_HOOK_EVENTS, MUSE_HOOK_COMMAND); + let mut legacy = legacy_muse_hook_config(); + for config in [&mut previous, &mut legacy] { + for groups in config["hooks"] + .as_object_mut() + .context("generated Muse hooks")? + .values_mut() + { + groups[0]["matcher"] = Value::from("*"); + } + } Ok(value == Value::Object(muse_hook_config()) + || value == Value::Object(previous) + || value == Value::Object(legacy) || value == Value::Object(legacy_muse_hook_config())) } Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), @@ -234,78 +322,81 @@ fn muse_hook_is_owned(path: &Path) -> anyhow::Result { } } -fn check_muse_hook_ownership(path: &Path) -> anyhow::Result<()> { - if path.exists() && !muse_hook_is_owned(path)? { - bail!("Muse hook file {} differs from Braintrust's generated hooks; refusing to overwrite or remove it", path.display()); - } - Ok(()) -} - fn setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { with_muse_settings_lock(config_dir, || { 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") { - if existing.as_str() != Some(owned_path.as_ref()) { - bail!( - "Muse already uses managed_hooks_path {}; refusing to overwrite it", - existing - ); - } + check_foreign_managed_hooks(config_dir, &settings)?; + let mut hooks = settings + .remove("hooks") + .unwrap_or_else(|| serde_json::json!({})); + remove_owned_hooks(&mut hooks)?; + if muse_contains_braintrust_hook(&hooks) { + bail!("Muse user hooks include an unrecognized Braintrust forwarder; refusing to duplicate it"); } - check_muse_hook_ownership(&hook_path)?; + append_hooks(&mut hooks, MUSE_HOOK_COMMAND)?; + settings.insert("hooks".into(), hooks); settings.entry("schema_version").or_insert(Value::from(1)); - settings.insert( - "managed_hooks_path".into(), - Value::String(owned_path.into_owned()), - ); - write_muse_object(&hook_path, muse_hook_config())?; - write_muse_object(&settings_path, settings) + // Migrate our old dedicated managed file; never edit an enterprise file. + let old_path = managed_hook_path(config_dir, &settings) + .filter(|path| path == &muse_hook_path(config_dir)); + let migrate = old_path + .as_ref() + .map(|path| muse_hook_is_owned(path)) + .transpose()? + .unwrap_or(false); + if migrate { + settings.remove("managed_hooks_path"); + } + write_muse_object(&settings_path, settings)?; + if migrate { + std::fs::remove_file(old_path.context("owned Muse hook path")?)?; + } + Ok(()) }) } fn disable_muse_at(config_dir: &Path) -> anyhow::Result<()> { with_muse_settings_lock(config_dir, || { let settings_path = config_dir.join("settings.json"); - let hook_path = muse_hook_path(config_dir); + if !settings_path.exists() { + return Ok(()); + } let mut settings = load_object(&settings_path)?; - check_muse_hook_ownership(&hook_path)?; - if settings.get("managed_hooks_path").and_then(Value::as_str) - == Some(hook_path.to_string_lossy().as_ref()) - { + if let Some(hooks) = settings.get_mut("hooks") { + remove_owned_hooks(hooks)?; + if hooks.as_object().is_some_and(Map::is_empty) { + settings.remove("hooks"); + } + } + let old_path = managed_hook_path(config_dir, &settings) + .filter(|path| path == &muse_hook_path(config_dir)); + let remove_old = old_path + .as_ref() + .map(|path| muse_hook_is_owned(path)) + .transpose()? + .unwrap_or(false); + if remove_old { settings.remove("managed_hooks_path"); - write_muse_object(&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() - ) - }), + write_muse_object(&settings_path, settings)?; + if remove_old { + std::fs::remove_file(old_path.context("owned Muse hook path")?)?; } + Ok(()) }) } fn update_muse_at(config_dir: &Path) -> anyhow::Result<()> { - with_muse_settings_lock(config_dir, || { - let settings = load_object(&config_dir.join("settings.json"))?; - let hook_path = muse_hook_path(config_dir); - if settings.get("managed_hooks_path").and_then(Value::as_str) - != Some(hook_path.to_string_lossy().as_ref()) - { - bail!("Muse tracing hooks are not installed; run `bt trace enable muse`"); - } - if !hook_path.exists() { - bail!("Muse tracing hook file is missing; run `bt trace enable muse`"); - } - check_muse_hook_ownership(&hook_path)?; - write_muse_object(&hook_path, muse_hook_config()) - }) + let settings = load_object(&config_dir.join("settings.json"))?; + let installed = settings + .get("hooks") + .is_some_and(muse_contains_braintrust_hook) + || muse_has_owned_managed_hooks(config_dir)?; + if !installed { + bail!("Muse tracing hooks are not installed; run `bt trace enable muse`"); + } + setup_muse_at(config_dir) } pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { @@ -315,8 +406,8 @@ pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { .output() { Ok(output) if output.status.success() => match muse_version(&output.stdout) { - Ok(version) if version < semver::Version::new(1, 1, 1) => warnings - .push("Muse Code 1.1.1 or newer is required for verified command hooks".into()), + Ok(version) if version < semver::Version::new(1, 3, 0) => warnings + .push("Muse Code 1.3.0 or newer is required for verified command hooks".into()), Ok(_) => {} Err(error) => warnings.push(error.to_string()), }, @@ -340,37 +431,30 @@ pub(crate) fn muse_doctor_warnings(config_dir: &Path) -> Vec { return warnings; } }; - let hook_path = muse_hook_path(config_dir); - if settings.get("managed_hooks_path").and_then(Value::as_str) - != Some(hook_path.to_string_lossy().as_ref()) - { - warnings.push("Muse managed_hooks_path does not point to Braintrust's hook file; run `bt trace enable muse`".into()); + if let Err(error) = check_foreign_managed_hooks(config_dir, &settings) { + warnings.push(error.to_string()); } - match std::fs::read(&hook_path) { - Ok(raw) => match serde_json::from_slice::(&raw) { - Ok(hooks) if hooks == Value::Object(muse_hook_config()) => {} - Ok(hooks) if hooks == Value::Object(legacy_muse_hook_config()) => warnings.push( - "Muse Braintrust hooks use an older command; run `bt trace update muse`".into(), - ), - Ok(_) => warnings.push("Muse Braintrust hook file differs from expected event definitions; restore it or run `bt trace enable muse`".into()), - Err(error) => warnings.push(format!("Muse Braintrust hook file is invalid JSON: {error}")), - }, - Err(error) => warnings.push(format!("Muse Braintrust hook file could not be read: {error}")), + let current = settings + .get("hooks") + .cloned() + .unwrap_or_else(|| serde_json::json!({})); + let mut remaining = current.clone(); + if remove_owned_hooks(&mut remaining).is_err() { + warnings.push("Muse user hooks are malformed".into()); + } else { + let mut upgraded = remaining; + if append_hooks(&mut upgraded, MUSE_HOOK_COMMAND).is_err() || current != upgraded { + warnings.push("Muse Braintrust hooks are missing or outdated; run `bt trace update muse` or `bt trace enable muse`".into()); + } } warnings } fn setup_muse() -> anyhow::Result<()> { - let output = crate::subprocess::background_command("muse") - .arg("--version") - .output() - .context("Muse Code is not installed; install Muse Code before enabling tracing")?; - if !output.status.success() { - bail!("Muse Code could not report its version; install Muse Code before enabling tracing"); - } - let installed = muse_version(&output.stdout)?; - if installed < semver::Version::new(1, 1, 1) { - bail!("Muse Code 1.1.1 or newer is required for verified command hooks"); + let installed = + installed_muse_version().context("install Muse Code before enabling tracing")?; + if installed < semver::Version::new(1, 3, 0) { + bail!("Muse Code 1.3.0 or newer is required for verified command hooks"); } setup_muse_at(&paths::muse_config_dir()) } @@ -397,207 +481,122 @@ impl Setup for Muse { update_muse() } } - #[cfg(test)] mod tests { use super::*; - #[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()); - assert_eq!(hooks["hooks"].as_object().unwrap().len(), 6); - assert!(hooks["hooks"].get("PreToolUse").is_none()); - assert_eq!( - hooks["hooks"]["PreLLMCall"][0]["hooks"][0]["command"], - "bt trace hook --source muse --flush-on-turn-end" - ); - assert!(config_dir.join(".settings.json.lock").exists()); - - 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_managed_run_replaces_only_owned_braintrust_hooks() { + fn setup_update_and_disable_preserve_enterprise_and_user_hooks() { let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - setup_muse_at(&config_dir).unwrap(); - let saved = std::fs::read(muse_hook_path(&config_dir)).unwrap(); - let generated = - muse_managed_hook_config(&config_dir, "'bt' 'hook' '--source' 'muse'").unwrap(); + let config = temp.path(); + let enterprise = serde_json::json!({"hooks":{"SessionStart":[{"hooks":[{"name":"policy","timeout_ms":1000,"command":"true"}]}]}}); + std::fs::write(config.join("enterprise.json"), enterprise.to_string()).unwrap(); + let foreign = serde_json::json!({"hooks":[{"type":"command","command":"echo unrelated"}]}); + let settings = serde_json::json!({"schema_version":1,"theme":"dark","managed_hooks_path":"enterprise.json","hooks":{"Stop":[foreign.clone()]}}); + std::fs::write(config.join("settings.json"), settings.to_string()).unwrap(); + setup_muse_at(config).unwrap(); + let first = std::fs::read(config.join("settings.json")).unwrap(); + setup_muse_at(config).unwrap(); + update_muse_at(config).unwrap(); + assert_eq!(first, std::fs::read(config.join("settings.json")).unwrap()); + let installed: Value = serde_json::from_slice(&first).unwrap(); + assert_eq!(installed["hooks"].as_object().unwrap().len(), 17); + assert_eq!(installed["hooks"]["Stop"][0], foreign); + // Unmatchable lifecycle events reject a matcher on current Muse. + assert!(installed["hooks"]["Stop"][1].get("matcher").is_none()); + disable_muse_at(config).unwrap(); assert_eq!( - generated["hooks"].as_object().unwrap().len(), - MUSE_HOOK_EVENTS.len() + Value::Object(load_object(&config.join("settings.json")).unwrap()), + settings ); assert_eq!( - generated["hooks"]["SessionStart"][0]["hooks"][0]["command"], - "'bt' 'hook' '--source' 'muse'" + std::fs::read_to_string(config.join("enterprise.json")).unwrap(), + enterprise.to_string() ); - assert_eq!(std::fs::read(muse_hook_path(&config_dir)).unwrap(), saved); } #[test] - fn muse_managed_run_keeps_foreign_handlers_and_rejects_duplicate_user_hooks() { + fn migrate_old_owned_managed_file_and_preserve_unknown_files() { 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.json"}"#, - ) - .unwrap(); - let original = r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"true"}]}]}}"#; - std::fs::write(config_dir.join("other.json"), original).unwrap(); - let generated = muse_managed_hook_config(&config_dir, "bt hook --source muse").unwrap(); - assert_eq!(generated["hooks"]["Stop"].as_array().unwrap().len(), 2); - assert_eq!(generated["hooks"]["Stop"][0]["hooks"][0]["command"], "true"); - assert_eq!( - std::fs::read_to_string(config_dir.join("other.json")).unwrap(), - original - ); - - std::fs::write( - config_dir.join("settings.json"), - r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","command":"bt trace hook --source muse"}]}]}}"#, - ) - .unwrap(); - assert!( - muse_managed_hook_config(&config_dir, "bt hook --source muse") - .unwrap_err() - .to_string() - .contains("user hooks already include Braintrust") - ); - - std::fs::write( - config_dir.join("settings.json"), - r#"{"schema_version":1,"managed_hooks_path":"other.json"}"#, - ) - .unwrap(); - std::fs::write( - config_dir.join("other.json"), - r#"{"schema_version":1,"hooks":{"Stop":[{"hooks":[{"type":"command","name":"foreign","timeout_ms":1000,"command":"true"}]}]}}"#, + let config = temp.path(); + let hook_path = muse_hook_path(config); + let mut old = muse_hook_config_for(PREVIOUS_MUSE_HOOK_EVENTS, MUSE_HOOK_COMMAND); + for groups in old["hooks"].as_object_mut().unwrap().values_mut() { + groups[0]["matcher"] = Value::from("*"); + } + write_muse_object(&hook_path, old).unwrap(); + write_muse_object( + &config.join("settings.json"), + serde_json::json!({"managed_hooks_path":hook_path}) + .as_object() + .unwrap() + .clone(), ) .unwrap(); - assert!( - muse_managed_hook_config(&config_dir, "bt hook --source muse") - .unwrap_err() - .to_string() - .contains("strict handler schema") - ); + assert!(muse_has_owned_managed_hooks(config).unwrap()); + update_muse_at(config).unwrap(); + assert!(!hook_path.exists()); + assert!(load_object(&config.join("settings.json")) + .unwrap() + .get("managed_hooks_path") + .is_none()); + let custom = r#"{"hooks":{"Stop":[]}}"#; + std::fs::write(&hook_path, custom).unwrap(); + setup_muse_at(config).unwrap(); + disable_muse_at(config).unwrap(); + assert_eq!(std::fs::read_to_string(&hook_path).unwrap(), custom); } #[test] - fn muse_managed_run_rejects_duplicate_project_hooks() { + fn invocation_overlay_replaces_owned_hooks_and_preserves_strict_enterprise() { let temp = tempfile::tempdir().unwrap(); - let project = temp.path().join("project"); - let nested = project.join("nested"); - std::fs::create_dir_all(project.join(".muse")).unwrap(); - std::fs::create_dir_all(&nested).unwrap(); - std::fs::write( - project.join(".muse/hooks.json"), - r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"'bt' 'trace' 'hook' '--source' 'muse'"}]}]}}"#, - ) - .unwrap(); - assert!(check_muse_project_hooks(&nested) - .unwrap_err() - .to_string() - .contains("project hooks")); + let config = temp.path(); + let enterprise = r#"{"hooks":{"SessionStart":[{"hooks":[{"name":"policy","timeout_ms":1000,"command":"true"}]}]}}"#; + std::fs::write(config.join("enterprise.json"), enterprise).unwrap(); + std::fs::write(config.join("settings.json"), r#"{"managed_hooks_path":"enterprise.json","hooks":{"Stop":[{"hooks":[{"type":"command","command":"echo unrelated"}]}]}}"#).unwrap(); + setup_muse_at(config).unwrap(); + let saved = std::fs::read(config.join("settings.json")).unwrap(); + let overlay = muse_managed_hook_config(config, "run-forwarder").unwrap(); + assert_eq!(overlay["Stop"].as_array().unwrap().len(), 2); + assert_eq!(overlay["Stop"][0]["hooks"][0]["command"], "echo unrelated"); + assert_eq!(overlay["Stop"][1]["hooks"][0]["command"], "run-forwarder"); + assert!(!overlay.to_string().contains(MUSE_HOOK_COMMAND)); + assert_eq!(std::fs::read(config.join("settings.json")).unwrap(), saved); + assert_eq!( + std::fs::read_to_string(config.join("enterprise.json")).unwrap(), + enterprise + ); } #[test] - fn muse_setup_refuses_another_managed_hook_owner() { + fn reject_unrecognized_braintrust_forwarders_and_project_duplicates() { 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")); + std::fs::write(temp.path().join("settings.json"), r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"custom-bt trace hook --source muse"}]}]}}"#).unwrap(); + assert!(setup_muse_at(temp.path()).is_err()); + assert!(muse_managed_hook_config(temp.path(), "managed").is_err()); + std::fs::create_dir(temp.path().join(".muse")).unwrap(); + std::fs::write(temp.path().join(".muse/hooks.json"), r#"{"hooks":{"Stop":[{"hooks":[{"type":"command","command":"bt trace hook --source muse"}]}]}}"#).unwrap(); + assert!(check_muse_project_hooks(temp.path()).is_err()); } #[test] - fn muse_setup_and_disable_preserve_foreign_hook_file() { + fn disable_before_install_does_not_create_invalid_native_settings() { let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - std::fs::create_dir_all(&config_dir).unwrap(); - let hook_path = muse_hook_path(&config_dir); - let custom = r#"{"schema_version":1,"hooks":{"Stop":[]}}"#; - std::fs::write(&hook_path, custom).unwrap(); - - assert!(setup_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - assert!(disable_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - assert_eq!(std::fs::read_to_string(&hook_path).unwrap(), custom); + disable_muse_at(temp.path()).unwrap(); + assert!(!temp.path().join("settings.json").exists()); } #[test] - fn muse_update_requires_an_owned_install_and_upgrades_legacy_command() { + fn version_and_update_diagnostics() { + assert_eq!( + muse_version(b"Muse Code 1.4.10 (1.4.10-R6427.1)\n").unwrap(), + semver::Version::new(1, 4, 10) + ); + assert!(muse_version(b"unexpected output").is_err()); let temp = tempfile::tempdir().unwrap(); - let config_dir = temp.path().join("muse"); - assert!(update_muse_at(&config_dir) + assert!(update_muse_at(temp.path()) .unwrap_err() .to_string() .contains("not installed")); - - setup_muse_at(&config_dir).unwrap(); - let hook_path = muse_hook_path(&config_dir); - write_muse_object(&hook_path, legacy_muse_hook_config()).unwrap(); - assert!(muse_doctor_warnings(&config_dir) - .iter() - .any(|warning| warning.contains("older command"))); - update_muse_at(&config_dir).unwrap(); - let installed = load_object(&hook_path).unwrap(); - assert_eq!(installed, muse_hook_config()); - - std::fs::write(&hook_path, r#"{"hooks":{}}"#).unwrap(); - assert!(update_muse_at(&config_dir) - .unwrap_err() - .to_string() - .contains("refusing to overwrite")); - } - - #[test] - fn muse_version_requires_a_parseable_native_version() { - assert_eq!( - muse_version(b"Muse Code 1.1.1 (1.1.1-R2514.1)\n").unwrap(), - semver::Version::new(1, 1, 1) - ); - assert!(muse_version(b"unexpected output").is_err()); } } diff --git a/bt-daemon/src/span_processor.rs b/bt-daemon/src/span_processor.rs index f706496e..4a2e07fa 100644 --- a/bt-daemon/src/span_processor.rs +++ b/bt-daemon/src/span_processor.rs @@ -311,6 +311,7 @@ fn process_inner( candidate.origin = row.origin.clone(); candidate.turn_span_id = row.turn_span_id.clone(); candidate.late_merge_key = row.late_merge_key.clone(); + candidate.late_merge_version = row.late_merge_version.clone(); row = candidate; } Ok(()) diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs index 5c91246d..055d52fb 100644 --- a/bt-daemon/src/transcript_import/muse.rs +++ b/bt-daemon/src/transcript_import/muse.rs @@ -67,6 +67,7 @@ struct NativePayload { kind: Option, run_id: Option, event: Option, + record: Option, } #[derive(Default)] @@ -74,12 +75,15 @@ struct RunState { active_model: Option, last_assistant: Option, model_ordinal: u64, + step_attempts: HashMap, } struct ModelState { request_id: String, usage: Option, text: Option, + response_id: Option, + completed_at: Option, } struct Reader<'a, F> { @@ -94,6 +98,11 @@ struct Reader<'a, F> { last_ts: Option, current_session_ended: bool, activity_count: usize, + metadata: Value, + tools: HashMap, + tool_errors: HashMap, + tool_tasks: HashMap, + run_metadata: HashMap, } impl Reader<'_, F> @@ -122,6 +131,9 @@ where "hook_event_name": "SessionStart", "source": "muse_export_v1", "trajectory_id": self.session.trajectory_id, + "model": self.metadata.get("model_id"), + "provider": self.metadata.get("provider_id"), + "cwd": self.metadata.get("workspace_root"), "root_session_id": self.session.root_session_id, "transcript_path": self.path.to_string_lossy(), }), @@ -136,16 +148,30 @@ where run_id: &str, ts: i64, request_id: Option<&str>, - messages: Value, + input: Value, + model_step: Option, ) -> Result<()> { + let model_metadata = self + .run_metadata + .get(run_id) + .unwrap_or(&self.metadata) + .clone(); let run = self.runs.entry(run_id.to_owned()).or_default(); if run.active_model.is_some() { bail!("Muse run {run_id} starts another model request before the previous response") } - let request_id = match request_id { - Some(id) if !id.is_empty() => id.to_owned(), - Some(_) => bail!("Muse model_input_trace_recorded has an empty request_record_id"), - None => { + let native_request_id = request_id.map(str::to_owned); + let request_id = match (model_step, request_id) { + (Some(step), _) => { + let attempt = run.step_attempts.entry(step).or_default(); + *attempt += 1; + format!("{run_id}:{step}:{attempt}") + } + (None, Some(id)) if !id.is_empty() => id.to_owned(), + (None, Some(_)) => { + bail!("Muse model_input_trace_recorded has an empty request_record_id") + } + (None, None) => { run.model_ordinal += 1; format!("{run_id}:model:{}", run.model_ordinal) } @@ -154,6 +180,8 @@ where request_id: request_id.clone(), usage: None, text: None, + response_id: None, + completed_at: None, }); self.emit( "PreLLMCall", @@ -163,7 +191,15 @@ where "hook_event_name": "PreLLMCall", "turn_id": run_id, "request_id": request_id, - "messages": messages, + "request_record_id": native_request_id, + "step": model_step, + "model": model_metadata.get("model_id"), + "provider": model_metadata.get("provider_id"), + "cwd": self.metadata.get("workspace_root"), + "messages": input.get("messages"), + "tools": input.get("tools"), + "options": input.get("options"), + "input_summary": input.get("bounded"), "source": "muse_export_v1", }), ) @@ -197,15 +233,17 @@ where }; self.emit( "PostLLMCall", - ts, + model.completed_at.unwrap_or(ts), json!({ "session_id": self.session_id, "hook_event_name": "PostLLMCall", "turn_id": run_id, "request_id": model.request_id, "usage": model.usage, + "usage_scope": "model", "output_text_preview": model.text, - "error": error, + "response_id": model.response_id, + "error": if model.completed_at.is_some() { None } else { error }, "source": "muse_export_v1", }), ) @@ -218,6 +256,74 @@ where let Some(payload) = native_envelope.payload else { return Ok(()); }; + if payload.kind.as_deref() == Some("run_model") { + if let Some(record) = payload.record { + if let Some(run) = record.pointer("/run_stream/id").and_then(Value::as_str) { + self.run_metadata.insert(run.to_owned(), record); + } + } + return Ok(()); + } + if payload.kind.as_deref() == Some("task") { + if let Some(event) = &payload.event { + let kind = event.get("kind").and_then(Value::as_str); + let task = event.get("task_id").and_then(Value::as_str); + if let (Some(task), Some(call)) = ( + task, + event + .get("idempotency_key") + .and_then(Value::as_str) + .and_then(|key| key.strip_prefix("tool:")), + ) { + self.tool_tasks.insert(task.to_owned(), call.to_owned()); + } + if kind == Some("rejected") { + if let (Some(id), Some(reason)) = ( + event + .pointer("/origin/tool_call_id") + .and_then(Value::as_str), + event.get("reason").and_then(Value::as_str), + ) { + self.tool_errors.insert(id.to_owned(), reason.to_owned()); + } + } + if let Some(call) = task.and_then(|task| self.tool_tasks.get(task)).cloned() { + if kind == Some("started") { + if let Some(tool) = self.tools.get(&call).cloned() { + let ts = item + .recorded_at + .or(native_envelope.recorded_at) + .context("Muse tool task has no recorded_at")? + / 1000; + self.emit("PreToolUse", ts, tool)?; + } + } else if matches!(kind, Some("failed" | "cancelled" | "canceled")) { + let reason = event + .get("reason") + .and_then(Value::as_str) + .unwrap_or("Muse tool task failed"); + self.tool_errors.insert(call, reason.to_owned()); + if let Some(task) = task { + self.tool_tasks.remove(task); + } + } else if kind == Some("completed") { + if let Some(tool) = self.tools.get_mut(&call) { + tool["status"] = json!("completed"); + } + if let Some(task) = task { + self.tool_tasks.remove(task); + } + } + } + } + return Ok(()); + } + if payload.kind.as_deref() == Some("metadata") { + if let Some(record) = payload.record { + self.metadata = record; + } + return Ok(()); + } let kind = payload .event .as_ref() @@ -234,6 +340,10 @@ where "started" | "model_input_trace_recorded" | "model_completed" + | "model_response_created" + | "assistant_tool_calls_committed" + | "tool_result_batch_committed" + | "model_failure_terminal_intent" | "resource_usage_sampled" | "assistant_message_committed" | "terminal" @@ -305,9 +415,88 @@ where &run_id, ts, Some(request_id), - event.get("bounded").cloned().unwrap_or(Value::Null), + json!({"bounded":event.get("bounded")}), + event.get("model_step").and_then(Value::as_u64), )?; } + "assistant_tool_calls_committed" => { + let calls = event + .get("tool_calls") + .and_then(Value::as_array) + .context("Muse assistant_tool_calls_committed has no tool_calls array")?; + let mut output_calls = Vec::new(); + for call in calls { + let id = call + .get("call_id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .context("Muse committed tool call has no call_id")?; + let name = call + .get("name") + .and_then(Value::as_str) + .context("Muse committed tool call has no name")?; + let args = call.get("args").cloned().unwrap_or(Value::Null); + let input = if let Some(args) = args.as_str() { + serde_json::from_str(args).unwrap_or_else(|_| args.into()) + } else { + args.clone() + }; + let payload = json!({"turn_id":run_id,"tool_use_id":id,"tool_name":name,"tool_input":input,"source":"muse_export_v1"}); + self.tools.insert(id.to_owned(), payload.clone()); + + output_calls.push(json!({"id":id,"type":"function","function":{"name":name,"arguments":args}})); + } + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|run| run.active_model.as_mut()) + { + model.text = Some(json!([{"role":"assistant","tool_calls":output_calls}])); + model.response_id = event.get("response_id").cloned(); + } + } + "tool_result_batch_committed" => { + let results = event + .get("results") + .and_then(Value::as_array) + .context("Muse tool_result_batch_committed has no results array")?; + for result in results { + let id = result + .get("tool_call_id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .context("Muse tool result has no tool_call_id")?; + let mut payload = self.tools.remove(id).unwrap_or_else( + || json!({"turn_id":run_id,"tool_use_id":id,"source":"muse_export_v1"}), + ); + payload["tool_response"] = result.get("text").cloned().unwrap_or(Value::Null); + let error = self.tool_errors.remove(id); + payload["error"] = json!(error); + if error.is_some() { + self.emit("PostToolUseFailure", ts, payload)?; + } else if payload.get("status").and_then(Value::as_str) == Some("completed") { + self.emit("PostToolUse", ts, payload)?; + } else { + self.emit("PostToolBatch", ts, json!({"turn_id":run_id,"tool_calls":[payload],"source":"muse_export_v1"}))?; + } + } + } + "model_failure_terminal_intent" => { + let reason = event + .get("reason") + .and_then(Value::as_str) + .unwrap_or("Muse model failed"); + self.finish_model(&run_id, ts, Some(reason))?; + } + "model_response_created" => { + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|run| run.active_model.as_mut()) + { + model.response_id = event.get("response_id").cloned(); + } + } "model_completed" => { if self .runs @@ -315,7 +504,7 @@ where .and_then(|run| run.active_model.as_ref()) .is_none() { - self.start_model(&run_id, ts, None, Value::Null)?; + self.start_model(&run_id, ts, None, Value::Null, None)?; } if let Some(model) = self .runs @@ -323,6 +512,7 @@ where .and_then(|run| run.active_model.as_mut()) { model.usage = event.get("usage").cloned(); + model.completed_at = Some(ts); } } "resource_usage_sampled" => { @@ -347,7 +537,7 @@ where .and_then(|run| run.active_model.as_ref()) .is_none() { - self.start_model(&run_id, ts, None, Value::Null)?; + self.start_model(&run_id, ts, None, Value::Null, None)?; } let run = self .runs @@ -358,6 +548,9 @@ where .as_mut() .context("Muse run has no active model")? .text = Some(text); + if let Some(model) = run.active_model.as_mut() { + model.response_id = event.get("response_id").cloned(); + } } "terminal" => { let terminal = event @@ -378,6 +571,7 @@ where }); self.finish_model(&run_id, ts, error.as_deref())?; let run = self.runs.remove(&run_id); + self.run_metadata.remove(&run_id); self.emit( "Stop", ts, @@ -524,6 +718,11 @@ where last_ts: None, current_session_ended: false, activity_count: 0, + metadata: json!({}), + tools: HashMap::new(), + tool_errors: HashMap::new(), + tool_tasks: HashMap::new(), + run_metadata: HashMap::new(), }; let mut de = serde_json::Deserializer::from_reader(BufReader::new( File::open(path).with_context(|| format!("read Muse export {}", path.display()))?, @@ -691,11 +890,125 @@ mod tests { "SessionEnd" ] ); + assert_eq!(events[3].ts_ms, 1002); assert_eq!(events[3].payload["usage"]["input_tokens"], 1); assert_eq!(events[3].payload["output_text_preview"], "world"); assert_eq!(events[5].ts_ms, 1009); } + #[test] + fn native_model_steps_match_live_hook_identity_and_distinguish_retries() { + let (_temp, path) = write_export(json!({ + "export_schema_version":1, + "sessions":[{"session_id":"session-1"}], + "events":[ + {"recorded_at":999000,"envelope":{"payload":{"kind":"metadata","record":{"model_id":"native-model","provider_id":"native-provider","workspace_root":"/workspace"}}}}, + {"recorded_at":1000000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"hello"}}}}, + {"recorded_at":1001000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","model_step":0,"request_record_id":"uuid-a","bounded":{"request_message_count":2}}}}}, + {"recorded_at":1002000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","model_step":0,"request_record_id":"uuid-b"}}}}, + {"recorded_at":1003000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"hello","response_id":"provider-response"}}}}, + {"recorded_at":1004000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}} + ] + })); + let events = envelopes(&path).unwrap(); + let starts: Vec<_> = events.iter().filter(|e| e.event == "PreLLMCall").collect(); + assert_eq!(starts[0].payload["model"], "native-model"); + assert_eq!(starts[0].payload["provider"], "native-provider"); + assert_eq!(events[0].payload["cwd"], "/workspace"); + assert_eq!(starts[0].payload["request_id"], "run-1:0:1"); + assert_eq!(starts[1].payload["request_id"], "run-1:0:2"); + assert_eq!(starts[0].payload["request_record_id"], "uuid-a"); + assert!(starts[0].payload["messages"].is_null()); + assert_eq!( + starts[0].payload["input_summary"]["request_message_count"], + 2 + ); + let post = events.iter().rfind(|e| e.event == "PostLLMCall").unwrap(); + assert_eq!(post.payload["response_id"], "provider-response"); + } + + #[test] + fn captured_modern_export_imports_native_tools_models_and_usage() { + // Reduced native 1.4.11 export captured against a deterministic Responses + // endpoint. Session/turn/task ids, paths and timestamps are sanitized. + let (_temp, path) = write_export( + serde_json::from_str(include_str!("../../tests/fixtures/muse-export-1.4.11.json")) + .unwrap(), + ); + let events = envelopes(&path).unwrap(); + let pre = events.iter().find(|e| e.event == "PreToolUse").unwrap(); + let post = events.iter().find(|e| e.event == "PostToolUse").unwrap(); + assert_eq!(pre.payload["tool_use_id"], "call_muse_1"); + assert_eq!(pre.payload["tool_input"]["command"], "printf MUSE_TOOL_OK"); + assert_eq!(post.payload["tool_response"], "MUSE_TOOL_OK"); + assert!(pre.ts_ms < post.ts_ms); + let registry = crate::translate::Registry::default_agents(); + let mut translator = registry.create("muse", "muse-export-fixture"); + let ctx = crate::translate::SessionCtx { + session_id: "muse-export-fixture".into(), + config: None, + }; + let mut rows = Vec::new(); + for event in &events { + rows.extend(translator.handle(event, &ctx).unwrap()); + } + assert_eq!(rows.iter().filter(|op|matches!(op,crate::translate::SpanOp::Insert(row) if row.span_type == crate::translate::SpanType::Tool)).count(),1); + assert_eq!(rows.iter().filter(|op|matches!(op,crate::translate::SpanOp::Insert(row) if row.span_type == crate::translate::SpanType::Llm)).count(),2); + let root = crate::ids::span_id( + &crate::ids::session_namespace("muse", "muse-export-fixture"), + "root", + ); + assert!(rows.iter().any(|op|matches!(op,crate::translate::SpanOp::Merge(row) if row.span_id == root && row.metrics.as_ref().is_some_and(|m|m["prompt_tokens"] == 20 && m["completion_tokens"] == 10 && m["tokens"] == 30)))); + assert!(rows.iter().any(|op|matches!(op,crate::translate::SpanOp::Merge(row) if row.output == Some(json!("MUSE_TOOL_OK")) && row.metadata.as_ref().is_some_and(|m|m["tool_outcome"] == "success")))); + } + + #[test] + fn captured_failed_tool_task_preserves_reason_and_output() { + let (_temp, path) = write_export( + serde_json::from_str(include_str!( + "../../tests/fixtures/muse-export-tool-failure-1.4.11.json" + )) + .unwrap(), + ); + let events = envelopes(&path).unwrap(); + let failed = events + .iter() + .find(|e| e.event == "PostToolUseFailure") + .expect("native failed task has a failure envelope"); + assert_eq!(failed.payload["tool_use_id"], "call_muse_failure"); + assert_eq!( + failed.payload["error"], + "exit status: 7\nstdout:\nMUSE_TOOL_FAILURE" + ); + assert_eq!( + failed.payload["tool_response"], + "tool failed: exit status: 7\nstdout:\nMUSE_TOOL_FAILURE" + ); + } + + #[test] + fn rejected_tool_and_model_failures_keep_explicit_native_reasons() { + let (_temp, path) = write_export( + json!({"export_schema_version":1,"sessions":[{"session_id":"failure-fixture"}],"events":[ + {"recorded_at":1000000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"started","prompt":"hello"}}}}, + {"recorded_at":1001000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"model_input_trace_recorded","model_step":0,"request_record_id":"uuid"}}}}, + {"recorded_at":1002000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"model_failure_terminal_intent","reason":"provider connection failed"}}}}, + {"recorded_at":1003000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"assistant_tool_calls_committed","tool_calls":[{"call_id":"call","name":"unknown","args":"{}"}]}}}}, + {"recorded_at":1004000,"envelope":{"payload":{"kind":"task","run_id":"turn","event":{"kind":"rejected","origin":{"tool_call_id":"call"},"reason":"unknown tool `unknown`"}}}}, + {"recorded_at":1005000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"tool_result_batch_committed","results":[{"tool_call_id":"call","text":"unavailable"}]}}}}, + {"recorded_at":1006000,"envelope":{"payload":{"kind":"run","run_id":"turn","event":{"kind":"terminal","terminal":"failed","reason":"provider connection failed"}}}} + ]}), + ); + let events = envelopes(&path).unwrap(); + assert!(events.iter().any( + |e| e.event == "PostLLMCall" && e.payload["error"] == "provider connection failed" + )); + assert!(events + .iter() + .any(|e| e.event == "PostToolUseFailure" + && e.payload["error"] == "unknown tool `unknown`")); + } + #[test] fn sparse_export_has_deterministic_model_and_failure() { let (_temp, path) = write_export(json!({ diff --git a/bt-daemon/src/transcript_import/muse_runtime.rs b/bt-daemon/src/transcript_import/muse_runtime.rs index f460a544..8c31a8d8 100644 --- a/bt-daemon/src/transcript_import/muse_runtime.rs +++ b/bt-daemon/src/transcript_import/muse_runtime.rs @@ -409,7 +409,13 @@ pub(super) async fn attach_muse_export( tokio::pin!(shutdown); loop { if reexport { - export_muse_session(session_id, export).await?; + tokio::select! { + result = &mut shutdown => { + result?; + break; + } + result = export_muse_session(session_id, export) => result?, + } } let (count, ended) = process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; @@ -428,13 +434,14 @@ pub(super) async fn attach_muse_export( _ = tokio::time::sleep(Duration::from_secs(1)) => {} } } - // Capture one final complete snapshot before interrupting any still-open - // spans, matching the finalization behavior of other attach sources. - if reexport { - export_muse_session(session_id, export).await?; - } - process_muse_snapshot(&mut processor, export, session_id, true, emitted).await?; - processor.finish().await + // Detaching stops observation, not the native session. Leave open spans + // unfinished so a later export can deliver their authoritative completion. + // Do not launch another potentially slow export after cancellation. + Ok(processor + .detach_session(session_id) + .await? + .into_iter() + .collect()) } #[cfg(test)] diff --git a/bt-daemon/src/transcript_import/pipeline.rs b/bt-daemon/src/transcript_import/pipeline.rs index b7725bcb..04dd6d4c 100644 --- a/bt-daemon/src/transcript_import/pipeline.rs +++ b/bt-daemon/src/transcript_import/pipeline.rs @@ -652,7 +652,8 @@ mod tests { .filter(|row| row .pointer("/Merge/span_id") .and_then(|value| value.as_str()) - == Some(id.as_str())) + == Some(id.as_str()) + && row.pointer("/Merge/end_ms").is_some()) .count(), 1, "missing or duplicate completion for {identity}" @@ -766,7 +767,8 @@ mod tests { .filter(|row| row .pointer("/Merge/span_id") .and_then(|value| value.as_str()) - == Some(root_id.as_str())) + == Some(root_id.as_str()) + && row.pointer("/Merge/end_ms").is_some()) .count(), 1 ); @@ -776,6 +778,226 @@ mod tests { assert_eq!(std::fs::read_to_string(&rows_path).unwrap(), completed); } + #[tokio::test] + async fn muse_live_capture_then_export_reuses_model_and_late_updates_resumed_root() { + let temp = tempfile::tempdir().unwrap(); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let export = temp.path().join("export.json"); + let make_processor = || { + ImportProcessor::new( + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + Some(ledger.clone()), + ) + .unwrap() + }; + let mut live = make_processor(); + for (name, payload, ts) in [ + ( + "UserPromptSubmit", + json!({"turn_id":"run-1","prompt":"first"}), + 1000, + ), + ( + "PreLLMCall", + json!({"turn_id":"run-1","request_id":"run-1:0:1"}), + 1001, + ), + ( + "PostLLMCall", + json!({"turn_id":"run-1","request_id":"run-1:0:1","output_text_preview":"preview"}), + 1003, + ), + ( + "Stop", + json!({"turn_id":"run-1","last_assistant_message":"first answer"}), + 1003, + ), + ("SessionEnd", json!({}), 1004), + ( + "UserPromptSubmit", + json!({"turn_id":"run-2","prompt":"second"}), + 2400, + ), + ( + "Stop", + json!({"turn_id":"run-2","last_assistant_message":"latest preview"}), + 2500, + ), + ("SessionEnd", json!({}), 2501), + ] { + live.process(vec![transcript_import::envelope( + "muse", + None, + "cross-ingress", + name, + ts, + payload, + )]) + .await + .unwrap(); + } + live.finish().await.unwrap(); + std::fs::write(&export, serde_json::to_vec(&json!({ + "export_schema_version":1,"sessions":[{"session_id":"cross-ingress"}], + "events":[ + {"recorded_at":1000000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"started","prompt":"first"}}}}, + {"recorded_at":1001000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"model_input_trace_recorded","request_record_id":"uuid-1","model_step":0}}}}, + {"recorded_at":1002000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"assistant_message_committed","text":"full exported answer"}}}}, + {"recorded_at":1003000,"envelope":{"payload":{"kind":"run","run_id":"run-1","event":{"kind":"terminal","terminal":"completed"}}}}, + {"recorded_at":1004000,"envelope":{"payload":{"kind":"session_end"}}}, + {"recorded_at":2000000,"envelope":{"payload":{"kind":"run","run_id":"run-2","event":{"kind":"started","prompt":"second"}}}}, + {"recorded_at":2001000,"envelope":{"payload":{"kind":"run","run_id":"run-2","event":{"kind":"assistant_message_committed","text":"latest answer"}}}}, + {"recorded_at":2002000,"envelope":{"payload":{"kind":"run","run_id":"run-2","event":{"kind":"terminal","terminal":"completed"}}}}, + {"recorded_at":2003000,"envelope":{"payload":{"kind":"session_end"}}} + ] + })).unwrap()).unwrap(); + let mut imported = make_processor(); + assert!(process_muse_export(&mut imported, &export, "cross-ingress") + .await + .unwrap()); + imported.finish().await.unwrap(); + let rows: Vec = + std::fs::read_to_string(output.join("spans/cross-ingress.ndjson")) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + let namespace = ids::session_namespace("muse", "cross-ingress"); + let model = ids::span_id(&namespace, "llm:run-1:0:1"); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(serde_json::Value::as_str) + == Some(model.as_str())) + .count(), + 1 + ); + assert!(rows + .iter() + .filter_map(|row| row.get("Merge")) + .any(|row| row["span_id"] == model && row["output"] == "full exported answer")); + let root = ids::span_id(&namespace, "root"); + let root_updates: Vec<_> = rows + .iter() + .filter_map(|row| row.get("Merge")) + .filter(|row| row["span_id"] == root) + .collect(); + assert!(root_updates + .iter() + .any(|row| row["output"] == "latest answer")); + assert_eq!(root_updates.last().unwrap()["end_ms"], 2003); + assert_eq!( + root_updates + .iter() + .rev() + .find(|row| row.get("output").is_some()) + .unwrap()["output"], + "latest answer" + ); + // An older complete export must not roll the root back to run-1. + let rows_path = output.join("spans/cross-ingress.ndjson"); + let before_stale = std::fs::read_to_string(&rows_path).unwrap(); + let mut stale: serde_json::Value = + serde_json::from_slice(&std::fs::read(&export).unwrap()).unwrap(); + stale["events"] + .as_array_mut() + .unwrap() + .retain(|event| event["recorded_at"].as_i64().unwrap() < 2_000_000); + std::fs::write(&export, serde_json::to_vec(&stale).unwrap()).unwrap(); + let mut old = make_processor(); + process_muse_export(&mut old, &export, "cross-ingress") + .await + .unwrap(); + old.finish().await.unwrap(); + assert_eq!(std::fs::read_to_string(rows_path).unwrap(), before_stale); + } + + #[tokio::test] + async fn muse_native_tool_export_reuses_delivered_live_tool_identity() { + let temp = tempfile::tempdir().unwrap(); + let output = temp.path().join("output"); + let ledger = temp.path().join("ledger"); + let make_processor = || { + ImportProcessor::new( + debug_serve_options("test", &output), + Some(import_test_config("project-a")), + Some(ledger.clone()), + ) + .unwrap() + }; + let mut live = make_processor(); + for (name, payload, ts) in [ + ( + "UserPromptSubmit", + json!({"turn_id":"fixture-turn","prompt":"Print an integration marker."}), + 1, + ), + ( + "PreToolUse", + json!({"turn_id":"fixture-turn","tool_use_id":"call_muse_1","tool_name":"shell","tool_input":{"command":"printf MUSE_TOOL_OK"}}), + 2, + ), + ( + "PostToolUse", + json!({"turn_id":"fixture-turn","tool_use_id":"call_muse_1","tool_name":"shell","tool_response":"MUSE_TOOL_OK"}), + 3, + ), + ( + "Stop", + json!({"turn_id":"fixture-turn","last_assistant_message":"MUSE_TOOL_OK"}), + 4, + ), + ("SessionEnd", json!({}), 5), + ] { + live.process(vec![transcript_import::envelope( + "muse", + Some("1.4.11".into()), + "muse-export-fixture", + name, + ts, + payload, + )]) + .await + .unwrap(); + } + live.finish().await.unwrap(); + let export = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/muse-export-1.4.11.json"); + let mut imported = make_processor(); + assert!( + process_muse_export(&mut imported, &export, "muse-export-fixture") + .await + .unwrap() + ); + imported.finish().await.unwrap(); + let rows: Vec = + std::fs::read_to_string(output.join("spans/muse-export-fixture.ndjson")) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect(); + let tool = ids::span_id( + &ids::session_namespace("muse", "muse-export-fixture"), + "tool:call_muse_1", + ); + assert_eq!( + rows.iter() + .filter(|row| row + .pointer("/Insert/span_id") + .and_then(serde_json::Value::as_str) + == Some(tool.as_str())) + .count(), + 1 + ); + assert!(rows + .iter() + .filter_map(|row| row.get("Merge")) + .any(|row| row["span_id"] == tool && row["output"] == "MUSE_TOOL_OK")); + } + #[derive(Debug, Parser)] struct ImportCli { #[command(flatten)] diff --git a/bt-daemon/src/translate/debug.rs b/bt-daemon/src/translate/debug.rs index 8d387afb..6af1a0b5 100644 --- a/bt-daemon/src/translate/debug.rs +++ b/bt-daemon/src/translate/debug.rs @@ -62,6 +62,7 @@ impl AgentTranslator for DebugTranslator { metrics: None, error: None, late_merge_key: None, + late_merge_version: None, turn_span_id: None, tags: None, is_turn: false, @@ -86,6 +87,7 @@ impl AgentTranslator for DebugTranslator { metrics: None, error: None, late_merge_key: None, + late_merge_version: None, turn_span_id: None, tags: None, is_turn: false, diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index d39e9a28..30f306bd 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -61,6 +61,14 @@ pub struct OriginSnapshot { pub source_version: Option, } +/// Causal watermark for an aggregate updated by independently replayed snapshots. +/// The same native operation may enrich a live preview despite capture clock lag. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct LateMergeVersion { + pub ts_ms: i64, + pub native_id: Option, +} + /// A resolved span row, ready for a sink to insert or merge. Field set is the /// subset every current plugin uses; extend as translators need more. #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -91,6 +99,9 @@ pub struct SpanRow { #[serde(skip)] #[doc(hidden)] pub late_merge_key: Option, + #[serde(skip)] + #[doc(hidden)] + pub late_merge_version: Option, /// Span id of the user turn that owns this row, recorded by the translator /// from its native turn tracking. [`turn_lineage`] publishes it as /// `metadata.turn_span_id`. diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs index a958b6b9..cc63e5c8 100644 --- a/bt-daemon/src/translate/muse.rs +++ b/bt-daemon/src/translate/muse.rs @@ -5,15 +5,18 @@ //! same event names through this translator. use super::git::GitMetadataCache; +use super::recent::{RecentMap, RecentSet}; + +#[path = "muse_events.rs"] +mod muse_events; use super::{ - local_username, root_tags, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, - TranslatorFactory, + local_username, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, }; use crate::ids; use crate::wire::Envelope; use serde::{Deserialize, Deserializer}; use serde_json::{json, Map, Value}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; impl TranslatorFactory for crate::agents::Muse { @@ -66,8 +69,14 @@ struct MuseTranslator { session_version: Option, turns: HashMap, llms: HashMap, - closed_turns: HashSet, - closed_llms: HashSet, + closed_turns: RecentSet, + closed_llms: RecentSet, + root_created: bool, + events: muse_events::Events, + total_metrics: Map, + model_metrics: RecentMap, + turn_usage: RecentMap, + last_completed_turn: Option, last_ts: i64, git: Arc, } @@ -84,8 +93,14 @@ impl MuseTranslator { session_version: None, turns: HashMap::new(), llms: HashMap::new(), - closed_turns: HashSet::new(), - closed_llms: HashSet::new(), + closed_turns: RecentSet::default(), + closed_llms: RecentSet::default(), + root_created: false, + events: muse_events::Events::default(), + total_metrics: Map::new(), + model_metrics: RecentMap::default(), + turn_usage: RecentMap::default(), + last_completed_turn: None, last_ts: 0, git, } @@ -97,6 +112,7 @@ impl MuseTranslator { for key in [ "cwd", "model", + "provider", "permission_mode", "trajectory_id", "root_session_id", @@ -108,7 +124,18 @@ impl MuseTranslator { } fn starts_trace(event: &str) -> bool { - matches!(event, "UserPromptSubmit" | "PreLLMCall") + matches!( + event, + "UserPromptSubmit" + | "PreLLMCall" + | "PreToolUse" + | "PermissionRequest" + | "PostToolUse" + | "PostToolUseFailure" + | "PostToolBatch" + | "SubagentStart" + | "PreCompact" + ) } fn ensure_root(&mut self, e: &Envelope, ctx: &SessionCtx) -> Vec { @@ -145,6 +172,10 @@ impl MuseTranslator { for (key, value) in &self.session_metadata { metadata.insert(key.clone(), value.clone()); } + if self.root_created { + return Vec::new(); + } + self.root_created = true; vec![SpanOp::Insert(SpanRow { span_id: self.root_id.clone(), root_span_id: self.effective_root_id.clone(), @@ -153,7 +184,6 @@ impl MuseTranslator { span_type: SpanType::Task, start_ms: Some(e.ts_ms), metadata: Some(Value::Object(metadata)), - tags: root_tags(ctx), ..Default::default() })] } @@ -178,11 +208,12 @@ impl MuseTranslator { output, metrics, error, + late_merge_key: Some(format!("muse:complete:{key}:{ts}")), ..Default::default() })] } fn close_all(&mut self, ts: i64, error: Option<&str>) -> Vec { - let mut out = Vec::new(); + let mut out = self.close_events(ts, error); for (map, closed) in [ (&mut self.llms, &mut self.closed_llms), (&mut self.turns, &mut self.closed_turns), @@ -256,11 +287,13 @@ impl AgentTranslator for MuseTranslator { } }; if (e.event == "UserPromptSubmit" - && ids.turn().is_some_and(|id| self.closed_turns.contains(id))) + && ids + .turn() + .is_some_and(|id| self.closed_turns.contains(&id.to_owned()))) || (e.event == "PreLLMCall" && ids .request() - .is_some_and(|id| self.closed_llms.contains(id))) + .is_some_and(|id| self.closed_llms.contains(&id.to_owned()))) { return Ok(Vec::new()); } @@ -306,6 +339,9 @@ impl AgentTranslator for MuseTranslator { "attempt", "message_count", "tool_count", + "options", + "request_record_id", + "input_summary", ] { if let Some(value) = p.get(key) { metadata.insert(key.into(), value.clone()); @@ -341,15 +377,74 @@ impl AgentTranslator for MuseTranslator { "PostLLMCall" => { if let Some(id) = ids.request() { let error = event_error(p); - let closed = Self::close( + let metrics = p.get("usage").filter(|v| v.is_object()).map(|usage| { + // Native hook usage is cumulative within a turn (verified + // against per-model durable records in Muse 1.4.11). + // Export envelopes explicitly carry per-model usage. + let mut per_model = usage.clone(); + if p.get("usage_scope").and_then(Value::as_str) != Some("model") + && self.llms.contains_key(id) + { + if let Some(turn) = ids.turn() { + if let Some(previous) = self.turn_usage.get(turn) { + if let Some(fields) = per_model.as_object_mut() { + for (key, value) in fields { + if let (Some(current), Some(previous)) = ( + value.as_u64(), + previous.get(key).and_then(Value::as_u64), + ) { + // A reset starts a new native counter window. + *value = json!(current + .checked_sub(previous) + .unwrap_or(current)); + } + } + } + } + self.turn_usage.insert(turn.to_owned(), usage.clone()); + } + } + usage_metrics(&per_model) + }); + if let Some(metrics) = &metrics { + if self.llms.contains_key(id) && self.model_metrics.get(id).is_none() { + for (key, value) in metrics.as_object().into_iter().flatten() { + if let Some(value) = value.as_u64() { + let previous = self + .total_metrics + .get(key) + .and_then(Value::as_u64) + .unwrap_or(0); + self.total_metrics + .insert(key.clone(), json!(previous.saturating_add(value))); + } + } + self.model_metrics.insert(id.to_owned(), metrics.clone()); + } + } + let mut closed = Self::close( &self.effective_root_id, &mut self.llms, id, e.ts_ms, p.get("output_text_preview").cloned(), - p.get("usage").map(usage_metrics), + metrics, error, ); + for op in &mut closed { + let (SpanOp::Merge(row) | SpanOp::Insert(row)) = op; + let metadata: Map = [ + "finish_reason", + "response_id", + "tool_call_count", + "status", + "options", + ] + .into_iter() + .filter_map(|key| p.get(key).map(|v| (key.to_owned(), v.clone()))) + .collect(); + row.metadata = Some(Value::Object(metadata)); + } if !closed.is_empty() { self.closed_llms.insert(id.into()); out.extend(closed); @@ -358,7 +453,7 @@ impl AgentTranslator for MuseTranslator { out.push(self.unpaired(e, "Muse PostLLMCall has no request_id")); } } - "Stop" => { + "Stop" | "StopFailure" => { if let Some(id) = ids.turn() { let closed = Self::close( &self.effective_root_id, @@ -371,7 +466,24 @@ impl AgentTranslator for MuseTranslator { ); if !closed.is_empty() { self.closed_turns.insert(id.into()); + self.last_completed_turn = Some(id.to_owned()); out.extend(closed); + out.push(SpanOp::Merge(SpanRow { + span_id: self.root_id.clone(), + root_span_id: self.effective_root_id.clone(), + output: p + .get("last_assistant_message") + .cloned() + .filter(|v| !v.is_null()), + metrics: Some(Value::Object(self.total_metrics.clone())), + error: event_error(p), + late_merge_key: Some(format!("muse:turn:{id}:{}", e.ts_ms)), + late_merge_version: Some(super::LateMergeVersion { + ts_ms: e.ts_ms, + native_id: Some(id.to_owned()), + }), + ..Default::default() + })); } } else if self.opened { out.push(self.unpaired(e, "Muse Stop has no turn_id")); @@ -386,12 +498,30 @@ impl AgentTranslator for MuseTranslator { end_ms: Some(e.ts_ms), metadata: p.get("reason").map(|reason| json!({"end_reason":reason})), error: event_error(p), + metrics: Some(Value::Object(self.total_metrics.clone())), + late_merge_key: Some(format!("muse:session:end:{}", e.ts_ms)), + late_merge_version: Some(super::LateMergeVersion { + ts_ms: e.ts_ms, + native_id: self.last_completed_turn.clone(), + }), ..Default::default() })); self.opened = false; } } - _ => {} + _ => out.extend(self.handle_events(e, &ids)), + } + // A native export can enrich a previously delivered preview at the + // same timestamp. Deduplicate the material completion, not only time. + for op in &mut out { + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op; + if let Some(key) = &row.late_merge_key { + let content = json!({"output":row.output,"metrics":row.metrics,"error":row.error,"metadata":row.metadata}); + row.late_merge_key = Some(format!( + "{key}:{}", + ids::span_id(&self.session_id, &content.to_string()) + )); + } } self.git .enrich_rows(p.get("cwd").and_then(Value::as_str), &mut out); @@ -415,21 +545,39 @@ impl AgentTranslator for MuseTranslator { } fn event_error(payload: &Value) -> Option { - payload - .get("error") - .and_then(|error| { - error - .as_str() - .map(str::to_owned) - .or_else(|| (!error.is_null()).then(|| error.to_string())) - }) - .or_else(|| { - payload - .get("status") - .and_then(Value::as_str) - .filter(|status| !matches!(*status, "success" | "completed")) - .map(|status| format!("Muse operation {status}")) - }) + let render = |value: &Value| match value { + Value::Null => None, + Value::String(text) => (!text.is_empty()).then(|| text.clone()), + value => Some(value.to_string()), + }; + let error = payload.get("error").and_then(render).or_else(|| { + payload + .get("status") + .and_then(Value::as_str) + // Native model completions include successful statuses such as + // `tool_calls`. Unknown future statuses are not evidence of failure. + .filter(|status| { + matches!( + *status, + "failed" + | "failure" + | "error" + | "cancelled" + | "canceled" + | "interrupted" + | "timeout" + | "timed_out" + | "denied" + | "rejected" + ) + }) + .map(|status| format!("Muse operation {status}")) + }); + match (error, payload.get("error_details").and_then(render)) { + (Some(error), Some(details)) => Some(format!("{error}: {details}")), + (Some(error), None) => Some(error), + (None, details) => details, + } } fn usage_metrics(usage: &Value) -> Value { @@ -437,7 +585,9 @@ fn usage_metrics(usage: &Value) -> Value { for (native, standard) in [ ("input_tokens", "prompt_tokens"), ("output_tokens", "completion_tokens"), + ("cache_read_tokens", "prompt_cached_tokens"), ("cached_tokens", "prompt_cached_tokens"), + ("cache_write_tokens", "prompt_cache_creation_tokens"), ("reasoning_tokens", "completion_reasoning_tokens"), ] { if let Some(value) = usage.get(native).and_then(Value::as_u64) { @@ -552,6 +702,126 @@ mod tests { assert!(translator.finalize(&ctx).unwrap().is_empty()); } + #[test] + fn resumed_session_keeps_one_root_and_publishes_latest_output_and_totals() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + let ctx = context(); + let mut rows = Vec::new(); + for (turn, ts) in [("t1", 1), ("t2", 10)] { + for event in [ + event( + "UserPromptSubmit", + json!({"turn_id":turn,"prompt":"hello"}), + ts, + ), + event( + "PreLLMCall", + json!({"turn_id":turn,"request_id":format!("{turn}:0:1"),"options":{"temperature":0.2}}), + ts + 1, + ), + event( + "PostLLMCall", + json!({"turn_id":turn,"request_id":format!("{turn}:0:1"),"finish_reason":"stop","response_id":"native-response","tool_call_count":0,"usage":{"input_tokens":3,"output_tokens":2}}), + ts + 2, + ), + event( + "Stop", + json!({"turn_id":turn,"last_assistant_message":turn}), + ts + 3, + ), + event("SessionEnd", json!({}), ts + 4), + ] { + rows.extend(translator.handle(&event, &ctx).unwrap()); + } + } + assert_eq!( + rows.iter() + .filter(|op| matches!(op, SpanOp::Insert(row) if row.name == "Muse Code")) + .count(), + 1 + ); + let root = ids::span_id("muse-fixture-session", "root"); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == root && row.output == Some(json!("t2")) && row.metrics.as_ref().unwrap()["tokens"] == 10 && row.late_merge_key.is_some()))); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.metadata.as_ref().is_some_and(|m| m["response_id"] == "native-response")))); + } + + #[test] + fn stop_failure_details_reach_turn_and_root() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + translator + .handle( + &event( + "UserPromptSubmit", + json!({"turn_id":"turn","prompt":"hello"}), + 1, + ), + &context(), + ) + .unwrap(); + let rows = translator.handle(&event("StopFailure", json!({"turn_id":"turn","error":"provider_error","error_details":"rate limited"}), 2), &context()).unwrap(); + assert_eq!(rows.iter().filter(|op| matches!(op, SpanOp::Merge(row) if row.error.as_deref() == Some("provider_error: rate limited"))).count(), 2); + } + + #[test] + fn native_cumulative_hook_usage_becomes_per_model_metrics_without_overcounting_root() { + // Muse 1.4.11 deterministic two-model native run: live PostLLMCall + // reports 10/5 then 20/10, while each durable model_completed is 10/5. + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + let ctx = context(); + translator + .handle( + &event( + "UserPromptSubmit", + json!({"turn_id":"turn","prompt":"hello"}), + 1, + ), + &ctx, + ) + .unwrap(); + for (step, input, output) in [(0, 10, 5), (1, 20, 10)] { + let request = format!("turn:{step}:1"); + translator + .handle( + &event( + "PreLLMCall", + json!({"turn_id":"turn","request_id":request}), + step + 2, + ), + &ctx, + ) + .unwrap(); + let complete = event( + "PostLLMCall", + json!({"turn_id":"turn","request_id":request,"usage":{"input_tokens":input,"output_tokens":output}}), + step + 3, + ); + let rows = translator.handle(&complete, &ctx).unwrap(); + assert!(rows.iter().any(|op|matches!(op,SpanOp::Merge(row) if row.metrics.as_ref().is_some_and(|m|m["prompt_tokens"] == 10 && m["completion_tokens"] == 5)))); + assert!(translator.handle(&complete, &ctx).unwrap().is_empty()); + } + let rows = translator + .handle( + &event( + "Stop", + json!({"turn_id":"turn","last_assistant_message":"done"}), + 10, + ), + &ctx, + ) + .unwrap(); + let root = ids::span_id("muse-fixture-session", "root"); + assert!(rows.iter().any(|op|matches!(op,SpanOp::Merge(row) if row.span_id == root && row.metrics.as_ref().is_some_and(|m|m["prompt_tokens"] == 20 && m["completion_tokens"] == 10 && m["tokens"] == 30)))); + } + #[test] fn passive_and_unknown_hooks_do_not_create_a_trace() { let mut translator = MuseTranslator::new( @@ -634,7 +904,7 @@ mod tests { let unsupported = translator .handle( &event( - "PermissionRequest", + "UnrecognizedEvent", json!({"turn_id":"t1","request_id":"p1"}), 1001, ), diff --git a/bt-daemon/src/translate/muse_events.rs b/bt-daemon/src/translate/muse_events.rs new file mode 100644 index 00000000..7ed7c471 --- /dev/null +++ b/bt-daemon/src/translate/muse_events.rs @@ -0,0 +1,474 @@ +//! Additional native Muse lifecycle events. Keep hook forwarding free of trace policy. +use super::*; + +#[derive(Default)] +pub(super) struct Events { + open: HashMap, + closed: RecentSet, +} + +impl MuseTranslator { + pub(super) fn handle_events(&mut self, e: &Envelope, hook_ids: &HookIds) -> Vec { + if !self.opened { + return Vec::new(); + } + let p = &e.payload; + match e.event.as_str() { + "PreToolUse" | "ToolUseStart" => self.start_tool(e, hook_ids, p), + "PostToolUse" | "PostToolUseFailure" => self.finish_tool(e, hook_ids, p), + "PostToolBatch" => p + .get("tool_calls") + .and_then(Value::as_array) + .map(|calls| { + calls + .iter() + .flat_map(|call| self.finish_tool(e, hook_ids, call)) + .collect() + }) + .unwrap_or_default(), + "PermissionRequest" | "Notification" => { + // Muse provides no permission operation id or decision event. Record the + // observation without claiming to know its duration or approval result. + let identity = format!("observation:{}:{}:{}", e.event, e.ts_ms, p); + let key = ids::span_id(&self.session_id, &identity); + if !self.events.closed.insert(key.clone()) { + return Vec::new(); + } + let mut row = self.event_row(e, hook_ids, &key); + row.name = e.event.clone(); + row.end_ms = Some(e.ts_ms); + row.input = p.get("tool_input").cloned(); + row.metadata = Some(selected( + p, + &["tool_name", "notification_type", "title", "message"], + )); + row.tags = Some(vec![if e.event == "Notification" { + "notification".into() + } else { + "permission-request".into() + }]); + vec![SpanOp::Insert(row)] + } + "SubagentStart" | "SubagentStop" => { + // subagent_id is an agent name (e.g. "skill-reminder"), so + // reusing it must not conflate multiple child sessions. + let Some(child) = string(p, "child_session_id") else { + return vec![self.unpaired(e, "Muse subagent event has no child identity")]; + }; + let key = format!("subagent:{child}"); + if e.event == "SubagentStart" { + let mut row = self.event_row(e, hook_ids, &key); + row.name = string(p, "subagent_id").unwrap_or("Subagent").into(); + row.metadata = Some(selected(p, &["subagent_id", "child_session_id"])); + self.open_event(&key, row) + } else { + self.finish_event( + &key, + e.ts_ms, + p.get("last_assistant_message").cloned(), + event_error(p), + None, + ) + } + } + "PreCompact" | "PostCompact" => { + // Compaction is serialized within a turn but has no native operation id. + // Pair by turn and trigger and retain that limitation on the span. + let key = format!( + "compact:{}:{}", + hook_ids.turn().unwrap_or("session"), + string(p, "trigger").unwrap_or("unknown") + ); + if e.event == "PreCompact" { + // A later compaction in the same turn is a distinct operation. + let key = format!("{key}:{}", e.ts_ms); + let mut row = self.event_row(e, hook_ids, &key); + row.name = "Context compaction".into(); + row.tags = Some(vec!["compaction".into()]); + row.metadata = Some( + json!({"trigger":p.get("trigger"), "native_operation_id_available":false}), + ); + self.open_event(&key, row) + } else { + // Choose only a uniquely matching open start; never guess across + // concurrent/malformed starts. + let matches: Vec<_> = self + .events + .open + .keys() + .filter(|candidate| candidate.starts_with(&format!("{key}:"))) + .cloned() + .collect(); + if matches.len() == 1 { + self.finish_event(&matches[0], e.ts_ms, None, event_error(p), None) + } else { + Vec::new() + } + } + } + _ => Vec::new(), + } + } + + fn event_row(&self, e: &Envelope, hook_ids: &HookIds, key: &str) -> SpanRow { + let parent = self.parent(hook_ids); + SpanRow { + span_id: ids::span_id(&self.session_id, key), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: vec![parent.clone()], + turn_span_id: (parent != self.root_id).then_some(parent), + start_ms: Some(e.ts_ms), + ..Default::default() + } + } + + fn open_event(&mut self, key: &str, row: SpanRow) -> Vec { + if self.events.closed.contains(&key.to_owned()) { + return Vec::new(); + } + Self::insert_open(&mut self.events.open, key, row) + } + + fn finish_event( + &mut self, + key: &str, + ts: i64, + output: Option, + error: Option, + metadata: Option, + ) -> Vec { + let mut out = Self::close( + &self.effective_root_id, + &mut self.events.open, + key, + ts, + output, + None, + error, + ); + if !out.is_empty() { + self.events.closed.insert(key.into()); + for op in &mut out { + let (SpanOp::Merge(row) | SpanOp::Insert(row)) = op; + row.metadata = metadata.clone(); + } + } + out + } + + fn start_tool(&mut self, e: &Envelope, hook_ids: &HookIds, p: &Value) -> Vec { + let Some(id) = string(p, "tool_use_id") else { + return vec![self.unpaired(e, "Muse tool event has no tool_use_id")]; + }; + let key = format!("tool:{id}"); + let mut row = self.event_row(e, hook_ids, &key); + row.name = string(p, "tool_name").unwrap_or("tool").into(); + row.span_type = SpanType::Tool; + row.input = p.get("tool_input").cloned(); + row.metadata = Some(json!({"tool_name":row.name})); + self.open_event(&key, row) + } + + fn finish_tool(&mut self, e: &Envelope, hook_ids: &HookIds, p: &Value) -> Vec { + let Some(id) = string(p, "tool_use_id") else { + return vec![self.unpaired(e, "Muse tool result has no tool_use_id")]; + }; + let key = format!("tool:{id}"); + if self.events.closed.contains(&key) { + return Vec::new(); + } + let failed = e.event == "PostToolUseFailure"; + let error = event_error(p).or_else(|| failed.then(|| "Muse tool failed".into())); + let outcome = if error.is_some() { + Some("error") + } else if e.event == "PostToolUse" { + Some("success") + } else { + // Batch results may contain unavailable/denied tools without a + // machine-readable failure field. Their prose is not an outcome API. + None + }; + let mut out = Vec::new(); + if !self.events.open.contains_key(&key) { + out.extend(self.start_tool(e, hook_ids, p)); + if let Some(SpanOp::Insert(row)) = out.first_mut() { + row.start_ms = p + .get("duration_ms") + .and_then(Value::as_i64) + .filter(|duration| *duration >= 0) + .map(|duration| e.ts_ms.saturating_sub(duration)) + .or(Some(e.ts_ms)); + row.metadata = Some(json!({"tool_name":row.name, "native_start_missing":true})); + } + } + out.extend(self.finish_event( + &key, + e.ts_ms, + p.get("tool_response").cloned(), + error, + outcome.map(|outcome| json!({"tool_outcome":outcome})), + )); + out + } + + pub(super) fn close_events(&mut self, ts: i64, error: Option<&str>) -> Vec { + let keys: Vec<_> = self.events.open.keys().cloned().collect(); + keys.into_iter() + .flat_map(|key| self.finish_event(&key, ts, None, error.map(str::to_owned), None)) + .collect() + } +} + +fn string<'a>(p: &'a Value, key: &str) -> Option<&'a str> { + p.get(key) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) +} +fn selected(p: &Value, keys: &[&str]) -> Value { + Value::Object( + keys.iter() + .filter_map(|key| p.get(key).map(|value| ((*key).into(), value.clone()))) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn model_tool_calls_status_is_success_and_batch_only_outcomes_stay_unknown() { + for status in [ + "success", + "completed", + "tool_calls", + "stop", + "future_status", + ] { + assert_eq!(event_error(&json!({"status":status,"error":null})), None); + } + assert_eq!(event_error(&json!({"status":"success","error":""})), None); + assert_eq!( + event_error(&json!({"status":"failed","error_details":"network unavailable"})), + Some("Muse operation failed: network unavailable".into()) + ); + let mut t = translator(); + let rows = t.handle(&event("PostToolBatch",json!({"turn_id":"t1","tool_calls":[{"tool_name":"bash","tool_use_id":"unavailable","tool_input":{},"tool_response":"tool unavailable: unknown tool `bash`"},{"tool_name":"shell","tool_use_id":"explicit-error","tool_input":{},"tool_response":"failed","error":"access denied"}]}),1010), &context()).unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == ids::span_id("muse-fixture-session","tool:unavailable") && row.error.is_none() && row.metadata.is_none()))); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == ids::span_id("muse-fixture-session","tool:explicit-error") && row.error.as_deref() == Some("access denied") && row.metadata.as_ref().and_then(|m|m.get("tool_outcome")) == Some(&json!("error"))))); + } + #[test] + fn captured_modern_hooks_keep_native_child_sessions_distinct() { + // Captured from Muse Code 1.4.10-R6427.1 with a managed `cat` hook, + // `muse exec --provider echo --no-session-log modern hook fixture`. + // Session/turn ids and the scratch workspace path are sanitized. + // Builtin reminder child hooks contain no parent session/turn identity. + let registry = super::super::super::Registry::default_agents(); + let mut translators = HashMap::new(); + let mut child_starts = 0; + let mut llms = 0; + for (index, line) in include_str!("../../tests/fixtures/muse-hooks-1.4.10.jsonl") + .lines() + .enumerate() + { + let payload: Value = serde_json::from_str(line).unwrap(); + let session = payload["session_id"].as_str().unwrap(); + let name = payload["hook_event_name"].as_str().unwrap(); + if name == "SubagentStart" { + child_starts += 1; + assert_eq!(payload["session_id"], payload["child_session_id"]); + assert!(payload.get("parent_session_id").is_none()); + } + let mut e = event(name, payload.clone(), 1000 + index as i64); + e.session_id = session.into(); + let ctx = SessionCtx { + session_id: session.into(), + config: None, + }; + let translator = translators + .entry(session.to_owned()) + .or_insert_with(|| registry.create("muse", session)); + let rows = translator.handle(&e, &ctx).unwrap(); + for op in rows { + if let SpanOp::Insert(row) = op { + if row.span_type == SpanType::Llm { + llms += 1; + assert_eq!( + row.root_span_id, + ids::span_id(&ids::session_namespace("muse", session), "root") + ); + } + } + } + } + assert!(child_starts > 0); + assert!(llms >= 2); + } + fn context() -> SessionCtx { + SessionCtx { + session_id: "muse-fixture-session".into(), + config: None, + } + } + fn event(name: &str, payload: Value, ts_ms: i64) -> Envelope { + Envelope { + source: "muse".into(), + source_version: Some("1.4.10-R6427.1".into()), + plugin_version: None, + session_id: "muse-fixture-session".into(), + event: name.into(), + ts_ms, + managed_run_id: None, + capture: None, + payload, + route: None, + config: None, + } + } + fn translator() -> MuseTranslator { + let mut t = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + t.handle( + &event( + "UserPromptSubmit", + json!({"turn_id":"t1","prompt":"x"}), + 1000, + ), + &context(), + ) + .unwrap(); + t + } + #[test] + fn parallel_tools_pair_by_id_and_batch_results_do_not_duplicate() { + let mut t = translator(); + for id in ["a", "b"] { + let rows = t.handle(&event("PreToolUse", json!({"turn_id":"t1","tool_use_id":id,"tool_name":"read_file","tool_input":{"path":id}}), 1010), &context()).unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.span_type == SpanType::Tool && row.end_ms.is_none() && row.turn_span_id.is_some()))); + } + let rows = t + .handle( + &event( + "PostToolUse", + json!({"turn_id":"t1","tool_use_id":"b","tool_response":{"text":"b"}}), + 1020, + ), + &context(), + ) + .unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.span_id == ids::span_id("muse-fixture-session", "tool:b") && row.metadata.as_ref().and_then(|m|m.get("tool_outcome")) == Some(&json!("success"))))); + let rows = t + .handle( + &event( + "PostToolUseFailure", + json!({"turn_id":"t1","tool_use_id":"a","error":"file absent"}), + 1021, + ), + &context(), + ) + .unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.error.as_deref() == Some("file absent") && row.metadata.as_ref().and_then(|m|m.get("tool_outcome")) == Some(&json!("error"))))); + assert!(t.handle(&event("PostToolBatch",json!({"turn_id":"t1","tool_calls":[{"tool_use_id":"a","tool_response":"ignored"},{"tool_use_id":"b","tool_response":"ignored"}]}),1022), &context()).unwrap().is_empty()); + assert!(t + .handle( + &event( + "PreToolUse", + json!({"turn_id":"t1","tool_use_id":"a"}), + 1023 + ), + &context() + ) + .unwrap() + .is_empty()); + } + #[test] + fn permission_is_an_observation_and_compaction_pairs_only_one_start() { + let mut t = translator(); + let rows = t.handle(&event("PermissionRequest",json!({"turn_id":"t1","tool_name":"shell","tool_input":{"command":"echo hello"}}),1010), &context()).unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.start_ms == row.end_ms && row.error.is_none() && row.metadata.as_ref().and_then(|m|m.get("tool_outcome")).is_none()))); + let notification = t + .handle( + &event( + "Notification", + json!({"turn_id":"t1","notification_type":"future_notice"}), + 1011, + ), + &context(), + ) + .unwrap(); + assert!(notification.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.tags.as_deref() == Some(&["notification".into()])))); + for (name, ts) in [("PreCompact", 1020), ("PostCompact", 1030)] { + assert_eq!( + t.handle( + &event(name, json!({"turn_id":"t1","trigger":"auto"}), ts), + &context() + ) + .unwrap() + .len(), + 1 + ); + } + assert!(t + .handle( + &event( + "PostCompact", + json!({"turn_id":"t1","trigger":"auto"}), + 1031 + ), + &context() + ) + .unwrap() + .is_empty()); + } + #[test] + fn subagent_lifecycle_and_pending_tool_close_once() { + let mut t = translator(); + let p = json!({"turn_id":"t1","subagent_id":"worker","child_session_id":"child"}); + let rows = t + .handle(&event("SubagentStart", p.clone(), 1010), &context()) + .unwrap(); + assert!(rows.iter().any(|op| matches!(op, SpanOp::Insert(row) if row.metadata.as_ref().and_then(|m|m.get("child_session_id")) == Some(&json!("child"))))); + assert!(t + .handle(&event("SubagentStop", p, 1020), &context()) + .unwrap() + .iter() + .any(|op| matches!(op, SpanOp::Merge(row) if row.end_ms == Some(1020)))); + t.handle( + &event( + "PreToolUse", + json!({"turn_id":"t1","tool_use_id":"pending"}), + 1030, + ), + &context(), + ) + .unwrap(); + let rows = t.finalize(&context()).unwrap(); + assert_eq!(rows.iter().filter(|op| matches!(op, SpanOp::Merge(row) if row.span_id == ids::span_id("muse-fixture-session","tool:pending") && row.error.is_some())).count(),1); + assert!(t.finalize(&context()).unwrap().is_empty()); + } + #[test] + fn orphan_failure_preserves_duration_and_missing_ids_cannot_collide() { + let mut t = translator(); + let rows = t.handle(&event("PostToolUseFailure",json!({"turn_id":"t1","tool_use_id":"orphan","error":"failed","duration_ms":20}),1050), &context()).unwrap(); + assert!(rows + .iter() + .any(|op| matches!(op, SpanOp::Insert(row) if row.start_ms == Some(1030)))); + let a = t + .handle( + &event("PreToolUse", json!({"turn_id":"t1","tool_name":"a"}), 1060), + &context(), + ) + .unwrap(); + let b = t + .handle( + &event("PreToolUse", json!({"turn_id":"t1","tool_name":"b"}), 1060), + &context(), + ) + .unwrap(); + let id = |rows: Vec| match rows.into_iter().next().unwrap() { + SpanOp::Insert(row) | SpanOp::Merge(row) => row.span_id, + }; + assert_ne!(id(a), id(b)); + } +} diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 98c24903..88d76837 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -61,6 +61,7 @@ fn row( metrics: None, error: None, late_merge_key: None, + late_merge_version: None, turn_span_id: None, tags: None, origin: Some(OriginSnapshot { diff --git a/bt-daemon/tests/fixtures/muse-export-1.4.11.json b/bt-daemon/tests/fixtures/muse-export-1.4.11.json new file mode 100644 index 00000000..20a1ecfd --- /dev/null +++ b/bt-daemon/tests/fixtures/muse-export-1.4.11.json @@ -0,0 +1,347 @@ +{ + "export_schema_version": 1, + "exporter_version": { + "semver": "1.4.11" + }, + "sessions": [ + { + "session_id": "muse-export-fixture", + "session_end": { + "exit_reason": "clean" + } + } + ], + "events": [ + { + "recorded_at": -685756, + "envelope": { + "payload": { + "kind": "metadata", + "record": { + "model_id": "mock-model", + "provider_id": "meta", + "workspace_root": "/workspace" + } + } + } + }, + { + "recorded_at": 438662, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "started", + "prompt": "Print an integration marker." + } + } + } + }, + { + "recorded_at": 544695, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "bounded": { + "request_message_count": 2 + }, + "kind": "model_input_trace_recorded", + "model_step": 0, + "provider_family": "meta", + "request_message_count": 2, + "request_prefix_message_count": 1, + "request_record_id": "fixture-request", + "schema_version": 2, + "scope": "full_request" + } + } + } + }, + { + "recorded_at": 586838, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "model_response_created", + "response_id": "resp_muse_tool" + } + } + } + }, + { + "recorded_at": 604709, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "duration_ms": 24, + "estimated_prompt_tokens": 10, + "kind": "model_completed", + "model": "mock-model", + "usage": { + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "input_tokens": 10, + "output_tokens": 5, + "reasoning_tokens": 0 + } + } + } + } + }, + { + "recorded_at": 633823, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "assistant_tool_calls_committed", + "message_id": "fixture-batch", + "response_id": "resp_muse_tool", + "tool_calls": [ + { + "args": "{\"command\":\"printf MUSE_TOOL_OK\"}", + "call_id": "call_muse_1", + "id": "fc_muse_1", + "name": "shell" + } + ] + } + } + } + }, + { + "recorded_at": 697679, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "kind": "proposed", + "task_id": "fixture-tool-task", + "task_kind": "tool.shell" + } + } + } + }, + { + "recorded_at": 743742, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "kind": "accepted", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 747823, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "idempotency_key": "tool:call_muse_1", + "kind": "scheduled", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 751752, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "cancellation_handle": null, + "idempotency_key": "tool:call_muse_1", + "kind": "side_effect_intent", + "operation": "tool:shell", + "parent_task_id": null, + "policy_decision": "allow:policy", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 758731, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "kind": "started", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 809169, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "chunk": "MUSE_TOOL_OK", + "final_result": true, + "kind": "output", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 814762, + "envelope": { + "payload": { + "kind": "task", + "run_id": "fixture-turn", + "task_id": "fixture-tool-task", + "event": { + "kind": "completed", + "task_id": "fixture-tool-task" + } + } + } + }, + { + "recorded_at": 858622, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "batch_id": "fixture-batch", + "kind": "tool_result_batch_committed", + "results": [ + { + "text": "MUSE_TOOL_OK", + "tool_call_id": "call_muse_1", + "tool_call_index": 0 + } + ] + } + } + } + }, + { + "recorded_at": 934575, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "bounded": { + "request_message_count": 4 + }, + "kind": "model_input_trace_recorded", + "model_step": 1, + "provider_family": "meta", + "request_message_count": 4, + "request_prefix_message_count": 1, + "request_record_id": "fixture-request", + "schema_version": 2, + "scope": "full_request" + } + } + } + }, + { + "recorded_at": 974464, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "model_response_created", + "response_id": "resp_mock_1" + } + } + } + }, + { + "recorded_at": 979674, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "assistant_message_committed", + "message_id": "f546e181-b8ed-45f2-af8b-fdb587d2a628", + "provider_item_id": "msg_mock_1", + "response_id": "resp_mock_1", + "text": "MUSE_MOCK_OK" + } + } + } + }, + { + "recorded_at": 996545, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "duration_ms": 26, + "estimated_prompt_tokens": 18, + "kind": "model_completed", + "model": "mock-model", + "usage": { + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "input_tokens": 10, + "output_tokens": 5, + "reasoning_tokens": 0 + } + } + } + } + }, + { + "recorded_at": 1096656, + "envelope": { + "payload": { + "kind": "run", + "run_id": "fixture-turn", + "event": { + "kind": "terminal", + "reason": null, + "terminal": "completed", + "time_to_first_token_ms": 603, + "turn_duration_ms": 720 + } + } + } + }, + { + "recorded_at": 1170037, + "envelope": { + "payload": { + "kind": "session_end" + } + } + } + ] +} diff --git a/bt-daemon/tests/fixtures/muse-export-tool-failure-1.4.11.json b/bt-daemon/tests/fixtures/muse-export-tool-failure-1.4.11.json new file mode 100644 index 00000000..acce6386 --- /dev/null +++ b/bt-daemon/tests/fixtures/muse-export-tool-failure-1.4.11.json @@ -0,0 +1,348 @@ +{ + "export_schema_version": 1, + "exporter_version": { + "semver": "1.4.11" + }, + "sessions": [ + { + "session_id": "muse-failure-fixture", + "session_end": { + "exit_reason": "clean" + } + } + ], + "events": [ + { + "recorded_at": 1000000, + "envelope": { + "payload": { + "kind": "metadata", + "record": { + "model_id": "mock-model", + "provider_id": "meta", + "workspace_root": "/workspace" + } + } + } + }, + { + "recorded_at": 2104594, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "started", + "prompt": "Run a failing integration command." + } + } + } + }, + { + "recorded_at": 2214403, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "bounded": { + "request_message_count": 2 + }, + "kind": "model_input_trace_recorded", + "model_step": 0, + "provider_family": "meta", + "request_message_count": 2, + "request_prefix_message_count": 1, + "request_record_id": "failure-request", + "schema_version": 2, + "scope": "full_request" + } + } + } + }, + { + "recorded_at": 2258957, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "model_response_created", + "response_id": "resp_call_muse_failure" + } + } + } + }, + { + "recorded_at": 2281286, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "duration_ms": 22, + "estimated_prompt_tokens": 10, + "kind": "model_completed", + "model": "mock-model", + "usage": { + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "input_tokens": 10, + "output_tokens": 5, + "reasoning_tokens": 0 + } + } + } + } + }, + { + "recorded_at": 2312338, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "assistant_tool_calls_committed", + "message_id": "failure-batch", + "response_id": "resp_call_muse_failure", + "tool_calls": [ + { + "args": "{\"command\":\"printf MUSE_TOOL_FAILURE; exit 7\"}", + "call_id": "call_muse_failure", + "id": "fc_call_muse_failure", + "name": "shell" + } + ] + } + } + } + }, + { + "recorded_at": 2373307, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "kind": "proposed", + "task_id": "failure-tool-task", + "task_kind": "tool.shell" + } + } + } + }, + { + "recorded_at": 2415592, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "kind": "accepted", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2419422, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "idempotency_key": "tool:call_muse_failure", + "kind": "scheduled", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2423304, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "cancellation_handle": null, + "idempotency_key": "tool:call_muse_failure", + "kind": "side_effect_intent", + "operation": "tool:shell", + "parent_task_id": null, + "policy_decision": "allow:policy", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2431463, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "kind": "started", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2481320, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "chunk": "tool failed: exit status: 7\nstdout:\nMUSE_TOOL_FAILURE", + "final_result": true, + "kind": "output", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2489421, + "envelope": { + "payload": { + "kind": "task", + "run_id": "failure-turn", + "task_id": "failure-tool-task", + "event": { + "kind": "failed", + "reason": "exit status: 7\nstdout:\nMUSE_TOOL_FAILURE", + "task_id": "failure-tool-task" + } + } + } + }, + { + "recorded_at": 2542413, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "batch_id": "failure-batch", + "kind": "tool_result_batch_committed", + "results": [ + { + "text": "tool failed: exit status: 7\nstdout:\nMUSE_TOOL_FAILURE", + "tool_call_id": "call_muse_failure", + "tool_call_index": 0 + } + ] + } + } + } + }, + { + "recorded_at": 2622460, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "bounded": { + "request_message_count": 4 + }, + "kind": "model_input_trace_recorded", + "model_step": 1, + "provider_family": "meta", + "request_message_count": 4, + "request_prefix_message_count": 1, + "request_record_id": "failure-request", + "schema_version": 2, + "scope": "full_request" + } + } + } + }, + { + "recorded_at": 2665876, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "model_response_created", + "response_id": "resp_mock_6" + } + } + } + }, + { + "recorded_at": 2672505, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "assistant_message_committed", + "message_id": "failure-batch", + "provider_item_id": "msg_mock_6", + "response_id": "resp_mock_6", + "text": "MUSE_FAILURE_HANDLED" + } + } + } + }, + { + "recorded_at": 2696355, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "duration_ms": 31, + "estimated_prompt_tokens": 20, + "kind": "model_completed", + "model": "mock-model", + "usage": { + "cache_read_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 0, + "input_tokens": 10, + "output_tokens": 5, + "reasoning_tokens": 0 + } + } + } + } + }, + { + "recorded_at": 2791452, + "envelope": { + "payload": { + "kind": "run", + "run_id": "failure-turn", + "event": { + "kind": "terminal", + "reason": null, + "terminal": "completed", + "time_to_first_token_ms": 637, + "turn_duration_ms": 756 + } + } + } + }, + { + "recorded_at": 2864316, + "envelope": { + "payload": { + "kind": "session_end" + } + } + } + ] +} diff --git a/bt-daemon/tests/fixtures/muse-hooks-1.4.10.jsonl b/bt-daemon/tests/fixtures/muse-hooks-1.4.10.jsonl new file mode 100644 index 00000000..230cd314 --- /dev/null +++ b/bt-daemon/tests/fixtures/muse-hooks-1.4.10.jsonl @@ -0,0 +1,11 @@ +{"hook_event_name":"SessionStart","source":"startup","session_id":"muse-modern-1","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"UserPromptSubmit","prompt":"modern hook fixture","session_id":"muse-modern-1","turn_id":"muse-modern-2","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"PreLLMCall","provider":"model.unknown.response","request_id":"muse-modern-2:0:1","attempt":1,"step":0,"messages":[{"role":"user","content":[{"type":"text","text":"modern hook fixture"}]}],"message_count":1,"tools":[{"name":"workflow","description":"Use this to orchestrate multi-agent work in deterministic JavaScript. Follow the current workflow availability context to decide whether to launch, propose, or abstain. For a new run, provide an inline `script` as a JavaScript module such as `export defaul...[truncated]","has_parameters":true,"strict":false},{"name":"read_file","description":"Read a line-numbered UTF-8 text file window, or attach a supported image or MP4/MOV video file as model-visible output.","has_parameters":true,"strict":false},{"name":"search","description":"Search files with native ripgrep semantics. Results are confined by the current filesystem policy and emitted through tool output. Prefer this tool over shelling out to `rg`, `find`, or `grep -r` via bash: it is policy-confined, output-bounded, and watchdo...[truncated]","has_parameters":true,"strict":false},{"name":"write_file","description":"Create or overwrite a complete UTF-8 file admitted by the current filesystem policy. For a LARGE file, write a small first chunk here and then grow it with edit_file \u2014 one huge write can exceed a single model response and fail to send.","has_parameters":true,"strict":false},{"name":"edit_file","description":"Replace one unique exact text match in a file admitted by the current filesystem policy. Also use this to GROW a large file in steps: match its current last line(s) and replace them with those line(s) plus more, so you never send one huge write_file that c...[truncated]","has_parameters":true,"strict":false},{"name":"read_memory","description":"Read a bounded line window from one local Markdown memory file. Use this when you need live memory content; reads never write to memory.","has_parameters":true,"strict":false},{"name":"add_memory","description":"Add Markdown content to local memory: creates the file when it is missing, appends to the end when it already exists, and does not overwrite existing content. Use edit_memory for exact replacements.","has_parameters":true,"strict":false},{"name":"edit_memory","description":"Replace one exact string in local Markdown memory. The edit fails unless old_str appears exactly once; use add_memory to append new content.","has_parameters":true,"strict":false},{"name":"work_stop","description":"Stop one runtime-owned work item by canonical Work ID, such as a launched workflow run or other long-running background work.","has_parameters":true,"strict":false},{"name":"work_list","description":"List current background work in this session, including Monitor, Bash, Workflow and native subagents. Use this to recover a lost canonical work_id for work_stop. Returns up to 100 items; pass next_after_work_id as after_work_id for the next page. stop_requ...[truncated]","has_parameters":true,"strict":null},{"name":"bash","description":"Run a bash-compatible shell command. By default the runtime waits at most 10 seconds in the foreground; for a slow build or test, pass a larger yield_time_ms (up to 300000) to wait for it to finish in this one call. Commands still running after the wait re...[truncated]","has_parameters":true,"strict":false},{"name":"bash_input","description":"Send input to or terminate a running bash PTY session using the internal session_id handle returned by bash \u2014 use it when a live interactive process needs input. Do not use it to poll a backgrounded command for completion: the final result is delivered aut...[truncated]","has_parameters":true,"strict":false},{"name":"monitor","description":"Monitor is for repeated events from one long-running source, never for a single completion: for one-shot work such as \"tell me when the build is done\", run the command once with bash and report. Sources: a shell command (each stdout line is an event; exit ...[truncated]","has_parameters":true,"strict":false},{"name":"read_skill","description":"Read one available SKILL.md body as a tool result.","has_parameters":true,"strict":false},{"name":"work_status","description":"Read the current state of one Work item by its canonical Work ID. This is a bounded, read-only lookup; use returned artifact references only when more detail is needed. Do not call it repeatedly to wait for a Work item to finish: completion is delivered au...[truncated]","has_parameters":true,"strict":false},{"name":"snooze_reminder","description":"Temporarily suppress matching async reminder notifications.","has_parameters":true,"strict":false},{"name":"write_todos","description":"Records the task's todo plan, which the user sees as live progress. Call it at the start of any task with three or more distinct steps, then update it as each step finishes. Always send the full list; keep exactly one item in_progress. Skip it for trivial ...[truncated]","has_parameters":true,"strict":false}],"tool_count":17,"options":{},"session_id":"muse-modern-1","turn_id":"muse-modern-2","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"SubagentStart","subagent_id":"skill-reminder","child_session_id":"muse-modern-3","session_id":"muse-modern-3","turn_id":"muse-modern-3","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"PostLLMCall","provider":"model.unknown.response","request_id":"muse-modern-2:0:1","attempt":1,"step":0,"status":"success","response_id":"muse-tui-echo","usage":{"input_tokens":0,"output_tokens":0,"cached_tokens":0,"reasoning_tokens":0},"finish_reason":null,"error":null,"output_text_preview":"echo: modern hook fixture","tool_call_count":0,"messages":[{"role":"developer","content":[{"type":"text","text":"\nWorkspace root: /workspace\nWorkspace-relative tool paths resolve against this root.\n\n\nSession permission mode (as of se...[truncated]"}]},{"role":"user","content":[{"type":"text","text":"modern hook fixture"}]}],"message_count":2,"tools":[{"name":"workflow","description":"Use this to orchestrate multi-agent work in deterministic JavaScript. Follow the current workflow availability context to decide whether to launch, propose, or abstain. For a new run, provide an inline `script` as a JavaScript module such as `export defaul...[truncated]","has_parameters":true,"strict":false},{"name":"read_file","description":"Read a line-numbered UTF-8 text file window, or attach a supported image or MP4/MOV video file as model-visible output.","has_parameters":true,"strict":false},{"name":"search","description":"Search files with native ripgrep semantics. Results are confined by the current filesystem policy and emitted through tool output. Prefer this tool over shelling out to `rg`, `find`, or `grep -r` via bash: it is policy-confined, output-bounded, and watchdo...[truncated]","has_parameters":true,"strict":false},{"name":"write_file","description":"Create or overwrite a complete UTF-8 file admitted by the current filesystem policy. For a LARGE file, write a small first chunk here and then grow it with edit_file \u2014 one huge write can exceed a single model response and fail to send.","has_parameters":true,"strict":false},{"name":"edit_file","description":"Replace one unique exact text match in a file admitted by the current filesystem policy. Also use this to GROW a large file in steps: match its current last line(s) and replace them with those line(s) plus more, so you never send one huge write_file that c...[truncated]","has_parameters":true,"strict":false},{"name":"read_memory","description":"Read a bounded line window from one local Markdown memory file. Use this when you need live memory content; reads never write to memory.","has_parameters":true,"strict":false},{"name":"add_memory","description":"Add Markdown content to local memory: creates the file when it is missing, appends to the end when it already exists, and does not overwrite existing content. Use edit_memory for exact replacements.","has_parameters":true,"strict":false},{"name":"edit_memory","description":"Replace one exact string in local Markdown memory. The edit fails unless old_str appears exactly once; use add_memory to append new content.","has_parameters":true,"strict":false},{"name":"work_stop","description":"Stop one runtime-owned work item by canonical Work ID, such as a launched workflow run or other long-running background work.","has_parameters":true,"strict":false},{"name":"work_list","description":"List current background work in this session, including Monitor, Bash, Workflow and native subagents. Use this to recover a lost canonical work_id for work_stop. Returns up to 100 items; pass next_after_work_id as after_work_id for the next page. stop_requ...[truncated]","has_parameters":true,"strict":null},{"name":"bash","description":"Run a bash-compatible shell command. By default the runtime waits at most 10 seconds in the foreground; for a slow build or test, pass a larger yield_time_ms (up to 300000) to wait for it to finish in this one call. Commands still running after the wait re...[truncated]","has_parameters":true,"strict":false},{"name":"bash_input","description":"Send input to or terminate a running bash PTY session using the internal session_id handle returned by bash \u2014 use it when a live interactive process needs input. Do not use it to poll a backgrounded command for completion: the final result is delivered aut...[truncated]","has_parameters":true,"strict":false},{"name":"monitor","description":"Monitor is for repeated events from one long-running source, never for a single completion: for one-shot work such as \"tell me when the build is done\", run the command once with bash and report. Sources: a shell command (each stdout line is an event; exit ...[truncated]","has_parameters":true,"strict":false},{"name":"read_skill","description":"Read one available SKILL.md body as a tool result.","has_parameters":true,"strict":false},{"name":"work_status","description":"Read the current state of one Work item by its canonical Work ID. This is a bounded, read-only lookup; use returned artifact references only when more detail is needed. Do not call it repeatedly to wait for a Work item to finish: completion is delivered au...[truncated]","has_parameters":true,"strict":false},{"name":"snooze_reminder","description":"Temporarily suppress matching async reminder notifications.","has_parameters":true,"strict":false},{"name":"write_todos","description":"Records the task's todo plan, which the user sees as live progress. Call it at the start of any task with three or more distinct steps, then update it as each step finishes. Always send the full list; keep exactly one item in_progress. Skip it for trivial ...[truncated]","has_parameters":true,"strict":false}],"tool_count":17,"options":{"meta.session_id":"muse-modern-1","meta.traceparent":"00-01a15a5b52037782948f81a15c429835-515ada85c4b955c0-01"},"session_id":"muse-modern-1","turn_id":"muse-modern-2","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"PreLLMCall","provider":"model.unknown.response","request_id":"muse-modern-3:0:1","attempt":1,"step":0,"messages":[{"role":"user","content":[{"type":"text","text":"You are a reminder observer for the main agent.\nDo not answer the user. Record exactly one decision this round by calling submit_reminder_decision. Use decision=\"remind\" only when the remind condition stated in your duty is met; otherwise use decision=\"non...[truncated]"}]}],"message_count":1,"tools":[{"name":"submit_reminder_decision","description":"Submit this reminder child's explicit decision.","has_parameters":true,"strict":true}],"tool_count":1,"options":{"anthropic.output_config.effort":"low","anthropic.thinking":"enabled","anthropic.thinking.budget_tokens":4096,"local.reasoning.effort":"low","meta.reasoning.effort":"low","openai.reasoning.effort":"low","openrouter.reasoning.effort":"low"},"session_id":"muse-modern-3","turn_id":"muse-modern-3","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"Stop","stop_hook_active":false,"last_assistant_message":"echo: modern hook fixture","session_id":"muse-modern-1","turn_id":"muse-modern-2","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"SubagentStart","subagent_id":"goal-reminder","child_session_id":"muse-modern-4","session_id":"muse-modern-4","turn_id":"muse-modern-4","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"SubagentStart","subagent_id":"verify-reminder","child_session_id":"muse-modern-5","session_id":"muse-modern-5","turn_id":"muse-modern-5","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"PreLLMCall","provider":"model.unknown.response","request_id":"muse-modern-3:1:1","attempt":1,"step":1,"messages":[{"role":"user","content":[{"type":"text","text":"You are a reminder observer for the main agent.\nDo not answer the user. Record exactly one decision this round by calling submit_reminder_decision. Use decision=\"remind\" only when the remind condition stated in your duty is met; otherwise use decision=\"non...[truncated]"}]},{"role":"assistant","content":[{"type":"assistant_message","text":"echo: You are a reminder observer for the main agent.\nDo not answer the user. Record exactly one decision this round by calling submit_reminder_decision. Use decision=\"remind\" only when the remind condition stated in your duty is met; otherwise use decisio...[truncated]","provider_item_id":null,"phase":"final_answer"}]}],"message_count":2,"tools":[{"name":"submit_reminder_decision","description":"Submit this reminder child's explicit decision.","has_parameters":true,"strict":true}],"tool_count":1,"options":{"anthropic.output_config.effort":"low","anthropic.thinking":"enabled","anthropic.thinking.budget_tokens":4096,"local.reasoning.effort":"low","meta.reasoning.effort":"low","openai.reasoning.effort":"low","openrouter.reasoning.effort":"low"},"session_id":"muse-modern-3","turn_id":"muse-modern-3","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} +{"hook_event_name":"SessionEnd","reason":"other","session_id":"muse-modern-1","cwd":"/workspace","transcript_path":null,"model":"unknown","permission_mode":"default"} diff --git a/bt-daemon/tests/muse_attach.rs b/bt-daemon/tests/muse_attach.rs new file mode 100644 index 00000000..d327ca7c --- /dev/null +++ b/bt-daemon/tests/muse_attach.rs @@ -0,0 +1,301 @@ +#![cfg(unix)] +#![allow( + clippy::unwrap_used, + clippy::expect_used, + reason = "Integration assertions fail immediately." +)] +#![allow( + clippy::disallowed_methods, + reason = "Integration fixtures control native child processes." +)] +mod support; + +use axum::middleware::{self, Next}; +use axum::{extract::Request, response::Response}; +use serde_json::{json, Value}; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::sync::Arc; +use std::time::Duration; +use support::inference::{MockReply, OpenAiMock, OpenAiTurn}; +use support::server::TestServer; +use tokio::process::{Child, Command}; +use tokio::sync::Notify; + +struct NativeSession { + root: tempfile::TempDir, + binary: std::ffi::OsString, +} +impl NativeSession { + fn new(base_url: &str) -> Self { + let root = tempfile::tempdir().unwrap(); + let config = root.path().join("config/muse"); + std::fs::create_dir_all(&config).unwrap(); + std::fs::create_dir_all(root.path().join("workspace")).unwrap(); + std::fs::create_dir_all(root.path().join("home")).unwrap(); + std::fs::write( + config.join("settings.json"), + json!({"schema_version":1,"endpoint_transport":{"base_url":base_url,"auth":"bearer"}}) + .to_string(), + ) + .unwrap(); + std::fs::write( + config.join("auth.json"), + json!({"schema_version":1,"providers":{"meta":{"api_key":"test-key"}}}).to_string(), + ) + .unwrap(); + Self { + root, + binary: std::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()), + } + } + fn command(&self, binary: impl AsRef) -> Command { + let mut command = Command::new(binary); + command + .kill_on_drop(true) + .current_dir(self.root.path().join("workspace")) + .env("HOME", self.root.path().join("home")) + .env("TBH_CREDENTIAL_BACKEND", "file") + .env("TBH_DISABLE_TELEMETRY", "1") + .env("META_API_KEY", "test-key") + .env("XDG_CONFIG_HOME", self.root.path().join("config")) + .env("XDG_DATA_HOME", self.root.path().join("data")) + .env( + "BRAINTRUST_DAEMON_DATA_DIR", + self.root.path().join("import"), + ) + .env("MUSE_BIN", &self.binary) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + for (key, _) in std::env::vars_os() { + if [ + "CODEX_", + "NODE_REPL_", + "CUA_REPL_", + "BROWSER_USE_", + "SKY_CUA_", + ] + .iter() + .any(|prefix| key.to_string_lossy().starts_with(prefix)) + { + command.env_remove(key); + } + } + command.env_remove("ELECTRON_RUN_AS_NODE"); + command + } + fn import(&self, session: &str, attach: bool) -> Command { + let mut command = self.command(env!("CARGO_BIN_EXE_bt-daemon")); + command.args(["--json", "import", "muse", session]); + if attach { + command.arg("--attach"); + } + command + } + fn rows_path(&self, session: &str) -> PathBuf { + self.root + .path() + .join("import/spans") + .join(format!("{session}.ndjson")) + } +} +fn rows(path: &Path) -> Vec { + std::fs::read_to_string(path) + .unwrap_or_default() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() +} +async fn detach(mut child: Child) { + // This is the public Ctrl-C path, delivered only to the observer process. + let pid = child.id().unwrap(); + assert_eq!(unsafe { libc::kill(pid as libc::pid_t, libc::SIGINT) }, 0); + let status = tokio::time::timeout(Duration::from_secs(10), child.wait()) + .await + .unwrap() + .unwrap(); + assert!(status.success(), "attach observer failed: {status}"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires native Muse Code 1.3+; uses deterministic loopback inference"] +async fn native_active_session_detach_reattach_preserves_authoritative_completion() { + let requested = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let inference = + OpenAiMock::new(|_, _| MockReply::response(OpenAiTurn::text("MUSE_ATTACH_COMPLETE"))); + let router = inference.router().layer(middleware::from_fn({ + let requested = requested.clone(); + let release = release.clone(); + move |request: Request, next: Next| { + let requested = requested.clone(); + let release = release.clone(); + async move { + if request.method() == axum::http::Method::POST + && request.uri().path().ends_with("/responses") + { + requested.notify_one(); + release.notified().await; + } + let response: Response = next.run(request).await; + response + } + } + })); + let server = TestServer::start(router).await; + let native = NativeSession::new(server.uri()); + let session = uuid::Uuid::new_v4().to_string(); + let mut command = native.command(&native.binary); + command.args([ + "exec", + "--session-id", + &session, + "--provider", + "meta", + "--model", + "mock-model", + "--no-foreign-personal-context", + "--disable-reminders", + "--disable-web-tools", + "--disable-request-user-input", + "--max-model-steps", + "1", + "Reply MUSE_ATTACH_COMPLETE.", + ]); + let mut agent = command.spawn().unwrap(); + tokio::time::timeout(Duration::from_secs(30), requested.notified()) + .await + .unwrap(); + let export_path = native.root.path().join("active.json"); + tokio::time::timeout(Duration::from_secs(20), async { + loop { + let exported = native + .command(&native.binary) + .args(["export", "--session", &session, "--out"]) + .arg(&export_path) + .output() + .await + .unwrap(); + if exported.status.success() { + break; + } + eprintln!( + "waiting for native retention: {}", + String::from_utf8_lossy(&exported.stderr) + ); + tokio::time::sleep(Duration::from_millis(500)).await; + } + }) + .await + .expect("Muse did not publish an active session export while inference was open"); + let active: Value = serde_json::from_slice(&std::fs::read(export_path).unwrap()).unwrap(); + assert!(active["events"] + .as_array() + .is_some_and(|events| !events.is_empty())); + let mut first = native.import(&session, true).spawn().unwrap(); + let path = native.rows_path(&session); + tokio::time::timeout(Duration::from_secs(20), async { + loop { + if rows(&path) + .iter() + .any(|row| row["Insert"]["name"] == "Muse Code") + { + break; + } + if let Some(status) = first.try_wait().unwrap() { + panic!("observer exited before active rows: {status}"); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .unwrap(); + detach(first).await; + assert!( + agent.try_wait().unwrap().is_none(), + "detaching killed native Muse" + ); + let before = rows(&path); + assert!( + before.iter().all(|row| row["Merge"]["error"].is_null()), + "detaching invented an interruption: {before:?}" + ); + assert!( + before.iter().all(|row| row["Merge"]["end_ms"].is_null()), + "detaching ended native work: {before:?}" + ); + let mut second = native.import(&session, true).spawn().unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + second.try_wait().unwrap().is_none(), + "reattach ended before native completion" + ); + release.notify_one(); + let output = tokio::time::timeout(Duration::from_secs(30), agent.wait_with_output()) + .await + .unwrap() + .unwrap(); + assert!( + output.status.success(), + "native Muse failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + let observed = tokio::time::timeout(Duration::from_secs(30), second.wait_with_output()) + .await + .unwrap() + .unwrap(); + assert!( + observed.status.success(), + "reattach failed: {}", + String::from_utf8_lossy(&observed.stderr) + ); + let completed = rows(&path); + let root_ids: std::collections::BTreeSet<_> = completed + .iter() + .filter(|row| row["Insert"]["name"] == "Muse Code") + .map(|row| row["Insert"]["span_id"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!( + root_ids.len(), + 1, + "reattach created a different root identity: {completed:?}" + ); + assert!( + completed.iter().any(|row| row["Merge"]["span_id"] + == root_ids.iter().next().unwrap().as_str() + && row["Merge"]["end_ms"].is_number() + && row["Merge"]["error"].is_null()), + "root did not receive authoritative completion: {completed:?}" + ); + assert!( + completed.iter().any(|row| row["Merge"]["output"] + .to_string() + .contains("MUSE_ATTACH_COMPLETE")), + "native final output was lost: {completed:?}" + ); + let imported = native.import(&session, false).output().await.unwrap(); + assert!( + imported.status.success(), + "final import: {}", + String::from_utf8_lossy(&imported.stderr) + ); + // The standalone debug importer has no destination-scoped delivery ledger. + // Replays may emit identical upserts; they must retain their span identities. + let identities = |rows: &[Value]| -> std::collections::BTreeSet { + rows.iter() + .filter_map(|row| row.get("Insert").or_else(|| row.get("Merge"))) + .filter_map(|row| row["span_id"].as_str().map(str::to_owned)) + .collect() + }; + assert_eq!( + identities(&rows(&path)), + identities(&completed), + "one-shot reimport changed native identities" + ); + assert_eq!( + identities(&before), + identities(&completed), + "reattach changed the open span identities" + ); + assert_eq!(inference.requests().len(), 1); +} diff --git a/bt-daemon/tests/muse_integration.rs b/bt-daemon/tests/muse_integration.rs new file mode 100644 index 00000000..06228499 --- /dev/null +++ b/bt-daemon/tests/muse_integration.rs @@ -0,0 +1,188 @@ +#![allow( + clippy::unwrap_used, + clippy::expect_used, + reason = "Integration test assertions fail immediately." +)] +mod support; + +use axum::http::StatusCode; +use serde_json::json; +use support::agent_process::AgentTestWorld; +use support::agents::MuseAgent; +use support::inference::{MockReply, OpenAiMock, OpenAiTurn}; +use support::ingest::IngestScenario; +use support::server::TestServer; + +fn muse_tool_call( + request: &support::inference::OpenAiRequest, + call_id: &str, + command: &str, +) -> MockReply { + let names = request.tool_names(); + if names.contains(&"muse") { + MockReply::response(OpenAiTurn::Events(vec![ + json!({"type":"response.created", "response":{"id":format!("resp_{call_id}")}}), + json!({"type":"response.function_call_arguments.done", "output_index":0, + "item_id":format!("fc_{call_id}"), "namespace":"muse", "name":"shell", "call_id":call_id, + "arguments":json!({"command":command}).to_string()}), + json!({"type":"response.completed", "response":{"id":format!("resp_{call_id}"), "status":"completed", + "output":[], "usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}}), + ])) + } else { + let name = ["shell", "bash"] + .into_iter() + .find(|name| names.contains(name)) + .unwrap_or_else(|| panic!("no shell tool: {names:?}")); + MockReply::response(OpenAiTurn::tool_call( + call_id, + name, + json!({"command": command, "description": "Print integration marker"}), + )) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires Muse Code 1.3+ installed; uses loopback inference and ingest"] +async fn muse_session_emits_traces() { + let inference = OpenAiMock::new(|context, request| { + assert_eq!(request.model(), Some("mock-model")); + match context.request_index { + 0 => { + assert!(request.contains_text("MUSE_TOOL_OK")); + muse_tool_call(&request, "call_muse_1", "printf MUSE_TOOL_OK") + } + 1 => { + assert!( + request.has_function_output("call_muse_1"), + "{}", + request.body + ); + MockReply::response(OpenAiTurn::text("MUSE_MOCK_OK")) + } + 2 => MockReply::http_error( + StatusCode::BAD_REQUEST, + json!({"error": { + "type": "invalid_request_error", "message": "deterministic Muse inference failure" + }}), + ), + 3 => MockReply::http_error( + StatusCode::SERVICE_UNAVAILABLE, + json!({"error": { + "type":"server_error", "message":"transient Muse inference failure" + }}), + ), + 4 => MockReply::response(OpenAiTurn::text("MUSE_RETRY_OK")), + 5 => muse_tool_call( + &request, + "call_muse_failure", + "printf MUSE_TOOL_FAILURE; exit 7", + ), + 6 => MockReply::response(OpenAiTurn::text("MUSE_FAILURE_HANDLED")), + n => panic!( + "unexpected Muse inference request {n}; model {:?}", + request.model() + ), + } + }); + let server = TestServer::start(inference.router()).await; + let world = AgentTestWorld::start().await; + assert!( + world.uses_mock_inference() && world.uses_mock_ingest(), + "Muse test requires mock backends" + ); + let muse = MuseAgent::new(&world, server.uri()); + let output = muse + .run( + &world, + "Run the shell command printf MUSE_TOOL_OK, then reply MUSE_MOCK_OK.", + ) + .await; + let export = muse.export_latest(&world).await; + assert_eq!(export["export_schema_version"], 1); + output.assert_success(); + output.assert_contains("MUSE_MOCK_OK"); + assert_eq!(inference.requests().len(), 2); + assert!( + inference.requests()[1].has_function_output("call_muse_1"), + "Muse did not return shell output" + ); + let failed = muse + .run(&world, "Trigger the deterministic inference error.") + .await; + failed.assert_failure(); + failed.assert_contains("deterministic Muse inference failure"); + assert_eq!(inference.requests().len(), 3); + let retried = muse.run(&world, "Reply after the transient failure.").await; + retried.assert_success(); + retried.assert_contains("MUSE_RETRY_OK"); + assert_eq!(inference.requests().len(), 5); + let tool_failed = muse + .run(&world, "Run a shell command that fails with exit 7.") + .await; + tool_failed.assert_success(); + tool_failed.assert_contains("MUSE_FAILURE_HANDLED"); + assert_eq!(inference.requests().len(), 7); + assert!(inference.requests()[6].has_function_output("call_muse_failure")); + let scenario = IngestScenario::new() + .expect("Muse version and provenance", |row| { + row["metadata"]["source"] == "muse" + && row["metadata"]["session_id"].is_string() + && row["metadata"]["muse_version"].is_string() + }) + .expect("Muse tool completion", |row| { + row["span_attributes"]["type"] == "tool" + && row["output"] + .as_str() + .is_some_and(|text| text.contains("MUSE_TOOL_OK")) + && row["error"].is_null() + && row["metrics"]["end"].is_number() + }) + .expect("Muse final LLM", |row| { + row["span_attributes"]["type"] == "llm" && row.to_string().contains("MUSE_MOCK_OK") + }) + .expect("Muse terminal session output", |row| { + row["span_attributes"]["name"] == "Muse Code" + && row["output"] == "MUSE_MOCK_OK" + && row["metrics"]["end"].is_number() + }) + .expect("Muse inference failure", |row| { + row.to_string() + .contains("deterministic Muse inference failure") + && row["error"].is_string() + }); + let scenario = scenario + .expect("Muse retries complete", |row| { + row["span_attributes"]["type"] == "llm" && row.to_string().contains("MUSE_RETRY_OK") + }) + .expect("Muse tool failures retain outcome", |row| { + row["span_attributes"]["type"] == "tool" + && row["metadata"]["tool_outcome"] == "error" + && row["error"].is_string() + }); + let rows = world.wait_for_mock_ingest_scenario(&scenario).await; + let root = rows + .iter() + .find(|row| { + row["span_attributes"]["name"] == "Muse Code" + && row["output"] == "MUSE_MOCK_OK" + && row["metrics"]["end"].is_number() + }) + .expect("successful session has authoritative terminal output"); + assert_eq!(root["metrics"]["prompt_tokens"], 20.0); + assert_eq!(root["metrics"]["completion_tokens"], 10.0); + assert_eq!(root["metrics"]["tokens"], 30.0); + let llms: std::collections::HashSet<_> = rows + .iter() + .filter(|row| { + row["span_attributes"]["type"] == "llm" && row["root_span_id"] == root["span_id"] + }) + .map(|row| row["span_id"].as_str().unwrap()) + .collect(); + assert_eq!(llms.len(), 2, "one LLM span per native model step"); + assert!( + rows.iter() + .filter(|row| row["root_span_id"] == root["span_id"]) + .all(|row| row["error"].is_null()), + "successful tool_calls must not become LLM errors" + ); +} diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index ec7265fd..81a6164c 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -2817,15 +2817,22 @@ esac #[cfg(all(feature = "cli", unix))] #[tokio::test] -#[ignore = "requires a local Muse Code 1.1.1 installation"] +#[ignore = "requires a local Muse Code 1.3+ installation"] async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { - let output = match std::process::Command::new("muse").arg("--version").output() { - Ok(output) if output.status.success() => output, - _ => return, - }; - if !String::from_utf8_lossy(&output.stdout).contains("1.1.1") { - return; - } + let _credential_backend = EnvVarGuard::set("TBH_CREDENTIAL_BACKEND", "file"); + let _telemetry = EnvVarGuard::set("TBH_DISABLE_TELEMETRY", "1"); + let muse_executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()); + let output = std::process::Command::new(&muse_executable) + .arg("--version") + .output() + .expect("install Muse Code before explicitly running this ignored test"); + assert!(output.status.success(), "Muse version check failed"); + let native_version = String::from_utf8(output.stdout) + .unwrap() + .split_whitespace() + .nth(2) + .unwrap() + .to_owned(); let tmp = tempfile::tempdir().unwrap(); let config_dir = tmp.path().join("muse"); std::fs::create_dir_all(&config_dir).unwrap(); @@ -2833,12 +2840,16 @@ async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { let foreign_command = format!("printf foreign >> '{}'", marker.display()); let managed = serde_json::json!({ "schema_version": 1, - "hooks": {"SessionStart": [{"matcher": "*", "hooks": [{"type": "command", "command": foreign_command}]}]} + "hooks": {"SessionStart": [{"hooks": [{"type": "command", "name": "enterprise-audit", "timeout_ms": 1000, "command": foreign_command}]}]} }); std::fs::write(config_dir.join("global.json"), managed.to_string()).unwrap(); - let settings = r#"{"schema_version":1,"managed_hooks_path":"global.json"}"#; - std::fs::write(config_dir.join("settings.json"), settings).unwrap(); + let user_marker = tmp.path().join("user-hook.txt"); + let settings = serde_json::json!({"schema_version":1,"managed_hooks_path":"global.json","hooks":{"Stop":[{"hooks":[{"type":"command","command":format!("printf user >> '{}'", user_marker.display())}]}]}}); + std::fs::write(config_dir.join("settings.json"), settings.to_string()).unwrap(); let _config = EnvVarGuard::set("XDG_CONFIG_HOME", tmp.path()); + let _data = EnvVarGuard::set("XDG_DATA_HOME", tmp.path().join("data")); + let workspace = tmp.path().join("workspace"); + std::fs::create_dir(&workspace).unwrap(); let (socket, daemon, flushes, daemon_dir) = start_tracking_daemon(env!("CARGO_PKG_VERSION")).await; let _socket = EnvVarGuard::set("BT_DAEMON_SOCKET", &socket); @@ -2866,6 +2877,8 @@ async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { "--provider".into(), "echo".into(), "--no-foreign-personal-context".into(), + "--workspace".into(), + workspace.clone().into_os_string(), "hello".into(), ], }, @@ -2883,9 +2896,10 @@ async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { .unwrap(); assert!(status.success()); assert_eq!(std::fs::read_to_string(&marker).unwrap(), "foreign"); + assert_eq!(std::fs::read_to_string(&user_marker).unwrap(), "user"); let saved: serde_json::Value = serde_json::from_slice(&std::fs::read(config_dir.join("settings.json")).unwrap()).unwrap(); - assert_eq!(saved["managed_hooks_path"], "global.json"); + assert_eq!(saved, settings); assert_eq!( std::fs::read_to_string(config_dir.join("global.json")).unwrap(), managed.to_string() @@ -2905,14 +2919,17 @@ async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { "muse-managed-test" ); assert!(row["managed_run_id"].as_str().is_some()); + assert_eq!(row["source_version"], native_version); main_session |= row["event"] == "SessionStart"; } } assert!(main_session); assert!(flushed.values().copied().sum::() >= 1); - let ordinary = std::process::Command::new("muse") + let ordinary = std::process::Command::new(&muse_executable) + .arg("exec") + .arg("--workspace") + .arg(&workspace) .args([ - "exec", "--provider", "echo", "--no-foreign-personal-context", @@ -2922,23 +2939,28 @@ async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { .unwrap(); assert!(ordinary.status.success()); assert_eq!(std::fs::read_to_string(&marker).unwrap(), "foreignforeign"); + assert_eq!(std::fs::read_to_string(&user_marker).unwrap(), "useruser"); shutdown(&socket).await; daemon.await.unwrap(); } #[cfg(all(feature = "cli", unix))] #[tokio::test] -#[ignore = "requires a local Muse Code 1.1.1 installation"] +#[ignore = "requires a local Muse Code 1.3+ installation"] async fn muse_managed_run_starts_its_isolated_daemon_before_sanitized_hooks() { - let output = match std::process::Command::new("muse").arg("--version").output() { - Ok(output) if output.status.success() => output, - _ => return, - }; - if !String::from_utf8_lossy(&output.stdout).contains("1.1.1") { - return; - } + let _credential_backend = EnvVarGuard::set("TBH_CREDENTIAL_BACKEND", "file"); + let _telemetry = EnvVarGuard::set("TBH_DISABLE_TELEMETRY", "1"); + let muse_executable = std::env::var_os("MUSE_BIN").unwrap_or_else(|| "muse".into()); + let output = std::process::Command::new(&muse_executable) + .arg("--version") + .output() + .expect("install Muse Code before explicitly running this ignored test"); + assert!(output.status.success(), "Muse version check failed"); let tmp = tempfile::tempdir().unwrap(); let _config = EnvVarGuard::set("XDG_CONFIG_HOME", tmp.path()); + let _data = EnvVarGuard::set("XDG_DATA_HOME", tmp.path().join("data")); + let workspace = tmp.path().join("workspace"); + std::fs::create_dir(&workspace).unwrap(); let _socket = EnvVarGuard::unset("BT_DAEMON_SOCKET"); let binary = env!("CARGO_BIN_EXE_bt-daemon"); let status = run_traced( @@ -2952,6 +2974,8 @@ async fn muse_managed_run_starts_its_isolated_daemon_before_sanitized_hooks() { "--provider".into(), "echo".into(), "--no-foreign-personal-context".into(), + "--workspace".into(), + workspace.clone().into_os_string(), "hello".into(), ], }, diff --git a/bt-daemon/tests/support/agents/mod.rs b/bt-daemon/tests/support/agents/mod.rs index 7cee545a..a8c6b024 100644 --- a/bt-daemon/tests/support/agents/mod.rs +++ b/bt-daemon/tests/support/agents/mod.rs @@ -1,5 +1,6 @@ mod claude; mod codex; +mod muse; mod opencode; mod pi; mod test_plugin; @@ -9,6 +10,8 @@ pub use claude::{ClaudeAgent, ClaudeRun}; #[allow(unused_imports)] pub use codex::{CodexAgent, CodexRun}; #[allow(unused_imports)] +pub use muse::MuseAgent; +#[allow(unused_imports)] pub use opencode::{OpenCodeAgent, OpenCodeRun}; #[allow(unused_imports)] pub use pi::{PiAgent, PiRun}; diff --git a/bt-daemon/tests/support/agents/muse.rs b/bt-daemon/tests/support/agents/muse.rs new file mode 100644 index 00000000..d3e90185 --- /dev/null +++ b/bt-daemon/tests/support/agents/muse.rs @@ -0,0 +1,120 @@ +use super::AgentOutput; +use crate::support::agent_process::AgentTestWorld; +use serde_json::json; +use std::path::PathBuf; + +/// Native hooks ship with bt; the CLI installs them for this invocation. +/// The seeded endpoint/auth shape follows Muse's official sdk-quickstart fake +/// provider, with dummy credentials and isolated configuration/data directories. +pub struct MuseAgent { + config_home: PathBuf, + data_home: PathBuf, +} + +impl MuseAgent { + pub fn new(world: &AgentTestWorld, base_url: &str) -> Self { + let config_home = world.temp_path("muse-config"); + let config = config_home.join("muse"); + std::fs::create_dir_all(&config).expect("create Muse config"); + std::fs::create_dir_all(world.temp_path("muse-home")).expect("create Muse home"); + std::fs::write( + config.join("settings.json"), + json!({ + "schema_version": 1, + "endpoint_transport": {"base_url": base_url, "auth": "bearer"} + }) + .to_string(), + ) + .expect("write Muse settings"); + std::fs::write( + config.join("auth.json"), + json!({ + "schema_version": 1, + "providers": {"meta": {"api_key": "test-key"}} + }) + .to_string(), + ) + .expect("write dummy Muse auth"); + Self { + config_home, + data_home: world.temp_path("muse-data"), + } + } + + pub async fn export_latest(&self, world: &AgentTestWorld) -> serde_json::Value { + let path = world.temp_path("muse-export.json"); + let mut command = super::command_from_env("MUSE_BIN", "muse"); + command + .args(["export", "--last", "--out"]) + .arg(&path) + .current_dir(world.workspace()) + .env_remove("CODEX_HOME") + .env("HOME", world.temp_path("muse-home")) + .env("TBH_CREDENTIAL_BACKEND", "file") + .env("TBH_DISABLE_TELEMETRY", "1") + .env("META_API_KEY", "test-key") + .env("XDG_CONFIG_HOME", &self.config_home) + .env("XDG_DATA_HOME", &self.data_home); + let output = world.output(&mut command).await; + assert!( + output.status.success(), + "Muse export: {}", + String::from_utf8_lossy(&output.stderr) + ); + serde_json::from_slice(&std::fs::read(path).expect("read native export")) + .expect("parse native export") + } + + pub async fn run(&self, world: &AgentTestWorld, prompt: &str) -> AgentOutput { + let mut command = tokio::process::Command::new(env!("CARGO_BIN_EXE_bt-daemon")); + command + .args([ + "run", + "--project", + "agent-e2e", + "muse", + "exec", + "--provider", + "meta", + "--model", + "mock-model", + "--no-foreign-personal-context", + "--disable-reminders", + "--disable-web-tools", + "--disable-request-user-input", + "--enable-shell-tool", + "--max-model-steps", + "3", + "--disable-approval", + "--disable-sandbox", + prompt, + ]) + .current_dir(world.workspace()) + .env_remove("CODEX_HOME") + .env("HOME", world.temp_path("muse-home")) + .env("TBH_CREDENTIAL_BACKEND", "file") + .env("TBH_DISABLE_TELEMETRY", "1") + .env("META_API_KEY", "test-key") + .env("XDG_CONFIG_HOME", &self.config_home) + .env("XDG_DATA_HOME", &self.data_home); + // Do not import the parent coding agent's injected tool runtime. + for (key, _) in std::env::vars_os() { + let name = key.to_string_lossy(); + if [ + "CODEX_", + "NODE_REPL_", + "CUA_REPL_", + "BROWSER_USE_", + "SKY_CUA_", + ] + .iter() + .any(|prefix| name.starts_with(prefix)) + { + command.env_remove(key); + } + } + command.env_remove("ELECTRON_RUN_AS_NODE"); + world.configure(&mut command); + world.output(&mut command).await.into() + } +} diff --git a/bt-daemon/tests/support/inference/openai.rs b/bt-daemon/tests/support/inference/openai.rs index 69ea7f95..c9b959a9 100644 --- a/bt-daemon/tests/support/inference/openai.rs +++ b/bt-daemon/tests/support/inference/openai.rs @@ -253,6 +253,8 @@ impl OpenAiMock { pub fn router(&self) -> Router { Router::new() .route("/v1/models", get(models)) + .route("/muse-code/models", get(models)) + .route("/responses", post(muse_responses)) .route("/v1/responses", post(responses)) .route("/v1/chat/completions", post(chat_completions)) .route("/backend-api/plugins/featured", get(featured_plugins)) @@ -267,7 +269,9 @@ impl OpenAiMock { async fn models() -> axum::Json { axum::Json(json!({ "object": "list", - "data": [{"id": "mock-model", "object": "model", "owned_by": "mock"}] + "data": [{"id": "mock-model", "object": "model", "owned_by": "mock", + "metadata": {"muse-code": {"release_date": "2026-01-01", "is_hidden": false, + "limit": {"context": 1000000, "output": 1024}}}}] })) } @@ -275,10 +279,29 @@ async fn featured_plugins() -> axum::Json { axum::Json(json!([])) } +// Muse validates the complete Responses envelope, including sequence numbers. +// Other agents intentionally also exercise the sparse mock response path. +async fn muse_responses( + State(state): State>, + headers: HeaderMap, + body: Bytes, +) -> axum::response::Response { + respond(state, headers, body, true).await +} + async fn responses( State(state): State>, headers: HeaderMap, body: Bytes, +) -> axum::response::Response { + respond(state, headers, body, false).await +} + +async fn respond( + state: Arc, + headers: HeaderMap, + body: Bytes, + complete: bool, ) -> axum::response::Response { let body = match decode_json_body(&headers, &body) { Ok(body) => body, @@ -297,11 +320,27 @@ async fn responses( }, request, ) { - MockReply::Response(turn) => raw_response( - StatusCode::OK, - "text/event-stream", - sse(&turn.events(index)), - ), + MockReply::Response(turn) => { + let mut events = turn.events(index); + if complete { + for (sequence, event) in events.iter_mut().enumerate() { + event["sequence_number"] = json!(sequence + 1); + let completed = event["type"] == "response.completed"; + if let Some(response) = event.get_mut("response").and_then(Value::as_object_mut) + { + response.entry("object").or_insert(json!("response")); + response.entry("model").or_insert(json!("mock-model")); + response.entry("status").or_insert(json!(if completed { + "completed" + } else { + "in_progress" + })); + response.entry("output").or_insert(json!([])); + } + } + } + raw_response(StatusCode::OK, "text/event-stream", sse(&events)) + } MockReply::HttpError { status, body } => json_response(status, body), MockReply::Raw { status,