diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4b3a66e0..fbf65f51 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -371,6 +371,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 "BRAINTRUST_MUSE_BIN=$MUSE_INSTALL_DIR/muse" >> "$GITHUB_ENV" + - name: Install Rust + run: | + rustup toolchain install 1.97.1 --profile minimal + rustup default 1.97.1 + - 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: @@ -381,6 +417,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/Cargo.lock b/bt-daemon/Cargo.lock index b061e11d..003148fd 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -307,6 +307,7 @@ dependencies = [ "clap", "fs2", "libc", + "muse-code-sdk", "parking_lot", "regex", "reqwest 0.12.28", @@ -359,6 +360,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" + [[package]] name = "cfg_aliases" version = "0.2.2" @@ -1351,6 +1358,31 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "muse-code-sdk" +version = "1.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95a9af6f1f1b8d7f706c3564f4ba7269c6d5a9a2030a1fb0461f69b9f3551126" +dependencies = [ + "nix", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", +] + +[[package]] +name = "nix" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab2156c4fce2f8df6c499cc1c763e4394b7482525bf2a9701c9d79d215f519e4" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases 0.1.1", + "libc", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -1491,7 +1523,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" dependencies = [ "bytes", - "cfg_aliases", + "cfg_aliases 0.2.2", "pin-project-lite", "quinn-proto", "quinn-udp", @@ -1533,7 +1565,7 @@ version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ - "cfg_aliases", + "cfg_aliases 0.2.2", "libc", "once_cell", "socket2", diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index b185ac05..0ecf9266 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -2,6 +2,7 @@ name = "bt-daemon" version = "0.1.0" edition = "2021" +rust-version = "1.97" license = "MIT" publish = false description = "Embeddable Braintrust coding-agent tracing daemon." @@ -25,6 +26,7 @@ expect_used = "deny" braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "45a2c88ca9d9ccc393ec0338ac438199fc3c5df2" } base64 = "0.22" anyhow = "1" +muse-code-sdk = "=1.4.12" parking_lot = "0.12" async-trait = "0.1" chrono = "0.4" @@ -34,7 +36,7 @@ regex = "1" rquickjs = "0.12.2" rquickjs-serde = "0.6.1" serde = { version = "1", features = ["derive", "rc"] } -serde_json = "1" +serde_json = { version = "1", features = ["raw_value"] } semver = "1" sha2 = "0.10" sysinfo = { version = "0.33.1", default-features = false, features = ["system"] } diff --git a/bt-daemon/README.md b/bt-daemon/README.md index 993d23eb..0a1c5127 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` @@ -281,15 +283,21 @@ 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. -Add `--attach` to keep following an active Codex, Claude, or Antigravity transcript until -Ctrl-C. `run [ARGS...]` launches the selected agent with +Muse Code 1.3+ supports live tracing, managed runs, session import, and attach +on macOS and Linux. Use `import muse --all` to import completed sessions. + +Add `--attach` to keep following an active Codex, Claude, Antigravity, or Muse +session until Ctrl-C. Detaching from Muse leaves the session running; attach +again or import it later to collect its completion. + +`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/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/delivery_ledger.rs b/bt-daemon/src/delivery_ledger.rs index 15952221..4c3a1de0 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, WeakOrigins}; +use crate::translate::{LateMergeVersion, OriginSnapshot, SpanOp, SpanRow, WeakOrigins}; mod origins; use crate::wire::SessionConfig; use origins::OriginIndex; @@ -34,6 +34,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. @@ -48,6 +50,8 @@ struct DeliveryLedger { // Keep only SDK deliveries awaiting a durable acknowledgement in memory. pending: HashSet, pending_late_merges: HashSet, + known_versions: HashMap, + pending_versions: HashMap, } impl DeliveryLedger { @@ -84,6 +88,8 @@ impl DeliveryLedger { journal_len: 0, pending: HashSet::new(), pending_late_merges: HashSet::new(), + known_versions: persisted.late_merge_versions, + pending_versions: HashMap::new(), }; let lock = ledger.lock().await?; ledger @@ -171,6 +177,9 @@ impl DeliveryLedger { self.origins .remember_completed(update.completed_span_ids, update.late_merge_ids) .await?; + 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 @@ -209,6 +218,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()); } @@ -229,7 +241,10 @@ impl DeliveryLedger { self.pending.retain(|id| !known.contains(id)); self.pending_late_merges .retain(|id| !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 @@ -238,6 +253,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, @@ -245,6 +261,9 @@ impl DeliveryLedger { .await?; self.pending.clear(); self.pending_late_merges.clear(); + for (span, version) in self.pending_versions.drain() { + remember_version(&mut self.known_versions, &span, &version); + } Ok(()) } @@ -285,6 +304,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) } @@ -474,7 +513,23 @@ impl Sink for LedgerSink { true } }); - filtered.retain(|op| ledger.should_emit(op, &known, &known_late_merges)); + let mut versions = ledger.known_versions.clone(); + for (span, version) in &ledger.pending_versions { + remember_version(&mut versions, span, version); + } + filtered.retain(|op| { + if !ledger.should_emit(op, &known, &known_late_merges) { + return false; + } + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op.as_ref(); + if let Some(version) = &row.late_merge_version { + if !version_allowed(versions.get(&row.span_id), version) { + return false; + } + remember_version(&mut versions, &row.span_id, version); + } + true + }); Some(lock) } else { None @@ -1111,7 +1166,7 @@ mod tests { assert_eq!(stored["span-1"].plugin_version.as_deref(), Some("v2")); assert_eq!(recovered.emit(&[terminal("span-0")]).await.unwrap(), 0); // Simulate upgrading a JSON checkpoint with versions and completion IDs. - let (_ledger, origins) = DeliveryLedger::load(temp.path(), "codex", "session-1", &config) + let (ledger, origins) = DeliveryLedger::load(temp.path(), "codex", "session-1", &config) .await .unwrap(); tokio::fs::write( @@ -1120,6 +1175,7 @@ mod tests { completed_span_ids: (0..256).map(|i| format!("span-{i}")).collect(), late_merge_ids: HashSet::new(), span_origins: origins, + late_merge_versions: ledger.known_versions, }) .unwrap(), ) @@ -1325,6 +1381,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 099ffd0e..9f5feb35 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,24 @@ 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(); + if args.source_version.is_none() { + args.source_version = context.source_version.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,18 +229,53 @@ 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(()) } 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); @@ -438,6 +494,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)] @@ -563,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/lib.rs b/bt-daemon/src/lib.rs index 7fb4b7cf..dd42434d 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -26,6 +26,7 @@ mod hook_executable; mod ids; mod journal; mod managed_run; +mod muse; mod pipeline; mod plugin_diagnostics; pub(crate) mod process; @@ -61,7 +62,7 @@ pub use hook::run_hook; pub use hook_executable::delegate_hook_executable; #[doc(hidden)] pub use journal::source_journal_path; -pub use managed_run::run_traced; +pub use managed_run::{run_traced, run_traced_with_host}; pub use plugin_diagnostics::PluginDiagnostic; pub use server::{ braintrust_serve_options, debug_serve_options, run_serve, AuthLease, AuthProvider, diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index 2af7390b..df7fe46c 100644 --- a/bt-daemon/src/main.rs +++ b/bt-daemon/src/main.rs @@ -7,10 +7,10 @@ use async_trait::async_trait; use bt_daemon::wire::{AuthSelection, AuthSource, BackendAuth, SessionRoute, TraceDestination}; use bt_daemon::{ - braintrust_serve_options, paths, run_hook, run_import, run_serve, run_status, run_traced, - AuthLease, AuthProvider, AuthResolveReason, BraintrustSinkConfig, DebugSinkFactory, HookArgs, - HostInfo, ImportArgs, OutputFormat, Registry, RunArgs, RunHookCommand, ServeArgs, ServeOptions, - StatusArgs, TraceCommandOutput, + braintrust_serve_options, paths, run_hook, run_import, run_serve, run_status, + run_traced_with_host, AuthLease, AuthProvider, AuthResolveReason, BraintrustSinkConfig, + DebugSinkFactory, HookArgs, HostInfo, ImportArgs, OutputFormat, Registry, RunArgs, + RunHookCommand, ServeArgs, ServeOptions, StatusArgs, TraceCommandOutput, }; use clap::{Args, Parser, Subcommand}; use std::ffi::OsString; @@ -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_with_host(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/managed_run/mod.rs b/bt-daemon/src/managed_run/mod.rs index f51b9703..66d87487 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,31 @@ pub(crate) trait ManagedRun: Agent { hook_command: &RunHookCommand, 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 { + 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 +88,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; @@ -86,9 +131,24 @@ pub async fn run_traced( } let (executable_env, default_executable) = agent.executable(); let executable = - std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); + env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); let managed_run_id = uuid::Uuid::new_v4().to_string(); - let mut 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 = (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 mut injection = + agent.inject_with_context(&hook_command, &managed_run_id, &route, &socket)?; if args.source == RunSource::Codex { let executable_for_trust = std::env::var_os("CODEX_BIN").unwrap_or_else(|| OsString::from("codex")); @@ -110,20 +170,36 @@ pub async fn run_traced( injection.args.push(OsString::from("-c")); injection.args.push(trust_override); } - // 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 - && env::var_os(paths::SOCKET_ENV).is_none(); - let isolated_runtime = isolate_daemon - .then(|| ManagedRunRuntime::new(&managed_run_id)) - .transpose()?; + if agent.needs_daemon_host() { + if let Some(host) = host.as_ref() { + // Sanitized hooks cannot spawn a daemon. Prepare the resolved + // endpoint even when the caller supplied a fresh explicit socket. + drop( + crate::client::ensure_daemon_in( + &socket, + host, + false, + isolated_runtime + .as_ref() + .map(|runtime| runtime.temp_dir.path()), + ) + .await?, + ); + } else { + // Existing callers without host information may use an already + // listening explicit daemon, but must never launch untraceable work. + drop(crate::transport::connect(&socket).await.map_err(|_| { + anyhow::anyhow!( + "Muse managed runs require run_traced_with_host to start a daemon at {}", + socket.display() + ) + })?); + } + } 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); @@ -134,9 +210,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); @@ -260,6 +345,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..b091786b --- /dev/null +++ b/bt-daemon/src/managed_run/muse.rs @@ -0,0 +1,183 @@ +//! 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 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, + #[serde(default)] + pub source_version: Option, +} + +impl ManagedRun for Muse { + fn executable(&self) -> (&'static str, &'static str) { + ("BRAINTRUST_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)?; + } + } + 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 + } + + 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(), + 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)?; + 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 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 { + 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/muse/config/hooks.rs b/bt-daemon/src/muse/config/hooks.rs new file mode 100644 index 00000000..7e869d18 --- /dev/null +++ b/bt-daemon/src/muse/config/hooks.rs @@ -0,0 +1,217 @@ +//! Known hook structure is typed; extension data and future handler formats are +//! preserved as opaque JSON. No settings orchestration navigates these values. +use super::{HOOK_COMMAND, LEGACY_HOOK_COMMAND}; +use crate::muse::field::Field; +use crate::muse::hooks::HookName; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Value; +use std::collections::BTreeMap; + +type Extensions = BTreeMap; + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct HookOverlay(BTreeMap>); + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +struct HookGroup { + hooks: Vec, + #[serde(default, skip_serializing_if = "Field::is_missing")] + matcher: Field, + #[serde(flatten)] + extra: Extensions, +} +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(untagged)] +enum HookHandler { + User(UserCommand), + Enterprise(EnterpriseCommand), + Unknown(Value), +} +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(super) enum CommandType { + #[serde(rename = "command")] + Command, +} +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +struct UserCommand { + #[serde(rename = "type")] + kind: CommandType, + command: String, + #[serde(flatten)] + extra: Extensions, +} +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +struct EnterpriseCommand { + command: String, + #[serde(default, skip_serializing_if = "Field::is_missing")] + name: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + timeout_ms: Field, + #[serde(flatten)] + extra: Extensions, +} +impl<'de> Deserialize<'de> for HookHandler { + fn deserialize>(d: D) -> Result { + // The opaque handler boundary selects a native schema. A malformed + // recognized command must not silently pass as a future handler kind. + let value = Value::deserialize(d)?; + if value.get("type").and_then(Value::as_str) == Some("command") { + serde_json::from_value(value) + .map(Self::User) + .map_err(serde::de::Error::custom) + } else if value.get("type").is_none() && value.get("command").is_some_and(Value::is_string) + { + serde_json::from_value(value) + .map(Self::Enterprise) + .map_err(serde::de::Error::custom) + } else { + Ok(Self::Unknown(value)) + } + } +} +impl HookHandler { + fn is_owned(&self) -> bool { + matches!(self, Self::User(handler) if handler.extra.is_empty() + && matches!(handler.command.as_str(), HOOK_COMMAND | LEGACY_HOOK_COMMAND)) + } + fn contains_braintrust(&self) -> bool { + match self { + Self::User(h) => is_forwarder(&h.command) || extensions_contain(&h.extra), + Self::Enterprise(h) => is_forwarder(&h.command) || extensions_contain(&h.extra), + Self::Unknown(raw) => opaque_contains_braintrust(raw), + } + } +} +impl HookOverlay { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + pub fn contains_braintrust_hook(&self) -> bool { + self.0.values().flatten().any(|g| { + extensions_contain(&g.extra) || g.hooks.iter().any(HookHandler::contains_braintrust) + }) + } + pub fn remove_owned(&mut self) { + self.0.retain(|_, groups| { + let original_len = groups.len(); + groups.retain_mut(|group| { + let original = group.hooks.len(); + group.hooks.retain(|h| !h.is_owned()); + // Foreign group attributes remain even if its only handler was ours. + original == group.hooks.len() + || !group.hooks.is_empty() + || !group.extra.is_empty() + || match &group.matcher { + Field::Missing => false, + Field::Value(value) => value != "*", + Field::Null => true, + } + }); + original_len == groups.len() || !groups.is_empty() + }); + } + pub fn append_owned(&mut self, command: &str) { + for event in HookName::ALL { + self.0 + .entry(event.as_str().into()) + .or_default() + .push(HookGroup { + matcher: Field::Missing, + extra: BTreeMap::new(), + hooks: vec![HookHandler::User(UserCommand { + kind: CommandType::Command, + command: command.into(), + extra: BTreeMap::new(), + })], + }); + } + } +} +fn is_forwarder(command: &str) -> bool { + command.contains("--source") && command.contains("muse") && command.contains("hook") +} +fn extensions_contain(extra: &Extensions) -> bool { + extra.iter().any(|(key, value)| { + (key == "command" && value.as_str().is_some_and(is_forwarder)) + || opaque_contains_braintrust(value) + }) +} +/// Conservative compatibility scan restricted to unknown extension/handler data. +pub(super) fn opaque_contains_braintrust(value: &Value) -> bool { + match value { + Value::Object(object) => object.iter().any(|(key, value)| { + (key == "command" && value.as_str().is_some_and(is_forwarder)) + || opaque_contains_braintrust(value) + }), + Value::Array(items) => items.iter().any(opaque_contains_braintrust), + _ => false, + } +} + +/// Read-only view of known project/enterprise files. Other root schemas are +/// retained by the inspection codec rather than being rewritten or normalized. +#[derive(Deserialize)] +pub(super) struct HookDocument { + #[serde(default)] + hooks: HookOverlay, + #[serde(flatten)] + extra: Extensions, +} +impl HookDocument { + pub fn contains_braintrust(&self) -> bool { + self.hooks.contains_braintrust_hook() || extensions_contain(&self.extra) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn foreign_handlers_and_groups_survive_install_remove_cycle() { + let original = json!({"FutureEvent":[{"matcher":null,"extension":{"nested":true},"hooks":[ + {"type":"future","payload":{"x":1}}, + {"name":"enterprise","timeout_ms":1000,"command":"echo policy","extension":true}, + {"type":"command","command":HOOK_COMMAND,"timeout":42} + ]}],"FutureEmpty":[]}); + let mut hooks: HookOverlay = serde_json::from_value(original.clone()).unwrap(); + hooks.append_owned(HOOK_COMMAND); + hooks.remove_owned(); + assert_eq!(serde_json::to_value(hooks).unwrap(), original); + } + #[test] + fn removing_our_handler_keeps_foreign_group_attributes() { + for group in [ + json!({"hooks":[{"type":"command","command":HOOK_COMMAND}],"future":true}), + json!({"hooks":[{"type":"command","command":HOOK_COMMAND}],"matcher":null}), + json!({"hooks":[{"type":"command","command":HOOK_COMMAND}],"matcher":"custom"}), + ] { + let mut hooks: HookOverlay = + serde_json::from_value(json!({"Stop":[group.clone()]})).unwrap(); + hooks.remove_owned(); + let mut expected = group; + expected["hooks"] = json!([]); + assert_eq!( + serde_json::to_value(hooks).unwrap(), + json!({"Stop":[expected]}) + ); + } + } + #[test] + fn conservative_forwarder_checks_cover_unknown_extensions() { + for handler in [ + json!({"type":"future","nested":{"command":HOOK_COMMAND}}), + json!({"type":"command","command":"echo unrelated","extension":{"command":HOOK_COMMAND}}), + ] { + let mut hooks: HookOverlay = + serde_json::from_value(json!({"Stop":[{"hooks":[handler]}]})).unwrap(); + hooks.remove_owned(); + assert!(hooks.contains_braintrust_hook()); + } + assert!(serde_json::from_value::( + json!({"Stop":[{"hooks":[{"type":"command","command":null}]}]}) + ) + .is_err()); + } +} diff --git a/bt-daemon/src/muse/config/mod.rs b/bt-daemon/src/muse/config/mod.rs new file mode 100644 index 00000000..d582727d --- /dev/null +++ b/bt-daemon/src/muse/config/mod.rs @@ -0,0 +1,13 @@ +//! Typed native settings, preservable hook trees, and strict owned-file detection. +mod hooks; +mod ownership; +mod settings; + +pub(crate) use hooks::HookOverlay; +#[cfg(test)] +pub(crate) use ownership::OwnedHookFile; +pub(crate) use ownership::{inspect_hook_file, HookFileInspection}; +pub(crate) use settings::MuseSettings; + +pub(crate) const HOOK_COMMAND: &str = "bt trace hook --source muse --flush-on-turn-end"; +const LEGACY_HOOK_COMMAND: &str = "bt trace hook --source muse --flush-on-turn-end || true"; diff --git a/bt-daemon/src/muse/config/ownership.rs b/bt-daemon/src/muse/config/ownership.rs new file mode 100644 index 00000000..c364844c --- /dev/null +++ b/bt-daemon/src/muse/config/ownership.rs @@ -0,0 +1,213 @@ +//! Strict owned-file recognition, separate from preservable foreign documents. +use super::hooks::{opaque_contains_braintrust, CommandType, HookDocument}; +use super::{HOOK_COMMAND, LEGACY_HOOK_COMMAND}; +use crate::muse::field::Field; +use crate::muse::hooks::HookName; +use anyhow::Context; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; +use std::path::Path; + +#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct OwnedHookFile { + schema_version: u32, + hooks: BTreeMap>, +} +#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct OwnedHookGroup { + #[serde(default, skip_serializing_if = "Field::is_missing")] + matcher: Field, + hooks: Vec, +} +#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct OwnedCommandHook { + #[serde(rename = "type")] + kind: CommandType, + command: String, +} +impl OwnedHookFile { + fn for_events(events: &[HookName], command: &str, matcher: bool) -> Self { + Self { + schema_version: 1, + hooks: events + .iter() + .map(|event| { + ( + *event, + vec![OwnedHookGroup { + matcher: if matcher { + Field::Value("*".into()) + } else { + Field::Missing + }, + hooks: vec![OwnedCommandHook { + kind: CommandType::Command, + command: command.into(), + }], + }], + ) + }) + .collect(), + } + } + pub fn current() -> Self { + Self::for_events(HookName::ALL, HOOK_COMMAND, false) + } + pub fn previous() -> Self { + Self::for_events( + &[ + HookName::SessionStart, + HookName::UserPromptSubmit, + HookName::PreLLMCall, + HookName::PostLLMCall, + HookName::Stop, + HookName::SessionEnd, + ], + HOOK_COMMAND, + true, + ) + } + fn legacy(matcher: bool) -> Self { + Self::for_events( + &[ + HookName::SessionStart, + HookName::UserPromptSubmit, + HookName::PreToolUse, + HookName::PermissionRequest, + HookName::PostToolUse, + HookName::PreLLMCall, + HookName::PostLLMCall, + HookName::PreCompact, + HookName::PostCompact, + HookName::SubagentStart, + HookName::SubagentStop, + HookName::Stop, + HookName::SessionEnd, + ], + LEGACY_HOOK_COMMAND, + matcher, + ) + } + fn is_recognized(&self) -> bool { + self == &Self::current() + || self == &Self::previous() + || self == &Self::legacy(true) + || self == &Self::legacy(false) + } +} +#[derive(Debug, PartialEq)] +pub(crate) enum HookFileInspection { + Missing, + Owned, + Foreign { contains_braintrust: bool }, +} +impl HookFileInspection { + pub fn is_owned(&self) -> bool { + matches!(self, Self::Owned) + } + pub fn contains_braintrust(&self) -> bool { + matches!( + self, + Self::Owned + | Self::Foreign { + contains_braintrust: true + } + ) + } +} +pub(crate) fn inspect_hook_file(path: &Path) -> anyhow::Result { + match std::fs::read(path) { + Ok(raw) => inspect(&raw) + .with_context(|| format!("read Muse hook configuration {}", path.display())), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + Ok(HookFileInspection::Missing) + } + Err(error) => { + Err(error).with_context(|| format!("read Muse hook configuration {}", path.display())) + } + } +} +fn inspect(raw: &[u8]) -> anyhow::Result { + // Schema mismatch means foreign, not invalid JSON. Validate syntax without + // allocating an untyped tree before trying the strict owned schema. + serde_json::from_slice::(raw)?; + if serde_json::from_slice::(raw).is_ok_and(|file| file.is_recognized()) { + return Ok(HookFileInspection::Owned); + } + let contains_braintrust = match serde_json::from_slice::(raw) { + Ok(document) => document.contains_braintrust(), + // Unknown root/group formats retain the conservative read-only scan. + Err(_) => opaque_contains_braintrust(&serde_json::from_slice(raw)?), + }; + Ok(HookFileInspection::Foreign { + contains_braintrust, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn recognize_all_exact_owned_versions() { + for file in [ + OwnedHookFile::current(), + OwnedHookFile::previous(), + OwnedHookFile::legacy(true), + OwnedHookFile::legacy(false), + ] { + assert_eq!( + inspect(&serde_json::to_vec(&file).unwrap()).unwrap(), + HookFileInspection::Owned + ); + } + } + #[test] + fn extensions_null_matchers_and_unknown_events_are_never_owned() { + let original = serde_json::to_value(OwnedHookFile::current()).unwrap(); + let mut variants = Vec::new(); + let mut file = original.clone(); + file["future"] = json!(true); + variants.push(file); + let mut file = original.clone(); + file["hooks"]["Stop"][0]["matcher"] = json!(null); + variants.push(file); + let mut file = original.clone(); + file["hooks"]["Stop"][0]["future"] = json!(true); + variants.push(file); + let mut file = original.clone(); + file["hooks"]["Stop"][0]["hooks"][0]["timeout"] = json!(1000); + variants.push(file); + let mut file = original.clone(); + file["hooks"]["FutureEvent"] = json!([]); + variants.push(file); + for file in variants { + assert_eq!( + inspect(&serde_json::to_vec(&file).unwrap()).unwrap(), + HookFileInspection::Foreign { + contains_braintrust: true + } + ); + } + assert!(inspect(b"{").is_err()); + assert_eq!( + inspect(b"[]").unwrap(), + HookFileInspection::Foreign { + contains_braintrust: false + } + ); + } + #[test] + fn unknown_file_schema_still_checks_nested_forwarders() { + let raw = serde_json::to_vec(&json!({"future_schema":[{"command":HOOK_COMMAND}]})).unwrap(); + assert_eq!( + inspect(&raw).unwrap(), + HookFileInspection::Foreign { + contains_braintrust: true + } + ); + } +} diff --git a/bt-daemon/src/muse/config/settings.rs b/bt-daemon/src/muse/config/settings.rs new file mode 100644 index 00000000..68327c0f --- /dev/null +++ b/bt-daemon/src/muse/config/settings.rs @@ -0,0 +1,97 @@ +//! Settings projection with lossless unrelated fields and explicit missing/null +//! semantics. The caller holds Muse's native lock while loading and writing. +use super::hooks::HookOverlay; +use crate::muse::field::Field; +use anyhow::Context; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +#[derive(Default, Serialize, Deserialize)] +pub(crate) struct MuseSettings { + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub schema_version: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub managed_hooks_path: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub hooks: Field, + #[serde(flatten)] + extra: BTreeMap, +} +impl MuseSettings { + pub fn load(path: &Path) -> anyhow::Result { + match std::fs::read(path) { + Ok(raw) => serde_json::from_slice(&raw) + .with_context(|| format!("read Muse settings {}", path.display())), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()), + Err(error) => { + Err(error).with_context(|| format!("read Muse settings {}", path.display())) + } + } + } + pub fn managed_path(&self, config_dir: &Path) -> Option { + self.managed_hooks_path.as_ref().map(|path| { + if path.is_absolute() { + path.clone() + } else { + config_dir.join(path) + } + }) + } + pub fn hook_overlay(&self) -> anyhow::Result { + match &self.hooks { + Field::Missing => Ok(HookOverlay::default()), + Field::Value(hooks) => Ok(hooks.clone()), + Field::Null => anyhow::bail!("Muse user hooks must be an object"), + } + } + pub fn ensure_schema_version(&mut self) { + if self.schema_version.is_missing() { + self.schema_version = Field::Value(1); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn settings_preserve_unrelated_data_and_missing_vs_null() { + for source in [ + json!({"theme":{"future":[true,null,2]},"managed_hooks_path":null,"schema_version":null}), + json!({"theme":"dark"}), + ] { + let settings: MuseSettings = serde_json::from_value(source.clone()).unwrap(); + assert_eq!(serde_json::to_value(settings).unwrap(), source); + } + let mut absent: MuseSettings = serde_json::from_value(json!({})).unwrap(); + absent.ensure_schema_version(); + assert_eq!( + serde_json::to_value(absent).unwrap(), + json!({"schema_version":1}) + ); + let mut null: MuseSettings = + serde_json::from_value(json!({"schema_version":null})).unwrap(); + null.ensure_schema_version(); + assert_eq!( + serde_json::to_value(null).unwrap(), + json!({"schema_version":null}) + ); + } + #[test] + fn managed_paths_resolve_relative_and_absolute_and_null_hooks_are_invalid() { + let base = tempfile::tempdir().unwrap(); + let relative: MuseSettings = + serde_json::from_value(json!({"managed_hooks_path":"enterprise.json"})).unwrap(); + assert_eq!( + relative.managed_path(base.path()).unwrap(), + base.path().join("enterprise.json") + ); + let absolute = base.path().join("absolute.json"); + let settings: MuseSettings = + serde_json::from_value(json!({"managed_hooks_path":absolute,"hooks":null})).unwrap(); + assert_eq!(settings.managed_path(base.path()).unwrap(), absolute); + assert!(settings.hook_overlay().is_err()); + } +} diff --git a/bt-daemon/src/muse/export_v1.rs b/bt-daemon/src/muse/export_v1.rs new file mode 100644 index 00000000..be0f01bc --- /dev/null +++ b/bt-daemon/src/muse/export_v1.rs @@ -0,0 +1,247 @@ +//! Durable export-v1 DTOs. This is a separate native schema from MSP views. +use super::hooks::{ErrorDetail, InputSummary, ToolArguments, ToolOutput}; +use super::usage::TokenUsage; +use serde::Deserialize; + +#[derive(Deserialize)] +pub(crate) struct Header { + pub export_schema_version: u64, + pub exporter_version: Option, + pub sessions: Vec, + #[serde(default)] + pub session_terminated_abnormally: bool, +} + +#[derive(Deserialize)] +#[serde(untagged)] +pub(crate) enum Version { + String(String), + Object { semver: String }, +} + +impl Version { + pub(crate) fn into_string(self) -> String { + match self { + Self::String(value) | Self::Object { semver: value } => value, + } + } +} + +#[derive(Clone, Deserialize)] +pub(crate) struct Session { + pub session_id: String, + pub trajectory_id: Option, + pub root_session_id: Option, + pub session_end: Option, +} + +#[derive(Clone, Deserialize)] +pub(crate) struct SessionEnd { + pub exit_reason: Option, +} + +/// Export timestamps are microseconds; daemon envelopes use milliseconds. +#[derive(Clone, Copy, Deserialize)] +#[serde(transparent)] +pub(crate) struct UnixMicros(i64); +impl UnixMicros { + pub fn millis(self) -> i64 { + self.0 / 1_000 + } +} +#[derive(Deserialize)] +pub(crate) struct ExportRecord { + pub recorded_at: Option, + pub envelope: Option, +} +#[derive(Deserialize)] +pub(crate) struct NativeEnvelope { + pub recorded_at: Option, + pub payload: Option, +} +pub(crate) enum NativePayload { + Metadata { record: Option }, + RunModel { record: Option }, + Run { run_id: String, event: RunEvent }, + Task { event: Option }, + SessionEnd, + Ignored, +} +// Some retained envelopes contain non-event payloads without a kind. Select +// the schema before decoding: unknown bodies stay opaque, malformed *known* +// bodies still fail. RawValue is confined to this codec and never enters state. +impl<'de> Deserialize<'de> for NativePayload { + fn deserialize>(d: D) -> Result { + use serde_json::value::RawValue; + #[derive(Default, Deserialize)] + #[serde(rename_all = "snake_case")] + enum Kind { + Metadata, + RunModel, + Run, + Task, + SessionEnd, + #[default] + #[serde(other)] + Ignored, + } + #[derive(Deserialize)] + struct Payload { + #[serde(default)] + kind: Option, + run_id: Option, + event: Option>, + record: Option>, + } + fn decode( + value: Option>, + ) -> Result, E> { + value + .map(|raw| serde_json::from_str(raw.get()).map_err(E::custom)) + .transpose() + } + let raw = Payload::deserialize(d)?; + Ok(match raw.kind.unwrap_or_default() { + Kind::Ignored => Self::Ignored, + Kind::SessionEnd => Self::SessionEnd, + Kind::Metadata => Self::Metadata { + record: decode(raw.record)?, + }, + Kind::RunModel => Self::RunModel { + record: decode(raw.record)?, + }, + Kind::Task => Self::Task { + event: decode(raw.event)?, + }, + Kind::Run => match decode(raw.event)? { + None | Some(RunEvent::Ignored) => Self::Ignored, + Some(event) => Self::Run { + run_id: raw + .run_id + .ok_or_else(|| serde::de::Error::missing_field("run_id"))?, + event, + }, + }, + }) + } +} +#[derive(Clone, Default, Deserialize)] +pub(crate) struct ExportMetadata { + pub model_id: Option, + pub provider_id: Option, + pub workspace_root: Option, +} +#[derive(Clone, Deserialize)] +pub(crate) struct RunModelRecord { + pub run_stream: RunStream, + #[serde(flatten)] + pub metadata: ExportMetadata, +} +#[derive(Clone, Deserialize)] +pub(crate) struct RunStream { + pub id: String, +} +#[derive(Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub(crate) enum RunEvent { + Started { + prompt: Option, + }, + ModelInputTraceRecorded { + request_record_id: String, + model_step: Option, + bounded: Option, + }, + ModelResponseCreated { + response_id: Option, + }, + ModelCompleted { + usage: Option, + }, + AssistantToolCallsCommitted { + tool_calls: Vec, + response_id: Option, + }, + ToolResultBatchCommitted { + results: Vec, + }, + ModelFailureTerminalIntent { + reason: Option, + }, + // Resource samples are not provider token usage. Their unused body is skipped. + ResourceUsageSampled, + AssistantMessageCommitted { + text: String, + response_id: Option, + }, + Terminal { + terminal: String, + reason: Option, + }, + #[serde(other)] + Ignored, +} +#[derive(Deserialize)] +pub(crate) struct ExportToolCall { + pub call_id: String, + pub name: String, + pub args: Option, +} +#[derive(Deserialize)] +pub(crate) struct ExportToolResult { + pub tool_call_id: String, + pub text: Option, +} +#[derive(Deserialize)] +pub(crate) struct TaskEvent { + #[serde(default)] + pub kind: TaskEventKind, + pub task_id: Option, + pub idempotency_key: Option, + pub origin: Option, + pub reason: Option, +} +#[derive(Default, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum TaskEventKind { + Scheduled, + Started, + Completed, + Failed, + #[serde(alias = "canceled")] + Cancelled, + Rejected, + #[default] + #[serde(other)] + Ignored, +} +#[derive(Deserialize)] +pub(crate) struct TaskOrigin { + pub tool_call_id: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn non_event_payloads_and_future_records_are_ignored_but_known_records_are_checked() { + for payload in [ + json!({"record":{"future":true}}), + json!({"kind":null,"record":{"future":true}}), + json!({"kind":"future_payload","event":{"future":true}}), + json!({"kind":"run","event":{"kind":"future_event"}}), + json!({"kind":"run"}), + ] { + assert!(matches!( + serde_json::from_str::(&payload.to_string()).unwrap(), + NativePayload::Ignored + )); + } + assert!(serde_json::from_str::(r#"{"kind":"run","run_id":"r","event":{"kind":"model_completed","usage":{"input_tokens":"wrong"}}}"#).is_err()); + assert!(serde_json::from_str::( + r#"{"kind":"run","event":{"kind":"started"}}"# + ) + .is_err()); + } +} diff --git a/bt-daemon/src/muse/field.rs b/bt-daemon/src/muse/field.rs new file mode 100644 index 00000000..229f061d --- /dev/null +++ b/bt-daemon/src/muse/field.rs @@ -0,0 +1,47 @@ +//! Optional wire fields whose explicit null differs from absence in authored +//! metadata and replay fingerprints. Ordinary optional data uses Option instead. +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) enum Field { + #[default] + Missing, + Null, + Value(T), +} +impl Field { + pub fn is_missing(&self) -> bool { + matches!(self, Self::Missing) + } + pub fn as_ref(&self) -> Option<&T> { + match self { + Self::Value(v) => Some(v), + _ => None, + } + } +} +impl Field { + pub fn as_deref(&self) -> Option<&T::Target> { + self.as_ref().map(std::ops::Deref::deref) + } +} +impl From> for Field { + fn from(value: Option) -> Self { + match value { + Some(v) => Self::Value(v), + None => Self::Null, + } + } +} +impl Serialize for Field { + fn serialize(&self, s: S) -> Result { + match self { + Self::Value(v) => v.serialize(s), + Self::Null | Self::Missing => s.serialize_none(), + } + } +} +impl<'de, T: Deserialize<'de>> Deserialize<'de> for Field { + fn deserialize>(d: D) -> Result { + Ok(Option::::deserialize(d)?.into()) + } +} diff --git a/bt-daemon/src/muse/hooks.rs b/bt-daemon/src/muse/hooks.rs new file mode 100644 index 00000000..99f79522 --- /dev/null +++ b/bt-daemon/src/muse/hooks.rs @@ -0,0 +1,302 @@ +//! Hook wire DTOs shared by live decoding and export synthesis. Unknown events +//! are ignored; known events decode once before any tracing policy executes. +use super::field::Field; +use super::ids::{RequestId, ToolCallId, TurnId}; +pub(crate) use super::messages::*; +pub(crate) use super::outcome::*; +use super::usage::{TokenUsage, UsageScope}; +use crate::wire::Envelope; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Value; + +macro_rules! hook_events { + ($($name:ident($body:ty)),+ $(,)?) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] + pub(crate) enum HookName { $($name),+ } + impl HookName { + pub const ALL: &'static [Self] = &[$(Self::$name),+]; + pub fn as_str(self) -> &'static str { match self { $(Self::$name => stringify!($name)),+ } } + pub fn parse(name: &str) -> Option { match name { + $(stringify!($name) => Some(Self::$name),)+ + "ToolUseStart" => Some(Self::PreToolUse), + _ => None, + }} + } + #[derive(Debug, Clone)] + pub(crate) enum MuseEvent { $($name(Hook<$body>)),+ } + impl MuseEvent { + pub fn name(&self) -> HookName { match self { $(Self::$name(_) => HookName::$name),+ } } + pub fn context(&self) -> &HookContext { match self { $(Self::$name(hook) => &hook.context),+ } } + pub fn payload(&self) -> serde_json::Result { match self { $(Self::$name(hook) => serde_json::to_value(hook)),+ } } + pub fn decode(envelope: &Envelope) -> serde_json::Result> { + Ok(match HookName::parse(&envelope.event) { + $(Some(HookName::$name) => Some(Self::$name(serde_json::from_value(envelope.payload.clone())?)),)+ + None => None, + }) + } + } + } +} +hook_events! { + SessionStart(Empty), UserPromptSubmit(PromptSubmitted), + PreToolUse(ToolStarted), PermissionRequest(PermissionObserved), + PostToolUse(ToolFinished), PostToolUseFailure(ToolFinished), + PostToolBatch(ToolBatchFinished), PreLLMCall(ModelStarted), PostLLMCall(ModelFinished), + PreCompact(Compaction), PostCompact(Compaction), + SubagentStart(SubagentStarted), SubagentStop(SubagentFinished), + Stop(TurnFinished), StopFailure(TurnFinished), SessionEnd(SessionFinished), + Notification(NotificationObserved), +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct Empty {} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct Hook { + #[serde(flatten)] + pub context: HookContext, + #[serde(flatten)] + pub body: B, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct HookContext { + #[serde(skip_serializing_if = "Option::is_none")] + pub session_id: Option, + #[serde( + default, + deserialize_with = "optional_id", + skip_serializing_if = "Option::is_none" + )] + pub turn_id: Option, + #[serde( + default, + deserialize_with = "optional_id", + skip_serializing_if = "Option::is_none" + )] + pub request_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub source: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub transcript_path: Option, + #[serde(flatten)] + pub session: SessionMetadata, +} +impl HookContext { + pub fn turn(&self) -> Option<&str> { + self.turn_id.as_deref().filter(|s| !s.is_empty()) + } + pub fn request(&self) -> Option<&str> { + self.request_id.as_deref().filter(|s| !s.is_empty()) + } +} +// The native optional identifiers have historically tolerated non-string values. +// Keep this compatibility at decoding, never inside correlation logic. +fn optional_id<'de, D: Deserializer<'de>, T: From>(d: D) -> Result, D::Error> { + #[derive(Deserialize)] + #[serde(untagged)] + enum Id { + Text(String), + Other(serde::de::IgnoredAny), + } + Ok(match Id::deserialize(d)? { + Id::Text(s) if !s.is_empty() => Some(T::from(s)), + _ => None, + }) +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct SessionMetadata { + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub cwd: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub model: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub provider: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub permission_mode: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub trajectory_id: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub root_session_id: Field, +} +impl SessionMetadata { + pub fn update(&mut self, incoming: &Self) { + macro_rules! update { ($($field:ident),+) => { $(if !incoming.$field.is_missing() { self.$field.clone_from(&incoming.$field); })+ }; } + update!( + cwd, + model, + provider, + permission_mode, + trajectory_id, + root_session_id + ); + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct PromptSubmitted { + pub prompt: Option, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ModelDetails { + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub step: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub attempt: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub message_count: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub tool_count: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub options: Field, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ModelStarted { + #[serde(flatten)] + pub details: ModelDetails, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub request_record_id: Field, + pub messages: Option>, + pub tools: Option>, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub input_summary: Field, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ModelFinished { + pub usage: Option, + #[serde(default)] + pub usage_scope: UsageScope, + pub output_text_preview: Option, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub response_id: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub finish_reason: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub tool_call_count: Field, + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub options: Field, + #[serde(flatten)] + pub outcome: NativeOutcome, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ToolStarted { + #[serde(default, deserialize_with = "optional_id")] + pub tool_use_id: Option, + pub tool_name: Option, + pub tool_input: Option, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ToolFinished { + #[serde(flatten)] + pub tool: ToolStarted, + pub tool_response: Option, + pub duration_ms: Option, + #[serde(flatten)] + pub outcome: NativeOutcome, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct ToolBatchFinished { + #[serde(default)] + pub tool_calls: Vec, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct PermissionObserved { + pub tool_name: Option, + pub tool_input: Option, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct NotificationObserved { + pub notification_type: Option, + pub title: Option, + pub message: Option, + #[serde(flatten)] + pub tool: PermissionObserved, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct SubagentStarted { + pub child_session_id: Option, + pub subagent_id: Option, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct SubagentFinished { + #[serde(flatten)] + pub child: SubagentStarted, + pub last_assistant_message: Option, + #[serde(flatten)] + pub outcome: NativeOutcome, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct Compaction { + pub trigger: Option, + #[serde(flatten)] + pub outcome: NativeOutcome, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct TurnFinished { + pub last_assistant_message: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub terminal: Option, + #[serde(flatten)] + pub outcome: NativeOutcome, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct SessionFinished { + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub reason: Field, + #[serde(flatten)] + pub outcome: NativeOutcome, +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + fn envelope(name: &str, payload: Value) -> Envelope { + Envelope { + source: "muse".into(), + session_id: "session".into(), + event: name.into(), + ts_ms: 100, + source_version: None, + plugin_version: None, + managed_run_id: None, + capture: None, + route: None, + config: None, + payload, + } + } + #[test] + fn all_hook_names_decode_and_optional_ids_remain_tolerant() { + for name in HookName::ALL { + let e = envelope( + name.as_str(), + json!({"turn_id":{},"request_id":4,"future_field":true}), + ); + let decoded = MuseEvent::decode(&e).unwrap().unwrap(); + assert_eq!(decoded.name(), *name); + assert_eq!(decoded.context().turn(), None); + assert_eq!(decoded.context().request(), None); + } + assert!(matches!( + MuseEvent::decode(&envelope("ToolUseStart", json!({}))).unwrap(), + Some(MuseEvent::PreToolUse(_)) + )); + assert!(MuseEvent::decode(&envelope("FutureEvent", json!({}))) + .unwrap() + .is_none()); + assert!(MuseEvent::decode(&envelope("PreLLMCall", json!({"step":"wrong"}))).is_err()); + } + #[test] + fn known_content_is_validated_and_unknown_content_is_lossless() { + let block = json!({"type":"image","source":{"url":"example"},"future":true}); + let decoded: ContentBlock = serde_json::from_value(block.clone()).unwrap(); + assert_eq!(serde_json::to_value(decoded).unwrap(), block); + assert!(serde_json::from_value::(json!({"type":"text","text":42})).is_err()); + let args = ToolArguments::Encoded("{\"x\":1}".into()); + assert_eq!(args.input().0, json!({"x":1})); + assert_eq!( + ToolArguments::Encoded("invalid json".into()).input().0, + json!("invalid json") + ); + } +} diff --git a/bt-daemon/src/muse/ids.rs b/bt-daemon/src/muse/ids.rs new file mode 100644 index 00000000..d7033557 --- /dev/null +++ b/bt-daemon/src/muse/ids.rs @@ -0,0 +1,15 @@ +//! Native correlation identifiers are separate domains even though their wire +//! representation is a string. Composite span keys are constructed by reducers. +use serde::{Deserialize, Serialize}; +macro_rules! identifier { + ($($name:ident),+ $(,)?) => {$ ( + #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] + #[serde(transparent)] + pub(crate) struct $name(String); + impl From for $name { fn from(value: String) -> Self { Self(value) } } + impl From<&str> for $name { fn from(value: &str) -> Self { Self(value.into()) } } + impl std::ops::Deref for $name { type Target = str; fn deref(&self) -> &str { &self.0 } } + impl std::fmt::Display for $name { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(&self.0) } } + )+}; +} +identifier!(TurnId, RequestId, ToolCallId); diff --git a/bt-daemon/src/muse/messages.rs b/bt-daemon/src/muse/messages.rs new file mode 100644 index 00000000..63f5549c --- /dev/null +++ b/bt-daemon/src/muse/messages.rs @@ -0,0 +1,196 @@ +//! Native hook message and provider payload contracts. +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Value; +use std::collections::BTreeMap; + +/// These values have schemas owned by the tool or provider, not by Muse tracing. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct ToolInput(pub Value); +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct ToolOutput(pub Value); +pub(crate) type ProviderOptions = BTreeMap; +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum ModelOutput { + Text(String), + Messages(Vec), +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct NativeMessage { + pub role: MessageRole, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(flatten)] + pub extensions: BTreeMap, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(from = "String", into = "String")] +pub(crate) enum MessageRole { + System, + Developer, + User, + Assistant, + Tool, + Other(String), +} +impl From for MessageRole { + fn from(s: String) -> Self { + match s.as_str() { + "system" => Self::System, + "developer" => Self::Developer, + "user" => Self::User, + "assistant" => Self::Assistant, + "tool" => Self::Tool, + _ => Self::Other(s), + } + } +} +impl From for String { + fn from(role: MessageRole) -> Self { + match role { + MessageRole::System => "system", + MessageRole::Developer => "developer", + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::Tool => "tool", + MessageRole::Other(s) => return s, + } + .into() + } +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum MessageContent { + Text(String), + Blocks(Vec), +} +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type")] +pub(crate) enum ContentBlock { + #[serde(rename = "text")] + Text { + text: String, + #[serde(flatten)] + extensions: BTreeMap, + }, + #[serde(rename = "assistant_message")] + AssistantMessage { + text: String, + #[serde(flatten)] + extensions: BTreeMap, + }, + #[serde(untagged)] + Other { + #[serde(rename = "type")] + kind: String, + #[serde(flatten)] + fields: BTreeMap, + }, +} +impl<'de> Deserialize<'de> for ContentBlock { + fn deserialize>(d: D) -> Result { + // Dispatch first so a malformed known block cannot hide in the fallback. + #[derive(Deserialize)] + struct Tag { + #[serde(rename = "type")] + kind: String, + } + #[derive(Deserialize)] + struct Text { + text: String, + #[serde(rename = "type")] + _kind: String, + #[serde(flatten)] + extensions: BTreeMap, + } + let value = Value::deserialize(d)?; + let tag: Tag = serde_json::from_value(value.clone()).map_err(serde::de::Error::custom)?; + match tag.kind.as_str() { + "text" | "assistant_message" => { + let body: Text = serde_json::from_value(value).map_err(serde::de::Error::custom)?; + Ok(if tag.kind == "text" { + Self::Text { + text: body.text, + extensions: body.extensions, + } + } else { + Self::AssistantMessage { + text: body.text, + extensions: body.extensions, + } + }) + } + _ => { + let mut fields: BTreeMap = + serde_json::from_value(value).map_err(serde::de::Error::custom)?; + fields.remove("type"); + Ok(Self::Other { + kind: tag.kind, + fields, + }) + } + } + } +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct AssistantToolCall { + pub id: String, + #[serde(rename = "type")] + pub call_type: FunctionCallType, + pub function: FunctionCall, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum FunctionCallType { + Function, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct FunctionCall { + pub name: String, + pub arguments: ToolArguments, +} +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum ToolArguments { + Encoded(String), + Json(Value), +} +impl ToolArguments { + pub fn input(&self) -> ToolInput { + ToolInput(match self { + Self::Encoded(s) => { + serde_json::from_str(s).unwrap_or_else(|_| Value::String(s.clone())) + } + Self::Json(v) => v.clone(), + }) + } +} +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct ToolDescription { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub has_parameters: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub strict: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub parameters: Option, + #[serde(flatten)] + pub extensions: BTreeMap, +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct InputSummary { + #[serde(skip_serializing_if = "Option::is_none")] + pub request_message_count: Option, + #[serde(flatten)] + pub extensions: BTreeMap, +} diff --git a/bt-daemon/src/muse/metadata.rs b/bt-daemon/src/muse/metadata.rs new file mode 100644 index 00000000..7cca8923 --- /dev/null +++ b/bt-daemon/src/muse/metadata.rs @@ -0,0 +1,102 @@ +//! Authored Braintrust row shapes. Serialization into SpanRow is the final JSON +//! boundary; lifecycle and metric policy never navigate these JSON values. +use super::field::Field; +use super::hooks::*; +use serde::Serialize; +use serde_json::Value; + +#[derive(Serialize)] +pub(crate) struct RootMetadata<'a> { + pub source: &'static str, + pub session_id: &'a str, + pub username: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub muse_version: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + pub adapter_version: Option<&'a str>, + #[serde(flatten)] + pub session: &'a SessionMetadata, +} +#[derive(Serialize)] +pub(crate) struct ModelInput<'a> { + pub messages: &'a Option>, + pub tools: &'a Option>, +} +#[derive(Serialize)] +pub(crate) struct ModelStartMetadata<'a> { + #[serde(skip_serializing_if = "Field::is_missing")] + pub model: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub provider: &'a Field, + #[serde(flatten)] + pub details: &'a ModelDetails, + #[serde(skip_serializing_if = "Field::is_missing")] + pub request_record_id: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub input_summary: &'a Field, +} +#[derive(Serialize)] +pub(crate) struct ModelEndMetadata<'a> { + #[serde(skip_serializing_if = "Field::is_missing")] + pub finish_reason: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub response_id: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub tool_call_count: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub status: &'a Field, + #[serde(skip_serializing_if = "Field::is_missing")] + pub options: &'a Field, +} +#[derive(Serialize)] +pub(crate) struct ToolStartMetadata<'a> { + pub tool_name: &'a str, + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub native_start_missing: bool, +} +#[derive(Serialize)] +pub(crate) struct ToolEndMetadata { + pub tool_outcome: ToolOutcomeLabel, +} +#[derive(Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ToolOutcomeLabel { + Success, + Error, +} +#[derive(Serialize)] +pub(crate) struct SubagentMetadata<'a> { + pub subagent_id: &'a Option, + pub child_session_id: &'a Option, +} +#[derive(Serialize)] +pub(crate) struct CompactMetadata<'a> { + pub trigger: &'a Option, + pub native_operation_id_available: bool, +} +#[derive(Serialize)] +pub(crate) struct SessionEndMetadata<'a> { + pub end_reason: &'a Field, +} +#[derive(Serialize)] +pub(crate) struct UnpairedMetadata { + pub unpaired: bool, +} +#[derive(Serialize)] +pub(crate) struct PermissionMetadata<'a> { + pub tool_name: &'a Option, +} +#[derive(Serialize)] +pub(crate) struct NotificationMetadata<'a> { + pub tool_name: &'a Option, + pub notification_type: &'a Option, + pub title: &'a Option, + pub message: &'a Option, +} +#[derive(Serialize)] +pub(crate) struct CompletionFingerprint<'a> { + pub output: &'a Option, + pub metrics: &'a Option, + pub error: &'a Option, + pub metadata: &'a Option, +} diff --git a/bt-daemon/src/muse/mod.rs b/bt-daemon/src/muse/mod.rs new file mode 100644 index 00000000..c8a1ab98 --- /dev/null +++ b/bt-daemon/src/muse/mod.rs @@ -0,0 +1,12 @@ +//! Native Muse contracts. MSP, hooks, exports, and settings have separate wire +//! schemas; only hooks and exports normalize into the same tracing events. +pub(crate) mod config; +pub(crate) mod export_v1; +pub(crate) mod field; +pub(crate) mod hooks; +pub(crate) mod ids; +pub(crate) mod messages; +pub(crate) mod metadata; +pub(crate) mod msp; +pub(crate) mod outcome; +pub(crate) mod usage; diff --git a/bt-daemon/src/muse/msp.rs b/bt-daemon/src/muse/msp.rs new file mode 100644 index 00000000..e3345518 --- /dev/null +++ b/bt-daemon/src/muse/msp.rs @@ -0,0 +1,357 @@ +//! Read-only MSP client. The SDK owns framing, correlation, and initialization; +//! we own process lifetime, timeouts, pagination, and the repo's spawn policy. +use anyhow::{bail, Context, Result}; +use muse_code_sdk::connection::{Connection, StdioTransport}; +use muse_code_sdk::handshake::Handshake; +use muse_code_sdk::msp::{ClientInfo, InitializeParams, SessionListParams, SessionListResult}; +use serde::{de::DeserializeOwned, Serialize}; +use std::collections::HashSet; +use std::ffi::OsStr; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::task::JoinHandle; + +const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const MAX_PAGES: usize = 1_000; +const DIAGNOSTIC_BYTES: usize = 64 * 1024; + +trait Method { + const NAME: &'static str; + type Params: Serialize; + type Response: DeserializeOwned; +} +struct ListSessions; +impl Method for ListSessions { + const NAME: &'static str = "session/list"; + type Params = SessionListParams; + type Response = SessionListResult; +} + +pub(crate) struct MuseMspClient { + connection: Connection, + host: tokio::process::Child, + driver: JoinHandle<()>, + diagnostics: Option>>, +} + +impl Drop for MuseMspClient { + fn drop(&mut self) { + // Cancellation must not leave a detached driver holding the child's pipes. + self.driver.abort(); + if let Some(diagnostics) = &self.diagnostics { + diagnostics.abort(); + } + // host was spawned with kill_on_drop. + } +} + +impl MuseMspClient { + pub(crate) async fn connect(executable: &OsStr) -> Result { + let mut command: tokio::process::Command = + crate::subprocess::background_command(executable).into(); + let mut host = command + .arg("serve") + .kill_on_drop(true) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .context("start Muse MSP host")?; + let stdin = host.stdin.take().context("Muse MSP stdin was not piped")?; + let stdout = host + .stdout + .take() + .context("Muse MSP stdout was not piped")?; + let stderr = host + .stderr + .take() + .context("Muse MSP stderr was not piped")?; + let (connection, driver) = Connection::new(StdioTransport::new(stdin, stdout)); + let client = Self { + connection, + host, + driver: tokio::spawn(driver.run()), + diagnostics: Some(tokio::spawn(drain_stderr(stderr))), + }; + let initialized = tokio::time::timeout( + REQUEST_TIMEOUT, + Handshake::new(client.connection.clone()).initialize(InitializeParams { + client_info: ClientInfo { + name: "braintrust_muse_import".into(), + version: env!("CARGO_PKG_VERSION").into(), + title: None, + }, + capabilities: None, + }), + ) + .await + .context("Muse MSP initialize timed out")? + .context("initialize Muse MSP connection")?; + if let Some(warning) = initialized.fingerprint_warning { + tracing::warn!(?warning, "Muse MSP schema fingerprint differs from SDK"); + } + Ok(client) + } + + async fn call(&self, params: &M::Params) -> Result { + call::(&self.connection, params, REQUEST_TIMEOUT).await + } + + pub(crate) async fn list_sessions( + &self, + params: &SessionListParams, + ) -> Result { + self.call::(params).await + } + + pub(crate) async fn session_ids(&self) -> Result> { + let mut sessions = Vec::new(); + let mut seen_cursors = HashSet::new(); + let mut cursor = None; + for page_index in 1..=MAX_PAGES { + let page = self + .list_sessions(&SessionListParams { + cursor, + limit: Some(200), + workspace_root: None, + updated_after: None, + filter: None, + }) + .await + .with_context(|| format!("Muse session/list page {page_index}"))?; + cursor = append_page(page, page_index, &mut sessions, &mut seen_cursors)?; + if cursor.is_none() { + sessions.sort(); + sessions.dedup(); + return Ok(sessions); + } + } + bail!("Muse session/list exceeded {MAX_PAGES} pages") + } + + pub(crate) async fn shutdown(mut self) -> Result<()> { + tokio::time::timeout(SHUTDOWN_TIMEOUT, self.connection.close()) + .await + .context("close Muse MSP connection timed out")?; + tokio::time::timeout(SHUTDOWN_TIMEOUT, &mut self.driver) + .await + .context("Muse MSP driver did not stop")??; + let status = match tokio::time::timeout(SHUTDOWN_TIMEOUT, self.host.wait()).await { + Ok(status) => status.context("wait for Muse MSP host")?, + Err(_) => { + self.host.start_kill().context("stop Muse MSP host")?; + self.host.wait().await.context("reap Muse MSP host")?; + bail!("Muse MSP host did not exit after stdin closed") + } + }; + let diagnostics = self + .diagnostics + .as_mut() + .context("Muse MSP diagnostics missing")?; + let stderr = tokio::time::timeout(SHUTDOWN_TIMEOUT, diagnostics) + .await + .context("collect Muse MSP diagnostics timed out")???; + if !status.success() { + bail!("Muse MSP session listing failed: {}", stderr.trim()); + } + Ok(()) + } +} + +/// The SDK's low-level transport is object-valued. Confine that representation +/// here: neither callers nor pagination can inspect untyped protocol values. +async fn call( + connection: &Connection, + params: &M::Params, + timeout: Duration, +) -> Result { + let serde_json::Value::Object(params) = serde_json::to_value(params)? else { + bail!("MSP {} parameters must serialize as an object", M::NAME) + }; + let response = tokio::time::timeout(timeout, connection.request(M::NAME, Some(params))) + .await + .with_context(|| format!("Muse MSP {} timed out", M::NAME))? + .with_context(|| format!("Muse MSP {} failed", M::NAME))?; + serde_json::from_value(serde_json::Value::Object(response)) + .with_context(|| format!("decode Muse MSP {} response", M::NAME)) +} + +fn append_page( + page: SessionListResult, + index: usize, + sessions: &mut Vec, + seen_cursors: &mut HashSet, +) -> Result> { + for session in page.sessions { + crate::transcript_import::validate_session_id(&session.session_id) + .with_context(|| format!("invalid Muse session id on session/list page {index}"))?; + sessions.push(session.session_id); + } + if let Some(cursor) = &page.next_cursor { + if !seen_cursors.insert(cursor.clone()) { + bail!("Muse session/list repeated cursor on page {index}") + } + } + Ok(page.next_cursor) +} + +/// Drain all bytes to prevent pipe backpressure, retaining bounded diagnostics. +/// Also used by the independent `muse export` subprocess. +pub(crate) async fn drain_stderr(mut stderr: tokio::process::ChildStderr) -> 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 = DIAGNOSTIC_BYTES.saturating_sub(output.len()); + output.extend_from_slice(&buffer[..count.min(remaining)]); + } + Ok(String::from_utf8_lossy(&output).into_owned()) +} + +#[cfg(test)] +mod tests { + use super::*; + use muse_code_sdk::DuplexTransport; + use serde_json::{json, Value}; + use std::future::Future; + use std::pin::Pin; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, ReadHalf, WriteHalf}; + + struct TestTransport { + reader: BufReader>, + writer: WriteHalf, + } + impl DuplexTransport for TestTransport { + fn read_chunk( + &mut self, + ) -> Pin>> + Send + '_>> { + Box::pin(async move { + let bytes = self.reader.fill_buf().await?; + if bytes.is_empty() { + return Ok(None); + } + let chunk = String::from_utf8(bytes.to_vec()).map_err(std::io::Error::other)?; + let len = bytes.len(); + self.reader.consume(len); + Ok(Some(chunk)) + }) + } + fn write<'a>( + &'a mut self, + frame: &str, + ) -> Pin> + Send + 'a>> { + let bytes = frame.as_bytes().to_vec(); + Box::pin(async move { self.writer.write_all(&bytes).await }) + } + } + fn params() -> SessionListParams { + SessionListParams { + cursor: None, + limit: Some(200), + workspace_root: None, + updated_after: None, + filter: None, + } + } + fn session(id: &str) -> Value { + json!({"sessionId":id,"path":"/sessions/example","status":"notLoaded","turnCount":0,"activeTurnId":null,"createdAt":"2026-10-23T00:00:00Z","updatedAt":"2026-10-23T00:00:00Z"}) + } + async fn exchange(result: Value, is_error: bool) -> Result { + let (client, server) = tokio::io::duplex(8192); + let (read, write) = tokio::io::split(client); + let (connection, driver) = Connection::new(TestTransport { + reader: BufReader::new(read), + writer: write, + }); + let driver = tokio::spawn(driver.run()); + let server = tokio::spawn(async move { + let (read, mut write) = tokio::io::split(server); + let request = BufReader::new(read) + .lines() + .next_line() + .await + .unwrap() + .unwrap(); + let request: Value = serde_json::from_str(&request).unwrap(); + assert_eq!(request["method"], "session/list"); + assert_eq!(request["params"], json!({"limit":200})); + let response = if is_error { + json!({"jsonrpc":"2.0","id":request["id"],"error":result}) + } else { + json!({"jsonrpc":"2.0","id":request["id"],"result":result}) + }; + write + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"progress\",\"params\":{}}\n") + .await + .unwrap(); + write + .write_all(format!("{response}\n").as_bytes()) + .await + .unwrap(); + }); + let result = call::(&connection, ¶ms(), Duration::from_secs(1)).await; + connection.close().await; + driver.await.unwrap(); + server.await.unwrap(); + result + } + #[tokio::test] + async fn typed_response_and_errors_through_sdk_connection() { + let page = exchange(json!({"sessions":[session("a")],"nextCursor":null}), false) + .await + .unwrap(); + assert_eq!(page.sessions[0].session_id, "a"); + assert!(exchange(json!({"sessions":[{"sessionId":"a"}]}), false) + .await + .unwrap_err() + .to_string() + .contains("decode Muse MSP session/list")); + assert!( + exchange(json!({"code":-32601,"message":"unsupported"}), true) + .await + .unwrap_err() + .to_string() + .contains("session/list failed") + ); + } + #[test] + fn pagination_validates_ids_and_rejects_cycles() { + let page = |id: &str| { + serde_json::from_value::( + json!({"sessions":[session(id)],"nextCursor":"next"}), + ) + .unwrap() + }; + let mut ids = Vec::new(); + let mut cursors = HashSet::new(); + assert_eq!( + append_page(page("session-a"), 1, &mut ids, &mut cursors).unwrap(), + Some("next".into()) + ); + assert_eq!(ids, ["session-a"]); + assert!(append_page(page("session-b"), 2, &mut ids, &mut cursors) + .unwrap_err() + .to_string() + .contains("repeated cursor")); + assert!(append_page(page("../escape"), 3, &mut ids, &mut HashSet::new()).is_err()); + } + #[tokio::test] + async fn timeout_is_bounded_when_host_does_not_answer() { + let (client, _server) = tokio::io::duplex(8192); + let (read, write) = tokio::io::split(client); + let (connection, driver) = Connection::new(TestTransport { + reader: BufReader::new(read), + writer: write, + }); + let driver = tokio::spawn(driver.run()); + let result = call::(&connection, ¶ms(), Duration::from_millis(10)).await; + assert!(result.unwrap_err().to_string().contains("timed out")); + connection.close().await; + driver.await.unwrap(); + } +} diff --git a/bt-daemon/src/muse/outcome.rs b/bt-daemon/src/muse/outcome.rs new file mode 100644 index 00000000..71968eda --- /dev/null +++ b/bt-daemon/src/muse/outcome.rs @@ -0,0 +1,134 @@ +//! Native operation outcomes, including open future statuses. +use super::field::Field; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum ErrorDetail { + Text(String), + Structured(Value), +} +impl ErrorDetail { + pub fn render(&self) -> Option { + match self { + Self::Text(s) => (!s.is_empty()).then(|| s.clone()), + Self::Structured(Value::Null) => None, + Self::Structured(value) => Some(value.to_string()), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(from = "String", into = "String")] +pub(crate) enum OperationStatus { + Success, + Completed, + ToolCalls, + Stop, + Failed, + Failure, + Error, + Cancelled, + Canceled, + Interrupted, + Timeout, + TimedOut, + Denied, + Rejected, + Other(String), +} +impl From for OperationStatus { + fn from(s: String) -> Self { + match s.as_str() { + "success" => Self::Success, + "completed" => Self::Completed, + "tool_calls" => Self::ToolCalls, + "stop" => Self::Stop, + "failed" => Self::Failed, + "failure" => Self::Failure, + "error" => Self::Error, + "cancelled" => Self::Cancelled, + "canceled" => Self::Canceled, + "interrupted" => Self::Interrupted, + "timeout" => Self::Timeout, + "timed_out" => Self::TimedOut, + "denied" => Self::Denied, + "rejected" => Self::Rejected, + _ => Self::Other(s), + } + } +} +impl From for String { + fn from(s: OperationStatus) -> Self { + match s { + OperationStatus::Success => "success", + OperationStatus::Completed => "completed", + OperationStatus::ToolCalls => "tool_calls", + OperationStatus::Stop => "stop", + OperationStatus::Failed => "failed", + OperationStatus::Failure => "failure", + OperationStatus::Error => "error", + OperationStatus::Cancelled => "cancelled", + OperationStatus::Canceled => "canceled", + OperationStatus::Interrupted => "interrupted", + OperationStatus::Timeout => "timeout", + OperationStatus::TimedOut => "timed_out", + OperationStatus::Denied => "denied", + OperationStatus::Rejected => "rejected", + OperationStatus::Other(s) => return s, + } + .into() + } +} +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct NativeOutcome { + #[serde(default, skip_serializing_if = "Field::is_missing")] + pub status: Field, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error_details: Option, +} +impl NativeOutcome { + pub fn error_message(&self) -> Option { + let error = self + .error + .as_ref() + .and_then(ErrorDetail::render) + .or_else(|| { + self.status + .as_ref() + .filter(|s| { + matches!( + s, + OperationStatus::Failed + | OperationStatus::Failure + | OperationStatus::Error + | OperationStatus::Canceled + | OperationStatus::TimedOut + | OperationStatus::Cancelled + | OperationStatus::Interrupted + | OperationStatus::Timeout + | OperationStatus::Denied + | OperationStatus::Rejected + ) + }) + .map(|s| format!("Muse operation {}", String::from(s.clone()))) + }); + match ( + error, + self.error_details.as_ref().and_then(ErrorDetail::render), + ) { + (Some(error), Some(details)) => Some(format!("{error}: {details}")), + (Some(error), None) => Some(error), + (None, details) => details, + } + } +} +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum ToolOutcome { + Success, + Failure(String), + Unknown, +} diff --git a/bt-daemon/src/muse/usage.rs b/bt-daemon/src/muse/usage.rs new file mode 100644 index 00000000..e2af3a67 --- /dev/null +++ b/bt-daemon/src/muse/usage.rs @@ -0,0 +1,99 @@ +//! Hook/export counters are snake_case, optional, and may be turn-cumulative. +//! They are not the SDK's camelCase, required-counter MSP view types. +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)] +pub(crate) struct TokenUsage { + #[serde(skip_serializing_if = "Option::is_none")] + pub input_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cached_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_write_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_tokens: Option, +} + +#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq)] +#[serde(rename_all = "snake_case")] +pub(crate) enum UsageScope { + #[default] + Turn, + Model, +} + +#[derive(Debug, Clone, Default, Serialize, PartialEq)] +pub(crate) struct TokenMetrics { + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_cached_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_cache_creation_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_reasoning_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tokens: Option, +} + +impl TokenUsage { + pub fn delta_since(&self, previous: &Self) -> Self { + fn delta(current: Option, previous: Option) -> Option { + current.map(|current| { + previous + .and_then(|p| current.checked_sub(p)) + .unwrap_or(current) + }) + } + Self { + input_tokens: delta(self.input_tokens, previous.input_tokens), + output_tokens: delta(self.output_tokens, previous.output_tokens), + cached_tokens: delta(self.cached_tokens, previous.cached_tokens), + cache_read_tokens: delta(self.cache_read_tokens, previous.cache_read_tokens), + cache_write_tokens: delta(self.cache_write_tokens, previous.cache_write_tokens), + reasoning_tokens: delta(self.reasoning_tokens, previous.reasoning_tokens), + } + } + pub fn metrics(&self) -> TokenMetrics { + TokenMetrics { + prompt_tokens: self.input_tokens, + completion_tokens: self.output_tokens, + // Preserve the native alias precedence used before this refactor. + prompt_cached_tokens: self.cached_tokens.or(self.cache_read_tokens), + prompt_cache_creation_tokens: self.cache_write_tokens, + completion_reasoning_tokens: self.reasoning_tokens, + tokens: self + .input_tokens + .zip(self.output_tokens) + .map(|(i, o)| i.saturating_add(o)), + } + } +} + +impl TokenMetrics { + pub(crate) fn add(&mut self, other: &Self) { + fn add(to: &mut Option, from: Option) { + if let Some(from) = from { + *to = Some(to.unwrap_or(0).saturating_add(from)); + } + } + add(&mut self.prompt_tokens, other.prompt_tokens); + add(&mut self.completion_tokens, other.completion_tokens); + add(&mut self.prompt_cached_tokens, other.prompt_cached_tokens); + add( + &mut self.prompt_cache_creation_tokens, + other.prompt_cache_creation_tokens, + ); + add( + &mut self.completion_reasoning_tokens, + other.completion_reasoning_tokens, + ); + add(&mut self.tokens, other.tokens); + } +} 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/pipeline/journal.rs b/bt-daemon/src/pipeline/journal.rs index 9d8067f0..0bbff985 100644 --- a/bt-daemon/src/pipeline/journal.rs +++ b/bt-daemon/src/pipeline/journal.rs @@ -51,6 +51,8 @@ struct StoredOp { version: u32, op: SpanOp, late_merge_key: Option, + #[serde(default)] + late_merge_version: Option, turn_span_id: Option, is_turn: bool, } @@ -63,6 +65,7 @@ impl StoredOp { version: u32, op: &'a SpanOp, late_merge_key: &'a Option, + late_merge_version: &'a Option, turn_span_id: &'a Option, is_turn: bool, } @@ -70,6 +73,7 @@ impl StoredOp { version: 1, op, late_merge_key: &row.late_merge_key, + late_merge_version: &row.late_merge_version, turn_span_id: &row.turn_span_id, is_turn: row.is_turn, })?) @@ -83,6 +87,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) @@ -1414,6 +1419,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/common.rs b/bt-daemon/src/setup/common.rs index b50f0461..f9471094 100644 --- a/bt-daemon/src/setup/common.rs +++ b/bt-daemon/src/setup/common.rs @@ -291,10 +291,19 @@ 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, +) -> anyhow::Result<()> { + write_json_atomic_unlocked(path, &object, access) +} + +/// Serialize a concrete configuration under a caller-held native settings lock. +pub(super) fn write_json_atomic_unlocked( + path: &Path, + value: &T, + access: FileAccess, ) -> anyhow::Result<()> { let parent = path .parent() @@ -305,7 +314,7 @@ fn write_object_atomic_unlocked( parent.display() ) })?; - let mut encoded = serde_json::to_string_pretty(&Value::Object(object))?; + let mut encoded = serde_json::to_string_pretty(value)?; encoded.push('\n'); let mut temporary = tempfile::NamedTempFile::new_in(parent) .with_context(|| format!("failed to create temporary file in {}", parent.display()))?; 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..49fac6ad --- /dev/null +++ b/bt-daemon/src/setup/muse.rs @@ -0,0 +1,409 @@ +//! Muse Code setup orchestration. Native settings and hook schemas live in +//! crate::muse::config; this module owns locks, migration, and installation. +use super::common::{write_json_atomic_unlocked, CommandRunner, FileAccess}; +use super::Setup; +use crate::agents::Muse; +use crate::muse::config::{ + inspect_hook_file, HookFileInspection, HookOverlay, MuseSettings, + HOOK_COMMAND as MUSE_HOOK_COMMAND, +}; +use crate::muse::field::Field; +use crate::paths; +use anyhow::{bail, Context}; +use fs2::FileExt; +use std::path::{Path, PathBuf}; + +fn muse_hook_path(config_dir: &Path) -> PathBuf { + config_dir.join("braintrust-hooks.json") +} + +pub(crate) fn muse_managed_hook_config( + config_dir: &Path, + command: &str, +) -> anyhow::Result { + let settings = MuseSettings::load(&config_dir.join("settings.json"))?; + check_foreign_managed_hooks(config_dir, &settings)?; + let mut hooks = settings.hook_overlay()?; + hooks.remove_owned(); + if hooks.contains_braintrust_hook() { + bail!("Muse user hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); + } + hooks.append_owned(command); + Ok(hooks) +} + +pub(crate) fn muse_has_owned_managed_hooks(config_dir: &Path) -> anyhow::Result { + let settings = MuseSettings::load(&config_dir.join("settings.json"))?; + Ok(owned_managed_path(config_dir, &settings)?.is_some()) +} +fn owned_managed_path( + config_dir: &Path, + settings: &MuseSettings, +) -> anyhow::Result> { + match settings.managed_path(config_dir) { + Some(path) + if path == muse_hook_path(config_dir) && inspect_hook_file(&path)?.is_owned() => + { + Ok(Some(path)) + } + _ => Ok(None), + } +} +fn check_foreign_managed_hooks(config_dir: &Path, settings: &MuseSettings) -> anyhow::Result<()> { + if let Some(path) = settings.managed_path(config_dir) { + let inspection = inspect_hook_file(&path)?; + if inspection.contains_braintrust() + && !(path == muse_hook_path(config_dir) && inspection.is_owned()) + { + bail!("Muse managed hooks already include an unrecognized Braintrust forwarder; cannot safely replace it"); + } + } + Ok(()) +} +pub(crate) fn check_muse_project_hooks(start: &Path) -> anyhow::Result<()> { + for dir in start.ancestors() { + let path = dir.join(".muse/hooks.json"); + let inspection = inspect_hook_file(&path)?; + if inspection == HookFileInspection::Missing { + continue; + } + if inspection.contains_braintrust() { + bail!("Muse project hooks {} already include Braintrust tracing; cannot safely run a managed trace", path.display()); + } + break; + } + Ok(()) +} + +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() + ) + }) +} + +/// 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("BRAINTRUST_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. +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 setup_muse_at(config_dir: &Path) -> anyhow::Result<()> { + with_muse_settings_lock(config_dir, || { + let settings_path = config_dir.join("settings.json"); + let mut settings = MuseSettings::load(&settings_path)?; + check_foreign_managed_hooks(config_dir, &settings)?; + let mut hooks = settings.hook_overlay()?; + hooks.remove_owned(); + if hooks.contains_braintrust_hook() { + bail!("Muse user hooks include an unrecognized Braintrust forwarder; refusing to duplicate it"); + } + hooks.append_owned(MUSE_HOOK_COMMAND); + settings.hooks = Field::Value(hooks); + settings.ensure_schema_version(); + let old_path = owned_managed_path(config_dir, &settings)?; + if old_path.is_some() { + settings.managed_hooks_path = Field::Missing; + } + write_muse_settings(&settings_path, &settings)?; + if let Some(path) = old_path { + std::fs::remove_file(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"); + if !settings_path.exists() { + return Ok(()); + } + let mut settings = MuseSettings::load(&settings_path)?; + if !settings.hooks.is_missing() { + let mut hooks = settings.hook_overlay()?; + hooks.remove_owned(); + settings.hooks = if hooks.is_empty() { + Field::Missing + } else { + Field::Value(hooks) + }; + } + let old_path = owned_managed_path(config_dir, &settings)?; + if old_path.is_some() { + settings.managed_hooks_path = Field::Missing; + } + write_muse_settings(&settings_path, &settings)?; + if let Some(path) = old_path { + std::fs::remove_file(path)?; + } + Ok(()) + }) +} +fn update_muse_at(config_dir: &Path) -> anyhow::Result<()> { + let settings = MuseSettings::load(&config_dir.join("settings.json"))?; + if !settings.hook_overlay()?.contains_braintrust_hook() + && owned_managed_path(config_dir, &settings)?.is_none() + { + 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 { + let mut warnings = Vec::new(); + match crate::subprocess::background_command( + crate::env::var_os("BRAINTRUST_MUSE_BIN").unwrap_or_else(|| "muse".into()), + ) + .arg("--version") + .output() + { + Ok(output) if output.status.success() => match muse_version(&output.stdout) { + 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()), + }, + _ => 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 MuseSettings::load(&config_dir.join("settings.json")) { + Ok(settings) => settings, + Err(error) => { + warnings.push(format!("Muse settings could not be read: {error}")); + return warnings; + } + }; + if let Err(error) = check_foreign_managed_hooks(config_dir, &settings) { + warnings.push(error.to_string()); + } + match settings.hook_overlay() { + Err(_) => warnings.push("Muse user hooks are malformed".into()), + Ok(current) => { + let mut upgraded = current.clone(); + upgraded.remove_owned(); + upgraded.append_owned(MUSE_HOOK_COMMAND); + if 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 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()) +} +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_settings(path: &Path, settings: &MuseSettings) -> anyhow::Result<()> { + write_json_atomic_unlocked(path, settings, 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::*; + use crate::muse::config::OwnedHookFile; + use crate::setup::common::load_object; + use serde_json::Value; + + #[test] + fn setup_update_and_disable_preserve_enterprise_and_user_hooks() { + let temp = tempfile::tempdir().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!( + Value::Object(load_object(&config.join("settings.json")).unwrap()), + settings + ); + assert_eq!( + std::fs::read_to_string(config.join("enterprise.json")).unwrap(), + enterprise.to_string() + ); + } + + #[test] + fn migrate_old_owned_managed_file_and_preserve_unknown_files() { + let temp = tempfile::tempdir().unwrap(); + let config = temp.path(); + let hook_path = muse_hook_path(config); + std::fs::write( + &hook_path, + serde_json::to_vec(&OwnedHookFile::previous()).unwrap(), + ) + .unwrap(); + std::fs::write( + config.join("settings.json"), + serde_json::to_vec(&serde_json::json!({"managed_hooks_path":hook_path})).unwrap(), + ) + .unwrap(); + 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 invocation_overlay_replaces_owned_hooks_and_preserves_strict_enterprise() { + let temp = tempfile::tempdir().unwrap(); + 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(); + let overlay = serde_json::to_value(overlay).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 reject_unrecognized_braintrust_forwarders_and_project_duplicates() { + let temp = tempfile::tempdir().unwrap(); + 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 disable_before_install_does_not_create_invalid_native_settings() { + let temp = tempfile::tempdir().unwrap(); + disable_muse_at(temp.path()).unwrap(); + assert!(!temp.path().join("settings.json").exists()); + } + + #[test] + fn malformed_known_fields_do_not_replace_saved_settings() { + for invalid in [ + serde_json::json!({"schema_version":"wrong"}), + serde_json::json!({"managed_hooks_path":42}), + serde_json::json!({"hooks":null}), + serde_json::json!({"hooks":{"Stop":[{"hooks":[{"type":"command","command":null}]}]}}), + ] { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("settings.json"); + let saved = serde_json::to_vec(&invalid).unwrap(); + std::fs::write(&path, &saved).unwrap(); + assert!(setup_muse_at(temp.path()).is_err()); + assert_eq!(std::fs::read(&path).unwrap(), saved); + assert!(disable_muse_at(temp.path()).is_err()); + assert_eq!(std::fs::read(&path).unwrap(), saved); + } + } + + #[test] + 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(); + assert!(update_muse_at(temp.path()) + .unwrap_err() + .to_string() + .contains("not installed")); + } +} diff --git a/bt-daemon/src/span_processor.rs b/bt-daemon/src/span_processor.rs index b67152c6..cc721f19 100644 --- a/bt-daemon/src/span_processor.rs +++ b/bt-daemon/src/span_processor.rs @@ -399,6 +399,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/trace_command.rs b/bt-daemon/src/trace_command.rs index e2d80014..4837b7d9 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -64,6 +64,8 @@ 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, #[value(name = "opencode", alias = "open-code")] @@ -78,6 +80,7 @@ impl DoctorAgent { 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(), @@ -128,6 +131,9 @@ pub struct UpdateArgs { pub enum SetupAgent { /// Install the published Codex tracing plugin. Codex, + /// Install the Muse Code tracing hooks. + #[command(alias = "muse-code")] + Muse, /// Install the published Claude Code tracing plugin. Claude, /// Configure the published OpenCode tracing plugin. @@ -147,6 +153,7 @@ 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(), diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index d3732c54..79e589ce 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, SetupAgent, 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; @@ -543,6 +543,11 @@ 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::muse_doctor_warnings( + &paths::muse_config_dir(), + )); + } let (route, route_source) = match settings.route { Some(route) => (Some(route), "settings_file".to_string()), @@ -697,6 +702,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 } @@ -755,7 +768,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 { @@ -1111,6 +1125,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 6d713459..36b010d3 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -11,6 +11,7 @@ mod antigravity; mod claude; mod codex; mod cursor; +pub(crate) mod muse; mod pipeline; pub use pipeline::{import_transcript, import_transcripts, run_import}; @@ -197,7 +198,7 @@ fn find_jsonl_files(directory: &Path, matches: &mut Vec) { } } -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() diff --git a/bt-daemon/src/transcript_import/muse.rs b/bt-daemon/src/transcript_import/muse.rs new file mode 100644 index 00000000..558085cd --- /dev/null +++ b/bt-daemon/src/transcript_import/muse.rs @@ -0,0 +1,1101 @@ +//! Streaming Muse export reader and correlation policy. Wire DTOs live in +//! muse::export_v1; synthesized hooks use the same typed events as live capture. +use super::{envelope, validate_session_id}; +use crate::muse::export_v1::*; +use crate::muse::hooks::*; +use crate::muse::usage::{TokenUsage, UsageScope}; +use crate::wire::Envelope; +use anyhow::{anyhow, bail, Context, Result}; +use serde::de::{DeserializeSeed, IgnoredAny, MapAccess, SeqAccess, Visitor}; +#[cfg(test)] +use serde_json::{json, Value}; +use std::collections::{BTreeMap, HashMap}; +use std::fmt; +use std::fs::File; +use std::io::BufReader; +use std::path::Path; + +#[derive(Clone, Copy)] +pub(crate) enum ReadMode { + /// Completed-file API: finalize observed unfinished work. + #[cfg(test)] + Historical, + /// One-shot observation: require traceable work, keep active spans open. + Import, + /// Live polling may observe a session before its first run. + Attach, +} + +#[derive(Default)] +struct RunState { + active_model: Option, + last_assistant: Option, + model_ordinal: u64, + step_attempts: HashMap, +} +struct ModelState { + request_id: String, + usage: Option, + output: Option, + response_id: Option, + completed_at: Option, +} +struct PendingTool { + turn_id: String, + tool: ToolStarted, + completed: bool, +} +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, + current_session_ended: bool, + activity_count: usize, + metadata: ExportMetadata, + tools: HashMap, + tool_errors: HashMap, + tool_tasks: HashMap, + run_metadata: HashMap, +} +impl Result<()>> Reader<'_, F> { + fn context(&self, turn: Option<&str>, request: Option<&str>) -> HookContext { + HookContext { + session_id: Some(self.session_id.clone()), + turn_id: turn.map(Into::into), + request_id: request.map(Into::into), + source: Some("muse_export_v1".into()), + ..Default::default() + } + } + fn emit(&mut self, ts: i64, event: MuseEvent) -> Result<()> { + (self.emit)(envelope( + "muse", + self.version.clone(), + &self.session_id, + event.name().as_str(), + ts, + event.payload()?, + )) + } + fn ensure_start(&mut self, ts: i64) -> Result<()> { + if self.first_ts.is_none() { + self.first_ts = Some(ts); + let mut context = self.context(None, None); + context.transcript_path = Some(self.path.to_string_lossy().into_owned()); + context.session = SessionMetadata { + trajectory_id: self.session.trajectory_id.clone().into(), + root_session_id: self.session.root_session_id.clone().into(), + model: self.metadata.model_id.clone().into(), + provider: self.metadata.provider_id.clone().into(), + cwd: self.metadata.workspace_root.clone().into(), + ..Default::default() + }; + self.emit( + ts, + MuseEvent::SessionStart(Hook { + context, + body: Empty {}, + }), + )?; + } + 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>, + bounded: Option, + model_step: Option, + ) -> Result<()> { + let 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 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) + } + }; + run.active_model = Some(ModelState { + request_id: request_id.clone(), + usage: None, + output: None, + response_id: None, + completed_at: None, + }); + let mut context = self.context(Some(run_id), Some(&request_id)); + context.session.model = metadata.model_id.into(); + context.session.provider = metadata.provider_id.into(); + context.session.cwd = self.metadata.workspace_root.clone().into(); + self.emit( + ts, + MuseEvent::PreLLMCall(Hook { + context, + body: ModelStarted { + request_record_id: native_request_id.into(), + details: ModelDetails { + step: model_step.into(), + ..Default::default() + }, + input_summary: bounded.into(), + ..Default::default() + }, + }), + ) + } + 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( + ts, + MuseEvent::UserPromptSubmit(Hook { + context: self.context(Some(run_id), None), + body: PromptSubmitted { prompt: None }, + }), + ) + } + 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(|r| r.active_model.take()) + else { + return Ok(()); + }; + self.emit( + model.completed_at.unwrap_or(ts), + MuseEvent::PostLLMCall(Hook { + context: self.context(Some(run_id), Some(&model.request_id)), + body: ModelFinished { + usage: model.usage, + usage_scope: UsageScope::Model, + output_text_preview: model.output, + response_id: model.response_id.into(), + outcome: NativeOutcome { + error: if model.completed_at.is_some() { + None + } else { + error.map(|e| ErrorDetail::Text(e.into())) + }, + ..Default::default() + }, + ..Default::default() + }, + }), + ) + } + fn process_task(&mut self, task: TaskEvent, at: Option) -> Result<()> { + if let (Some(id), Some(call)) = ( + task.task_id.as_ref(), + task.idempotency_key + .as_deref() + .and_then(|s| s.strip_prefix("tool:")), + ) { + self.tool_tasks.insert(id.clone(), call.into()); + } + if matches!(task.kind, TaskEventKind::Rejected) { + if let (Some(id), Some(reason)) = ( + task.origin.and_then(|o| o.tool_call_id), + task.reason.as_ref(), + ) { + self.tool_errors.insert(id, reason.clone()); + } + } + if let Some(call) = task + .task_id + .as_ref() + .and_then(|id| self.tool_tasks.get(id)) + .cloned() + { + match task.kind { + TaskEventKind::Started => { + if let Some(tool) = self.tools.get(&call) { + let event = MuseEvent::PreToolUse(Hook { + context: self.context(Some(&tool.turn_id), None), + body: tool.tool.clone(), + }); + self.emit( + at.context("Muse tool task has no recorded_at")?.millis(), + event, + )?; + } + } + TaskEventKind::Failed | TaskEventKind::Cancelled => { + self.tool_errors.insert( + call, + task.reason + .unwrap_or_else(|| "Muse tool task failed".into()), + ); + if let Some(id) = task.task_id { + self.tool_tasks.remove(&id); + } + } + TaskEventKind::Completed => { + if let Some(tool) = self.tools.get_mut(&call) { + tool.completed = true; + } + if let Some(id) = task.task_id { + self.tool_tasks.remove(&id); + } + } + _ => {} + } + } + Ok(()) + } + fn process(&mut self, item: ExportRecord) -> Result<()> { + let Some(native) = item.envelope else { + return Ok(()); + }; + let at = item.recorded_at.or(native.recorded_at); + let Some(payload) = native.payload else { + return Ok(()); + }; + let (run_id, event) = match payload { + NativePayload::RunModel { record } => { + if let Some(record) = record { + self.run_metadata + .insert(record.run_stream.id, record.metadata); + } + return Ok(()); + } + NativePayload::Metadata { record } => { + if let Some(record) = record { + self.metadata = record; + } + return Ok(()); + } + NativePayload::Task { event } => { + if let Some(event) = event { + self.process_task(event, at)?; + } + return Ok(()); + } + NativePayload::SessionEnd => { + let ts = at + .context("Muse session_end record has no recorded_at timestamp")? + .millis(); + self.ensure_start(ts)?; + self.current_session_ended = true; + return Ok(()); + } + NativePayload::Run { + event: RunEvent::Ignored, + .. + } + | NativePayload::Ignored => return Ok(()), + NativePayload::Run { run_id, event } => (run_id, event), + }; + if run_id.is_empty() { + bail!("Muse run record has no run_id") + } + let ts = at + .context("Muse run record has no recorded_at timestamp")? + .millis(); + self.ensure_start(ts)?; + self.current_session_ended = false; + self.activity_count += 1; + if !matches!(event, RunEvent::Started { .. }) { + self.ensure_run(&run_id, ts)?; + } + match event { + RunEvent::Started { prompt } => { + 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( + ts, + MuseEvent::UserPromptSubmit(Hook { + context: self.context(Some(&run_id), None), + body: PromptSubmitted { prompt }, + }), + )?; + } + RunEvent::ModelInputTraceRecorded { + request_record_id, + model_step, + bounded, + } => { + if self + .runs + .get(&run_id) + .and_then(|r| r.active_model.as_ref()) + .is_some() + { + self.finish_model(&run_id, ts, None)?; + } + self.start_model(&run_id, ts, Some(&request_record_id), bounded, model_step)?; + } + RunEvent::AssistantToolCallsCommitted { + tool_calls, + response_id, + } => { + let mut output_calls = Vec::new(); + for call in tool_calls { + if call.call_id.is_empty() { + bail!("Muse committed tool call has no call_id") + } + let args = call + .args + .unwrap_or(ToolArguments::Json(serde_json::Value::Null)); + self.tools.insert( + call.call_id.clone(), + PendingTool { + turn_id: run_id.clone(), + tool: ToolStarted { + tool_use_id: Some(call.call_id.clone().into()), + tool_name: Some(call.name.clone()), + tool_input: Some(args.input()), + }, + completed: false, + }, + ); + output_calls.push(AssistantToolCall { + id: call.call_id, + call_type: FunctionCallType::Function, + function: FunctionCall { + name: call.name, + arguments: args, + }, + }); + } + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|r| r.active_model.as_mut()) + { + model.output = Some(ModelOutput::Messages(vec![NativeMessage { + role: MessageRole::Assistant, + content: None, + tool_calls: Some(output_calls), + tool_call_id: None, + name: None, + extensions: BTreeMap::new(), + }])); + model.response_id = response_id; + } + } + RunEvent::ToolResultBatchCommitted { results } => { + for result in results { + if result.tool_call_id.is_empty() { + bail!("Muse tool result has no tool_call_id") + } + let tool = self + .tools + .remove(&result.tool_call_id) + .unwrap_or(PendingTool { + turn_id: run_id.clone(), + tool: ToolStarted { + tool_use_id: Some(result.tool_call_id.clone().into()), + ..Default::default() + }, + completed: false, + }); + let error = self.tool_errors.remove(&result.tool_call_id); + let context = self.context(Some(&tool.turn_id), None); + let body = ToolFinished { + tool: tool.tool, + tool_response: result.text, + outcome: NativeOutcome { + error: error.clone().map(ErrorDetail::Text), + ..Default::default() + }, + ..Default::default() + }; + let event = if error.is_some() { + MuseEvent::PostToolUseFailure(Hook { context, body }) + } else if tool.completed { + MuseEvent::PostToolUse(Hook { context, body }) + } else { + MuseEvent::PostToolBatch(Hook { + context, + body: ToolBatchFinished { + tool_calls: vec![body], + }, + }) + }; + self.emit(ts, event)?; + } + } + RunEvent::ModelFailureTerminalIntent { reason } => self.finish_model( + &run_id, + ts, + Some(reason.as_deref().unwrap_or("Muse model failed")), + )?, + RunEvent::ModelResponseCreated { response_id } => { + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|r| r.active_model.as_mut()) + { + model.response_id = response_id; + } + } + RunEvent::ModelCompleted { usage } => { + if self + .runs + .get(&run_id) + .and_then(|r| r.active_model.as_ref()) + .is_none() + { + self.start_model(&run_id, ts, None, None, None)?; + } + if let Some(model) = self + .runs + .get_mut(&run_id) + .and_then(|r| r.active_model.as_mut()) + { + model.usage = usage; + model.completed_at = Some(ts); + } + } + RunEvent::ResourceUsageSampled => {} + RunEvent::AssistantMessageCommitted { text, response_id } => { + if self + .runs + .get(&run_id) + .and_then(|r| r.active_model.as_ref()) + .is_none() + { + self.start_model(&run_id, ts, None, None, None)?; + } + let run = self + .runs + .get_mut(&run_id) + .context("Muse model start did not create a run")?; + run.last_assistant = Some(text.clone()); + let model = run + .active_model + .as_mut() + .context("Muse run has no active model")?; + model.output = Some(ModelOutput::Text(text)); + model.response_id = response_id; + } + RunEvent::Terminal { terminal, reason } => { + let error = (terminal != "completed").then(|| { + reason + .as_ref() + .and_then(ErrorDetail::render) + .unwrap_or_else(|| terminal.clone()) + }); + self.finish_model(&run_id, ts, error.as_deref())?; + let run = self.runs.remove(&run_id); + self.run_metadata.remove(&run_id); + self.emit( + ts, + MuseEvent::Stop(Hook { + context: self.context(Some(&run_id), None), + body: TurnFinished { + last_assistant_message: run + .and_then(|r| r.last_assistant) + .map(ModelOutput::Text), + terminal: Some(terminal), + outcome: NativeOutcome { + error: error.map(ErrorDetail::Text), + ..Default::default() + }, + }, + }), + )?; + } + RunEvent::Ignored => {} + } + Ok(()) + } + fn finish(mut self, mode: ReadMode) -> Result { + if self.activity_count == 0 { + if matches!(mode, ReadMode::Attach) { + return Ok(self.current_session_ended); + } + bail!( + "Muse export {} has no traceable run events", + self.path.display() + ); + } + let end_ts = self + .last_ts + .ok_or_else(|| anyhow!("Muse export {} has no timestamps", self.path.display()))?; + let ended = self.current_session_ended; + let keep_active_open = match mode { + #[cfg(test)] + ReadMode::Historical => false, + _ => true, + }; + if keep_active_open && !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) + .context("Muse run disappeared while finalizing")?; + self.emit( + end_ts, + MuseEvent::Stop(Hook { + context: self.context(Some(&run_id), None), + body: TurnFinished { + last_assistant_message: run.last_assistant.map(ModelOutput::Text), + outcome: NativeOutcome { + error: Some(ErrorDetail::Text( + "Run ended without a terminal event".into(), + )), + ..Default::default() + }, + ..Default::default() + }, + }), + )?; + } + let error = if !ended { + Some("Session has no terminal end after its latest activity") + } else if self.terminated_abnormally { + Some("Session terminated abnormally") + } else { + None + }; + self.emit( + end_ts, + MuseEvent::SessionEnd(Hook { + context: self.context(None, None), + body: SessionFinished { + reason: (if ended { + self.session + .session_end + .as_ref() + .and_then(|e| e.exit_reason.clone()) + .map(ErrorDetail::Text) + } else { + None + }) + .into(), + outcome: NativeOutcome { + error: error.map(|e| ErrorDetail::Text(e.into())), + ..Default::default() + }, + }, + }), + )?; + Ok(ended) + } +} + +#[cfg(test)] +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 session_id(path: &Path) -> Result { + let header = read_header(path)?; + 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| { + events.push(event); + Ok(()) + })?; + Ok(events) +} + +#[cfg(test)] +pub(crate) fn for_each_envelope_for_session( + path: &Path, + expected_id: &str, + emit: F, +) -> Result<()> +where + F: FnMut(Envelope) -> Result<()>, +{ + for_each_snapshot_envelope_for_session(path, expected_id, ReadMode::Historical, 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, + mode: ReadMode, + 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, + current_session_ended: false, + activity_count: 0, + metadata: ExportMetadata::default(), + 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()))?, + )); + EventsSeed(&mut reader) + .deserialize(&mut de) + .map_err(|error| anyhow!("parse Muse export {}: {error}", path.display()))?; + de.end() + .map_err(|error| anyhow!("parse Muse export {}: {error}", path.display()))?; + reader.finish(mode) +} + +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 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; + loop { + let item = seq.next_element::().map_err(|error| { + serde::de::Error::custom(format!("event {}: {error}", index + 1)) + })?; + let Some(item) = item else { break }; + 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::*; + + 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(&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","terminal":"completed"}}}}, + {"recorded_at": 1_009_000, "envelope":{"payload":{"kind":"session_end"}}} + ] + })); + let events = envelopes_for_session(&path, "session-1").unwrap(); + assert_eq!( + events + .iter() + .map(|event| event.event.as_str()) + .collect::>(), + [ + "SessionStart", + "UserPromptSubmit", + "PreLLMCall", + "PostLLMCall", + "Stop", + "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!({ + "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":{"cpu_time_ms":2,"peak_rss_bytes":4096}}}}}, + {"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!(events[3].payload["usage"].is_null()); + 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: missing field `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" + ); + } + + #[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", ReadMode::Attach, |event| { + events.push(event.event); + Ok(()) + }) + .unwrap(); + assert!(!ended); + assert_eq!(events, ["SessionStart", "UserPromptSubmit", "PreLLMCall"]); + } + + #[test] + fn one_shot_empty_import_fails_even_when_a_session_end_was_observed() { + for events in [ + json!([]), + json!([{"recorded_at":1000,"envelope":{"payload":{"kind":"session_end"}}}]), + ] { + let (_temp, path) = write_export( + json!({"export_schema_version":1,"sessions":[{"session_id":"session-1"}],"events":events}), + ); + let error = for_each_snapshot_envelope_for_session( + &path, + "session-1", + ReadMode::Import, + |_| Ok(()), + ) + .unwrap_err(); + assert!(error.to_string().contains("no traceable run events")); + assert!(for_each_snapshot_envelope_for_session( + &path, + "session-1", + ReadMode::Attach, + |_| Ok(()) + ) + .is_ok()); + } + } + + #[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", ReadMode::Attach, |event| { + events.push(event); + Ok(()) + }) + .unwrap(); + assert!(!ended); + assert!(events.is_empty()); + } +} 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..5a3ae7c1 --- /dev/null +++ b/bt-daemon/src/transcript_import/muse_runtime.rs @@ -0,0 +1,251 @@ +//! Muse exports are streamed through the shared import processor. +use super::*; +use crate::muse::msp::{drain_stderr, MuseMspClient}; +use crate::transcript_import::muse::ReadMode; +use anyhow::Context; +use std::ffi::OsString; +use std::process::Stdio; +use std::time::Duration; +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 = + crate::env::var_os("BRAINTRUST_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(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 = + crate::env::var_os("BRAINTRUST_MUSE_BIN").unwrap_or_else(|| OsString::from("muse")); + let client = MuseMspClient::connect(&executable).await?; + let listing = client.session_ids().await; + let shutdown = client.shutdown().await; + let sessions = listing?; + shutdown?; + Ok(sessions) +} + +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, ReadMode::Import, 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, + mode: ReadMode, + 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, + mode, + |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 { + 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, + ReadMode::Attach, + 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)) => {} + } + } + // 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()) +} diff --git a/bt-daemon/src/transcript_import/pipeline.rs b/bt-daemon/src/transcript_import/pipeline.rs index 7ccfdf12..e3066bce 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,490 @@ 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(()) + } + + 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, } } } #[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", + crate::transcript_import::muse::ReadMode::Attach, + 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", + crate::transcript_import::muse::ReadMode::Attach, + 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", + crate::transcript_import::muse::ReadMode::Attach, + 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()) + && row.pointer("/Merge/end_ms").is_some()) + .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()) + && row.pointer("/Merge/end_ms").is_some()) + .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); + } + + #[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 { 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 d3c87f9a..40780482 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -18,6 +18,7 @@ mod debug; mod decode; mod git; mod grok; +mod muse; mod opencode; mod pi; pub(crate) use claude::successful_subagent_handback_message; @@ -66,6 +67,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)] @@ -96,6 +105,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`. @@ -300,6 +312,7 @@ mod tests { "cursor", "debug", "grok", + "muse", "opencode", "pi" ] @@ -315,6 +328,7 @@ mod tests { ("open-code", "opencode"), ("opencode", "opencode"), ("codex", "codex"), + ("muse-code", "muse"), ("debug", "debug"), ] { assert_eq!(registry.canonical_source(alias), Some(source)); diff --git a/bt-daemon/src/translate/muse.rs b/bt-daemon/src/translate/muse.rs new file mode 100644 index 00000000..ff757aee --- /dev/null +++ b/bt-daemon/src/translate/muse.rs @@ -0,0 +1,892 @@ +//! Muse hook reducer. Native hook decoding and export DTOs live in crate::muse; +//! this module owns span correlation, ordering, and replay policy. +use super::git::GitMetadataCache; +use super::recent::RecentSet; +#[path = "muse_events.rs"] +mod muse_events; +use super::{ + local_username, AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, +}; +use crate::ids; +use crate::muse::hooks::*; +use crate::muse::metadata::*; +#[path = "muse_usage.rs"] +mod muse_usage; +use crate::wire::Envelope; +use muse_usage::UsageTracker; +#[cfg(test)] +use serde_json::json; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::Arc; + +type Result = anyhow::Result; +impl TranslatorFactory for crate::agents::Muse { + fn create(&self, session_id: &str) -> Box { + Box::new(MuseTranslator::new(session_id, GitMetadataCache::shared())) + } +} +#[derive(Clone)] +struct OpenOp { + id: String, + parents: Vec, + turn_span_id: Option, +} +struct MuseTranslator { + session_id: String, + root_id: String, + effective_root_id: String, + external_parent: Option, + opened: bool, + root_created: bool, + session_metadata: SessionMetadata, + session_version: Option, + turns: HashMap, + llms: HashMap, + closed_turns: RecentSet, + closed_llms: RecentSet, + events: muse_events::Events, + usage: UsageTracker, + last_completed_turn: Option, + 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, + session_metadata: SessionMetadata::default(), + session_version: None, + turns: HashMap::new(), + llms: HashMap::new(), + closed_turns: RecentSet::default(), + closed_llms: RecentSet::default(), + root_created: false, + events: muse_events::Events::default(), + usage: UsageTracker::default(), + last_completed_turn: None, + last_ts: 0, + git, + } + } + fn ensure_root(&mut self, e: &Envelope, ctx: &SessionCtx) -> Result> { + if self.opened { + return Ok(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()); + if self.root_created { + return Ok(Vec::new()); + } + let authored = serde_json::to_value(RootMetadata { + source: "muse", + session_id: &ctx.session_id, + username: local_username(), + muse_version: self.session_version.as_deref(), + adapter_version: e.plugin_version.as_deref(), + session: &self.session_metadata, + })?; + // Arbitrary user metadata is an intentional JSON merge boundary. + let mut metadata = ctx + .config + .as_ref() + .and_then(|c| c.additional_metadata.as_ref()) + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + if let Value::Object(authored) = authored { + metadata.extend(authored); + } + self.root_created = true; + Ok(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() + })]) + } + // Values here are already serialized SpanRow fields, never native payloads. + fn close( + root: &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::Merge(SpanRow { + span_id: open.id, + root_span_id: root.into(), + parent_span_ids: open.parents, + turn_span_id: open.turn_span_id, + end_ms: Some(ts), + 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 = self.close_events(ts, error); + for (map, closed) in [ + (&mut self.llms, &mut self.closed_llms), + (&mut self.turns, &mut self.closed_turns), + ] { + 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, + turn_span_id: open.turn_span_id, + end_ms: Some(ts), + error: error.map(str::to_owned), + ..Default::default() + })); + } + } + out + } + fn parent(&self, ids: &HookContext) -> 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(), + turn_span_id: row.turn_span_id.clone(), + }, + ); + vec![SpanOp::Insert(row)] + } + fn unpaired(&self, e: &Envelope, reason: &str) -> Result { + // Raw payload is used only as an identity fingerprint, not inspected. + let identity = format!("unpaired:{}:{}:{}", e.event, e.ts_ms, e.payload); + Ok(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(serde_json::to_value(UnpairedMetadata { unpaired: true })?), + error: Some(reason.into()), + ..Default::default() + })) + } + fn start_model(&mut self, e: &Envelope, hook: Hook) -> Result> { + let Some(id) = hook.context.request() else { + return Ok(vec![self.unpaired(e, "Muse PreLLMCall has no request_id")?]); + }; + let parent = self.parent(&hook.context); + let model = hook.context.session.model.as_deref(); + Ok(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.clone()], + turn_span_id: (parent != self.root_id).then_some(parent), + name: model + .filter(|s| !s.is_empty() && *s != "unknown") + .unwrap_or("llm") + .into(), + span_type: SpanType::Llm, + start_ms: Some(e.ts_ms), + input: Some(serde_json::to_value(ModelInput { + messages: &hook.body.messages, + tools: &hook.body.tools, + })?), + metadata: Some(serde_json::to_value(ModelStartMetadata { + model: &hook.context.session.model, + provider: &hook.context.session.provider, + details: &hook.body.details, + request_record_id: &hook.body.request_record_id, + input_summary: &hook.body.input_summary, + })?), + ..Default::default() + }, + )) + } + fn finish_model(&mut self, e: &Envelope, hook: Hook) -> Result> { + let Some(id) = hook.context.request() else { + return if self.opened { + Ok(vec![self.unpaired(e, "Muse PostLLMCall has no request_id")?]) + } else { + Ok(Vec::new()) + }; + }; + let metrics = hook.body.usage.as_ref().map(|usage| { + self.usage.record( + hook.context.turn(), + id, + usage, + hook.body.usage_scope, + self.llms.contains_key(id), + ) + }); + let mut closed = Self::close( + &self.effective_root_id, + &mut self.llms, + id, + e.ts_ms, + hook.body + .output_text_preview + .as_ref() + .map(serde_json::to_value) + .transpose()?, + metrics.map(serde_json::to_value).transpose()?, + hook.body.outcome.error_message(), + ); + for op in &mut closed { + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op; + row.metadata = Some(serde_json::to_value(ModelEndMetadata { + finish_reason: &hook.body.finish_reason, + response_id: &hook.body.response_id, + tool_call_count: &hook.body.tool_call_count, + status: &hook.body.outcome.status, + options: &hook.body.options, + })?); + } + if !closed.is_empty() { + self.closed_llms.insert(id.into()); + } + Ok(closed) + } + fn finish_turn(&mut self, e: &Envelope, hook: Hook) -> Result> { + let Some(id) = hook.context.turn() else { + return if self.opened { + Ok(vec![self.unpaired(e, "Muse Stop has no turn_id")?]) + } else { + Ok(Vec::new()) + }; + }; + let output = hook + .body + .last_assistant_message + .map(serde_json::to_value) + .transpose()?; + let error = hook.body.outcome.error_message(); + let mut out = Self::close( + &self.effective_root_id, + &mut self.turns, + id, + e.ts_ms, + output.clone(), + None, + error.clone(), + ); + if !out.is_empty() { + self.closed_turns.insert(id.into()); + self.last_completed_turn = Some(id.into()); + out.push(SpanOp::Merge(SpanRow { + span_id: self.root_id.clone(), + root_span_id: self.effective_root_id.clone(), + output, + metrics: Some(serde_json::to_value(&self.usage.totals)?), + error, + 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.into()), + }), + ..Default::default() + })); + } + Ok(out) + } +} +impl AgentTranslator for MuseTranslator { + fn handle(&mut self, e: &Envelope, ctx: &SessionCtx) -> Result> { + self.last_ts = self.last_ts.max(e.ts_ms); + let event = match MuseEvent::decode(e) { + Ok(Some(event)) => event, + Ok(None) => return Ok(Vec::new()), + Err(error) => { + tracing::warn!(source = "muse", event = %e.event, %error, "invalid Muse hook payload"); + return Ok(Vec::new()); + } + }; + if let Some(version) = &e.source_version { + self.session_version = Some(version.clone()); + } + self.session_metadata.update(&event.context().session); + let kind = event.name(); + let ids = event.context(); + if (kind == HookName::UserPromptSubmit + && ids + .turn() + .is_some_and(|id| self.closed_turns.contains(&id.to_owned()))) + || (kind == HookName::PreLLMCall + && ids + .request() + .is_some_and(|id| self.closed_llms.contains(&id.to_owned()))) + { + return Ok(Vec::new()); + } + let mut out = if matches!( + kind, + HookName::UserPromptSubmit + | HookName::PreLLMCall + | HookName::PreToolUse + | HookName::PermissionRequest + | HookName::PostToolUse + | HookName::PostToolUseFailure + | HookName::PostToolBatch + | HookName::SubagentStart + | HookName::PreCompact + ) { + self.ensure_root(e, ctx)? + } else { + Vec::new() + }; + let cwd = event.context().session.cwd.clone(); + match event { + MuseEvent::UserPromptSubmit(hook) => { + if let Some(id) = hook.context.turn() { + let turn_span_id = ids::span_id(&self.session_id, &format!("turn:{id}")); + out.extend(Self::insert_open( + &mut self.turns, + id, + SpanRow { + span_id: turn_span_id.clone(), + root_span_id: self.effective_root_id.clone(), + parent_span_ids: vec![self.root_id.clone()], + name: "Turn".into(), + span_type: SpanType::Task, + start_ms: Some(e.ts_ms), + input: hook.body.prompt.map(serde_json::to_value).transpose()?, + is_turn: true, + turn_span_id: Some(turn_span_id), + ..Default::default() + }, + )); + } else { + out.push(self.unpaired(e, "Muse UserPromptSubmit has no turn_id")?); + } + } + MuseEvent::PreLLMCall(hook) => out.extend(self.start_model(e, hook)?), + MuseEvent::PostLLMCall(hook) => out.extend(self.finish_model(e, hook)?), + MuseEvent::Stop(hook) | MuseEvent::StopFailure(hook) => { + out.extend(self.finish_turn(e, hook)?) + } + MuseEvent::SessionEnd(hook) => { + out.extend(self.close_all(e.ts_ms, Some("Interrupted before completion"))); + if self.opened { + 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: if hook.body.reason.is_missing() { + None + } else { + Some(serde_json::to_value(SessionEndMetadata { + end_reason: &hook.body.reason, + })?) + }, + error: hook.body.outcome.error_message(), + metrics: Some(serde_json::to_value(&self.usage.totals)?), + 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; + } + } + event => out.extend(self.handle_events(e, event)?), + } + for op in &mut out { + let (SpanOp::Insert(row) | SpanOp::Merge(row)) = op; + if let Some(key) = &row.late_merge_key { + let content = serde_json::to_value(CompletionFingerprint { + 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(cwd.as_deref(), &mut out); + Ok(out) + } + fn finalize(&mut self, _ctx: &SessionCtx) -> Result> { + let mut out = self.close_all(self.last_ts, Some("Interrupted before completion")); + if self.opened { + 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), + error: Some("Interrupted before completion".into()), + ..Default::default() + })); + self.opened = false; + } + Ok(out) + } +} +#[cfg(test)] +fn event_error(payload: &Value) -> Option { + serde_json::from_value::(payload.clone()) + .unwrap() + .error_message() +} +#[cfg(test)] +mod tests { + use super::*; + use crate::translate::Registry; + + fn event(name: &str, payload: Value, ts_ms: i64) -> Envelope { + Envelope { + source: "muse".into(), + source_version: Some("1.1.1-R2514.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 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 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", "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 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 { + 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 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 completion_metadata_preserves_absent_null_and_native_status_aliases() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + translator + .handle( + &event("PreLLMCall", json!({"request_id":"r"}), 1), + &context(), + ) + .unwrap(); + let completed = translator + .handle( + &event( + "PostLLMCall", + json!({ + "request_id":"r", "status":"error", "response_id":null, + "finish_reason":null, "options":null, + }), + 2, + ), + &context(), + ) + .unwrap(); + let row = completed + .iter() + .find_map(|op| match op { + SpanOp::Merge(row) => Some(row), + _ => None, + }) + .unwrap(); + assert_eq!( + row.metadata, + Some(json!({"status":"error","response_id":null,"finish_reason":null,"options":null})) + ); + assert_eq!(row.error.as_deref(), Some("Muse operation error")); + let ended = translator + .handle(&event("SessionEnd", json!({"reason":null}), 3), &context()) + .unwrap(); + assert!(ended.iter().any(|op| matches!(op, SpanOp::Merge(row) if row.metadata == Some(json!({"end_reason":null}))))); + } + + #[test] + fn malformed_known_hook_is_a_noop_and_does_not_lose_open_state() { + let mut translator = MuseTranslator::new( + "muse-fixture-session", + Arc::new(GitMetadataCache::default()), + ); + translator + .handle( + &event("PreLLMCall", json!({"request_id":"r"}), 1), + &context(), + ) + .unwrap(); + assert!(translator + .handle( + &event( + "PostLLMCall", + json!({"request_id":"r","usage":{"input_tokens":"wrong"}}), + 2 + ), + &context() + ) + .unwrap() + .is_empty()); + assert!(translator.llms.contains_key("r")); + assert!(translator + .handle( + &event( + "PostLLMCall", + json!({"request_id":"r","usage":{"input_tokens":1}}), + 3 + ), + &context() + ) + .unwrap() + .iter() + .any(|op| matches!(op, SpanOp::Merge(_)))); + } + + #[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":[{"role":"user","content":"a"}],"request_id":42}), + 1000, + ), + &ctx, + ) + .unwrap(); + let second = translator + .handle( + &event( + "PreLLMCall", + json!({"messages":[{"role":"user","content":"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!({"turn_id":"t1","prompt":"x"}), + 1000, + ), + &ctx, + ) + .unwrap(); + let unsupported = translator + .handle( + &event( + "UnrecognizedEvent", + 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(); + 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/src/translate/muse_events.rs b/bt-daemon/src/translate/muse_events.rs new file mode 100644 index 00000000..2e3de370 --- /dev/null +++ b/bt-daemon/src/translate/muse_events.rs @@ -0,0 +1,550 @@ +//! Typed tool, observation, subagent, and compaction reducers. +use super::*; +#[derive(Default)] +pub(super) struct Events { + open: HashMap, + closed: RecentSet, +} +impl MuseTranslator { + pub(super) fn handle_events(&mut self, e: &Envelope, event: MuseEvent) -> Result> { + if !self.opened { + return Ok(Vec::new()); + } + match event { + MuseEvent::PreToolUse(hook) => self.start_tool(e, &hook.context, &hook.body, false), + MuseEvent::PostToolUse(hook) => { + let outcome = hook + .body + .outcome + .error_message() + .map(ToolOutcome::Failure) + .unwrap_or(ToolOutcome::Success); + self.finish_tool(e, &hook.context, hook.body, outcome) + } + MuseEvent::PostToolUseFailure(hook) => { + let outcome = ToolOutcome::Failure( + hook.body + .outcome + .error_message() + .unwrap_or_else(|| "Muse tool failed".into()), + ); + self.finish_tool(e, &hook.context, hook.body, outcome) + } + MuseEvent::PostToolBatch(hook) => { + let mut out = Vec::new(); + for call in hook.body.tool_calls { + let outcome = call + .outcome + .error_message() + .map(ToolOutcome::Failure) + .unwrap_or(ToolOutcome::Unknown); + out.extend(self.finish_tool(e, &hook.context, call, outcome)?); + } + Ok(out) + } + MuseEvent::PermissionRequest(hook) => { + let metadata = serde_json::to_value(PermissionMetadata { + tool_name: &hook.body.tool_name, + })?; + self.observe( + e, + &hook.context, + hook.body.tool_input.map(|v| v.0), + metadata, + "permission-request", + ) + } + MuseEvent::Notification(hook) => { + let metadata = serde_json::to_value(NotificationMetadata { + tool_name: &hook.body.tool.tool_name, + notification_type: &hook.body.notification_type, + title: &hook.body.title, + message: &hook.body.message, + })?; + self.observe( + e, + &hook.context, + hook.body.tool.tool_input.map(|v| v.0), + metadata, + "notification", + ) + } + MuseEvent::SubagentStart(hook) => { + let Some(child) = hook + .body + .child_session_id + .as_deref() + .filter(|s| !s.is_empty()) + else { + return Ok(vec![ + self.unpaired(e, "Muse subagent event has no child identity")? + ]); + }; + // subagent_id names an agent; only child_session_id identifies an invocation. + let key = format!("subagent:{child}"); + let mut row = self.event_row(e, &hook.context, &key); + row.name = hook + .body + .subagent_id + .as_deref() + .unwrap_or("Subagent") + .into(); + row.metadata = Some(serde_json::to_value(SubagentMetadata { + subagent_id: &hook.body.subagent_id, + child_session_id: &hook.body.child_session_id, + })?); + Ok(self.open_event(&key, row)) + } + MuseEvent::SubagentStop(hook) => { + let Some(child) = hook + .body + .child + .child_session_id + .as_deref() + .filter(|s| !s.is_empty()) + else { + return Ok(vec![ + self.unpaired(e, "Muse subagent event has no child identity")? + ]); + }; + Ok(self.finish_event( + &format!("subagent:{child}"), + e.ts_ms, + hook.body + .last_assistant_message + .map(serde_json::to_value) + .transpose()?, + hook.body.outcome.error_message(), + None, + )) + } + MuseEvent::PreCompact(hook) => { + let key = format!( + "compact:{}:{}:{}", + hook.context.turn().unwrap_or("session"), + hook.body.trigger.as_deref().unwrap_or("unknown"), + e.ts_ms + ); + let mut row = self.event_row(e, &hook.context, &key); + row.name = "Context compaction".into(); + row.tags = Some(vec!["compaction".into()]); + row.metadata = Some(serde_json::to_value(CompactMetadata { + trigger: &hook.body.trigger, + native_operation_id_available: false, + })?); + Ok(self.open_event(&key, row)) + } + MuseEvent::PostCompact(hook) => { + let prefix = format!( + "compact:{}:{}:", + hook.context.turn().unwrap_or("session"), + hook.body.trigger.as_deref().unwrap_or("unknown") + ); + // There is no native operation id. Pair only a uniquely matching start. + let matches: Vec<_> = self + .events + .open + .keys() + .filter(|key| key.starts_with(&prefix)) + .cloned() + .collect(); + Ok(if matches.len() == 1 { + self.finish_event( + &matches[0], + e.ts_ms, + None, + hook.body.outcome.error_message(), + None, + ) + } else { + Vec::new() + }) + } + _ => Ok(Vec::new()), + } + } + fn observe( + &mut self, + e: &Envelope, + context: &HookContext, + input: Option, + metadata: Value, + tag: &str, + ) -> Result> { + // Native observation has neither a decision nor an operation duration. + let key = ids::span_id( + &self.session_id, + &format!("observation:{}:{}:{}", e.event, e.ts_ms, e.payload), + ); + if !self.events.closed.insert(key.clone()) { + return Ok(Vec::new()); + } + let mut row = self.event_row(e, context, &key); + row.name = e.event.clone(); + row.end_ms = Some(e.ts_ms); + row.input = input; + row.metadata = Some(metadata); + row.tags = Some(vec![tag.into()]); + Ok(vec![SpanOp::Insert(row)]) + } + fn event_row(&self, e: &Envelope, context: &HookContext, key: &str) -> SpanRow { + let parent = self.parent(context); + 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, + context: &HookContext, + tool: &ToolStarted, + missing_start: bool, + ) -> Result> { + let Some(id) = tool.tool_use_id.as_deref().filter(|s| !s.is_empty()) else { + return Ok(vec![self.unpaired(e, "Muse tool event has no tool_use_id")?]); + }; + let key = format!("tool:{id}"); + let mut row = self.event_row(e, context, &key); + row.name = tool + .tool_name + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or("tool") + .into(); + row.span_type = SpanType::Tool; + row.input = tool.tool_input.as_ref().map(|v| v.0.clone()); + row.metadata = Some(serde_json::to_value(ToolStartMetadata { + tool_name: &row.name, + native_start_missing: missing_start, + })?); + Ok(self.open_event(&key, row)) + } + fn finish_tool( + &mut self, + e: &Envelope, + context: &HookContext, + tool: ToolFinished, + outcome: ToolOutcome, + ) -> Result> { + let Some(id) = tool.tool.tool_use_id.as_deref().filter(|s| !s.is_empty()) else { + return Ok(vec![ + self.unpaired(e, "Muse tool result has no tool_use_id")? + ]); + }; + let key = format!("tool:{id}"); + if self.events.closed.contains(&key) { + return Ok(Vec::new()); + } + let mut out = Vec::new(); + if !self.events.open.contains_key(&key) { + out.extend(self.start_tool(e, context, &tool.tool, true)?); + if let Some(SpanOp::Insert(row)) = out.first_mut() { + row.start_ms = Some( + tool.duration_ms + .map(|d| e.ts_ms.saturating_sub(i64::try_from(d).unwrap_or(i64::MAX))) + .unwrap_or(e.ts_ms), + ); + } + } + let (error, label) = match outcome { + ToolOutcome::Success => (None, Some(ToolOutcomeLabel::Success)), + ToolOutcome::Failure(error) => (Some(error), Some(ToolOutcomeLabel::Error)), + ToolOutcome::Unknown => (None, None), + }; + out.extend( + self.finish_event( + &key, + e.ts_ms, + tool.tool_response.map(|v| v.0), + error, + label + .map(|tool_outcome| serde_json::to_value(ToolEndMetadata { tool_outcome })) + .transpose()?, + ), + ); + Ok(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() + } +} +#[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/src/translate/muse_usage.rs b/bt-daemon/src/translate/muse_usage.rs new file mode 100644 index 00000000..4a889322 --- /dev/null +++ b/bt-daemon/src/translate/muse_usage.rs @@ -0,0 +1,71 @@ +//! Bounded, counted-once Muse reducer state. Native counter DTOs and arithmetic +//! are independent of the translator in crate::muse::usage. +use super::super::recent::RecentMap; +use crate::muse::usage::{TokenMetrics, TokenUsage, UsageScope}; + +#[derive(Default)] +pub(super) struct UsageTracker { + turn_usage: RecentMap, + model_metrics: RecentMap, + pub totals: TokenMetrics, +} +impl UsageTracker { + pub fn record( + &mut self, + turn: Option<&str>, + request: &str, + usage: &TokenUsage, + scope: UsageScope, + request_was_open: bool, + ) -> TokenMetrics { + let mut counters = usage.clone(); + if scope == UsageScope::Turn && request_was_open { + if let Some(turn) = turn { + if let Some(previous) = self.turn_usage.get(turn) { + counters = usage.delta_since(previous); + } + self.turn_usage.insert(turn.to_owned(), usage.clone()); + } + } + let metrics = counters.metrics(); + if request_was_open && self.model_metrics.get(request).is_none() { + self.totals.add(&metrics); + self.model_metrics + .insert(request.to_owned(), metrics.clone()); + } + metrics + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + #[test] + fn cumulative_usage_resets_missing_counters_and_replay() { + let mut tracker = UsageTracker::default(); + let first: TokenUsage = + serde_json::from_value(json!({"input_tokens":10,"output_tokens":5})).unwrap(); + let second: TokenUsage = + serde_json::from_value(json!({"input_tokens":14,"output_tokens":8})).unwrap(); + tracker.record(Some("turn"), "a", &first, UsageScope::Turn, true); + let delta = tracker.record(Some("turn"), "b", &second, UsageScope::Turn, true); + assert_eq!(delta.tokens, Some(7)); + tracker.record(Some("turn"), "b", &second, UsageScope::Turn, false); + assert_eq!(tracker.totals.tokens, Some(22)); + let reset = tracker.record(Some("turn"), "c", &first, UsageScope::Turn, true); + assert_eq!(reset.tokens, Some(15)); + assert_eq!(reset.prompt_cached_tokens, None); + } + #[test] + fn per_model_alias_precedence_and_saturation() { + let usage: TokenUsage = serde_json::from_value(json!({"input_tokens":u64::MAX,"output_tokens":3,"cached_tokens":0,"cache_read_tokens":4})).unwrap(); + let metrics = usage.metrics(); + assert_eq!(metrics.tokens, Some(u64::MAX)); + assert_eq!(metrics.prompt_cached_tokens, Some(0)); + let mut tracker = UsageTracker::default(); + tracker.record(Some("turn"), "a", &usage, UsageScope::Model, true); + tracker.record(Some("turn"), "b", &usage, UsageScope::Model, true); + assert_eq!(tracker.totals.tokens, Some(u64::MAX)); + } +} diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 99962754..4b54ba75 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( 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.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"} 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..3a39d71a --- /dev/null +++ b/bt-daemon/tests/muse_attach.rs @@ -0,0 +1,322 @@ +#![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("BRAINTRUST_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("BRAINTRUST_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() { + let mut stderr = String::new(); + tokio::io::AsyncReadExt::read_to_string( + &mut first.stderr.take().unwrap(), + &mut stderr, + ) + .await + .unwrap(); + panic!("observer exited before active rows: {status}: {stderr}"); + } + 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" + ); + // Exercise the SDK handshake and full SessionListResult against the native + // durable index, then reimport through the same production reader. + let all = native + .command(env!("CARGO_BIN_EXE_bt-daemon")) + .args(["--json", "import", "muse", "--all"]) + .output() + .await + .unwrap(); + assert!( + all.status.success(), + "MSP import all: {}", + String::from_utf8_lossy(&all.stderr) + ); + assert!(String::from_utf8_lossy(&all.stdout).contains(&session)); + 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/muse_managed_startup.rs b/bt-daemon/tests/muse_managed_startup.rs new file mode 100644 index 00000000..56d0dba0 --- /dev/null +++ b/bt-daemon/tests/muse_managed_startup.rs @@ -0,0 +1,105 @@ +#![cfg(all(feature = "cli", unix))] +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::disallowed_methods, + reason = "Isolated subprocess regression fixtures." +)] +use bt_daemon::wire::{SessionRoute, TraceDestination}; +use bt_daemon::{run_traced_with_host, HostInfo, RunArgs, RunHookCommand, RunSource}; +use std::ffi::OsString; +use std::os::unix::fs::PermissionsExt; +use std::path::Path; + +struct Env(&'static str, Option); +impl Env { + fn set(key: &'static str, value: impl AsRef) -> Self { + let previous = std::env::var_os(key); + std::env::set_var(key, value); + Self(key, previous) + } +} +impl Drop for Env { + fn drop(&mut self) { + if let Some(value) = &self.1 { + std::env::set_var(self.0, value); + } else { + std::env::remove_var(self.0); + } + } +} +fn assert_running(marker: &Path) { + let status: serde_json::Value = + serde_json::from_slice(&std::fs::read(marker).unwrap()).unwrap(); + assert_eq!(status["running"], true); +} +#[tokio::test] +async fn muse_starts_fresh_explicit_endpoint_and_standalone_supplies_its_host() { + let tmp = tempfile::tempdir().unwrap(); + let fake = tmp.path().join("muse-override.sh"); + std::fs::write( + &fake, + r#"#!/bin/sh +if [ "$1" = "--version" ]; then echo 'Muse Code 1.4.11'; exit 0; fi +"$MUSE_TEST_DAEMON" --json status --socket "$BRAINTRUST_DAEMON_SOCKET" > "$MUSE_TEST_MARKER" +"#, + ) + .unwrap(); + std::fs::set_permissions(&fake, std::fs::Permissions::from_mode(0o755)).unwrap(); + let binary = env!("CARGO_BIN_EXE_bt-daemon"); + let marker = tmp.path().join("explicit-status.json"); + let socket = tmp.path().join("fresh.sock"); + let _bin = Env::set("BRAINTRUST_MUSE_BIN", &fake); + let _config = Env::set("XDG_CONFIG_HOME", tmp.path().join("config")); + let _data = Env::set("BRAINTRUST_DAEMON_DATA_DIR", tmp.path().join("data")); + let _socket = Env::set("BRAINTRUST_DAEMON_SOCKET", &socket); + let _daemon = Env::set("MUSE_TEST_DAEMON", binary); + let _marker = Env::set("MUSE_TEST_MARKER", &marker); + let result = run_traced_with_host( + RunArgs { + source: RunSource::Muse, + additional_metadata: None, + tags: Vec::new(), + plugin: Vec::new(), + agent_args: vec!["exec".into()], + }, + RunHookCommand { + program: binary.into(), + args: vec!["hook".into()], + }, + SessionRoute { + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("test".into()), + }), + ..Default::default() + }, + HostInfo { + serve_argv: vec![binary.into(), "serve".into(), "--debug-sink".into()], + version: env!("CARGO_PKG_VERSION").into(), + }, + ) + .await; + let _ = bt_daemon::shutdown_daemon(&socket).await; + assert!(result.unwrap().success()); + assert_running(&marker); + + // Exercise the actual standalone entry point with no explicit endpoint. + // This child emits no model events, so no credentials/network are required. + let marker = tmp.path().join("standalone-status.json"); + let output = tokio::process::Command::new(binary) + .args(["run", "--project", "test", "muse", "exec"]) + .env_remove("BRAINTRUST_DAEMON_SOCKET") + .env_remove("BT_DAEMON_SOCKET") + .env("MUSE_TEST_MARKER", &marker) + .current_dir(tmp.path()) + .output() + .await + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + assert_running(&marker); +} diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 58f1bd5c..e6ee0146 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}; @@ -2727,6 +2727,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")] @@ -2828,6 +2835,7 @@ esac args: Vec::new(), }, route(), + dummy_host(), ) }; @@ -2865,6 +2873,216 @@ esac handle.await.unwrap(); } +#[cfg(all(feature = "cli", unix))] +#[tokio::test] +#[ignore = "requires a local Muse Code 1.3+ installation"] +async fn muse_managed_run_preserves_saved_settings_and_foreign_hooks() { + let _credential_backend = EnvVarGuard::set("TBH_CREDENTIAL_BACKEND", "file"); + let _telemetry = EnvVarGuard::set("TBH_DISABLE_TELEMETRY", "1"); + let muse_executable = std::env::var_os("BRAINTRUST_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(); + 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": [{"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 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); + 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(), + "--workspace".into(), + workspace.clone().into_os_string(), + "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"); + 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, settings); + 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()); + 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_executable) + .arg("exec") + .arg("--workspace") + .arg(&workspace) + .args([ + "--provider", + "echo", + "--no-foreign-personal-context", + "ordinary", + ]) + .output() + .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.3+ installation"] +async fn muse_managed_run_starts_its_isolated_daemon_before_sanitized_hooks() { + let _credential_backend = EnvVarGuard::set("TBH_CREDENTIAL_BACKEND", "file"); + let _telemetry = EnvVarGuard::set("TBH_DISABLE_TELEMETRY", "1"); + let muse_executable = std::env::var_os("BRAINTRUST_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"); + for explicit in [false, true] { + 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 socket = tmp.path().join("fresh-muse.sock"); + let _canonical_socket = if explicit { + EnvVarGuard::set("BRAINTRUST_DAEMON_SOCKET", &socket) + } else { + EnvVarGuard::unset("BRAINTRUST_DAEMON_SOCKET") + }; + let data = tmp.path().join("daemon-data"); + let _daemon_data = EnvVarGuard::set("BRAINTRUST_DAEMON_DATA_DIR", &data); + 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(), + "--workspace".into(), + workspace.clone().into_os_string(), + "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()); + if explicit { + shutdown(&socket).await; + assert!( + std::fs::read_dir(data.join("spans")) + .unwrap() + .any(|file| std::fs::read_to_string(file.unwrap().path()) + .unwrap() + .contains("Muse Code")), + "fresh explicit daemon did not accept native hook events" + ); + } + } +} + #[cfg(all(feature = "cli", unix))] struct FlushTrackingDebugSinkFactory { inner: DebugSinkFactory, diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 5ff935a1..9b7d4ee2 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -20,6 +20,148 @@ 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","terminal":"completed"}}}}, + {"recorded_at": 1_005_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-import.ndjson")); + assert_eq!(inserted(&output_rows, "llm"), 1); + assert!(output_rows + .iter() + .any(|row| { row.pointer("/Merge/output").and_then(Value::as_str) == Some("traced") })); + assert!(output_rows.iter().any(|row| { + row.pointer("/Merge/metrics/prompt_tokens") + .and_then(Value::as_i64) + == Some(3) + })); +} + +#[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(); + 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":{"cpu_time_ms":3, "peak_rss_bytes":4096}}}}}, + {"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").is_none() + && 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; 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..4f3ddca7 --- /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("BRAINTRUST_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, diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 00000000..57ef39fe --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "1.99.1" +profile = "minimal" +components = ["clippy", "rustfmt"]