diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index b76b55a..d404440 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -9,7 +9,7 @@ description = "Embeddable Braintrust coding-agent tracing daemon." [features] default = [] # Standalone development/test binary. Production embeds the library in `bt`. -cli = ["dep:tracing-subscriber"] +cli = [] [[bin]] name = "bt-daemon" @@ -28,15 +28,15 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" thiserror = "2" +tempfile = "3" tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util", "sync", "time", "process", "signal", "fs"] } tracing = "0.1" -tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } +tracing-subscriber = { version = "0.3", features = ["env-filter"] } uuid = { version = "1", features = ["v4", "v5"] } [dev-dependencies] axum = "0.8" bytes = "1" reqwest = { version = "0.12", default-features = false, features = ["json", "stream"] } -tempfile = "3" wiremock = "0.6" zstd = "0.13" diff --git a/bt-daemon/README.md b/bt-daemon/README.md index d6b38fc..642c96a 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -16,6 +16,10 @@ One self-contained Cargo crate, liftable to its own repo by copying - `src/wire` — the wire protocol module: envelope types + JSON-RPC framing. - `src/translate` and `src/sink` — agent state machines and Braintrust output. - `src/lib.rs` — the embeddable library: clap `Args` + async entry points +- `src/trace_command.rs`, `src/trace_runtime.rs`, and `src/setup.rs` — the + complete mounted `bt trace` command schema, dispatch, daemon lifecycle, and + agent-specific persistent setup behavior. Hosts supply only credential and + destination-resolution services. (`run_serve`, `run_hook`, `run_status`, `run_import`, `run_traced`). This is what `bt` depends on. - `src/main.rs` — the standalone **`bt-daemon` binary**, compiled only with diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index ac53f71..7b7dc30 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -19,7 +19,10 @@ mod ids; mod journal; mod server; mod settings; +mod setup; mod sink; +mod trace_command; +mod trace_runtime; mod transcript_import; mod translate; mod transport; @@ -30,7 +33,10 @@ pub use command_output::{ OutputFormat, SetupCommandOutput, StatusCommandOutput, StopCommandOutput, TraceCommandOutput, }; pub use server::{AuthLease, AuthProvider, AuthResolveReason, ServeOptions}; +pub use setup::run_setup; pub use sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, Sink, SinkFactory}; +pub use trace_command::{SetupAgent, SetupArgs, StopArgs, TraceArgs, TraceCommand}; +pub use trace_runtime::{run_trace, TraceHostContext, TraceHostServices}; pub use translate::{ AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, }; diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs new file mode 100644 index 0000000..af7a26b --- /dev/null +++ b/bt-daemon/src/setup.rs @@ -0,0 +1,545 @@ +//! Persistent installation and configuration for coding-agent tracing plugins. + +use crate::paths; +use crate::trace_command::{SetupAgent, SetupArgs}; +use crate::wire::SessionRoute; +use crate::TraceCommandOutput; +use anyhow::{bail, Context}; +use serde_json::{Map, Value}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::process::Command as ProcessCommand; + +const CODEX_MARKETPLACE: &str = "braintrust-codex-plugins"; +const CODEX_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-codex-plugin"; +const CODEX_PLUGIN: &str = "trace-codex@braintrust-codex-plugins"; +const CLAUDE_MARKETPLACE: &str = "braintrust-claude-plugin"; +const CLAUDE_MARKETPLACE_SOURCE: &str = "braintrustdata/braintrust-claude-plugin"; +const CLAUDE_PLUGIN: &str = "trace-claude-code@braintrust-claude-plugin"; +const OPENCODE_PLUGIN: &str = "@braintrust/trace-opencode@^1"; +const PI_PLUGIN: &str = "npm:@braintrust/pi-extension@^1"; + +trait CommandRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result; + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()>; +} + +struct SystemCommandRunner; + +impl CommandRunner for SystemCommandRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { + let output = ProcessCommand::new(program) + .args(args) + .output() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("`{program} {}` failed: {}", args.join(" "), stderr.trim()); + } + serde_json::from_slice(&output.stdout) + .with_context(|| format!("`{program} {}` returned invalid JSON", args.join(" "))) + } + + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { + let status = ProcessCommand::new(program) + .args(args) + .status() + .with_context(|| { + format!("failed to run `{program}`; install {program} and ensure it is on PATH") + })?; + if !status.success() { + bail!("`{program} {}` failed with {status}", args.join(" ")); + } + Ok(()) + } +} + +fn github_repo_matches(source: &str, expected: &str) -> bool { + let source = source.trim().trim_end_matches('/'); + let source = source.strip_suffix(".git").unwrap_or(source); + let source = source + .strip_prefix("https://github.com/") + .or_else(|| source.strip_prefix("git@github.com:")) + .unwrap_or(source); + source == expected +} + +fn codex_marketplace(value: &Value) -> Option<&Value> { + value + .get("marketplaces") + .and_then(Value::as_array)? + .iter() + .find(|item| item.get("name").and_then(Value::as_str) == Some(CODEX_MARKETPLACE)) +} + +fn codex_marketplace_is_published(item: &Value) -> bool { + item.get("marketplaceSource") + .and_then(|source| source.get("source")) + .and_then(Value::as_str) + .is_some_and(|source| github_repo_matches(source, CODEX_MARKETPLACE_SOURCE)) +} + +fn setup_codex(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let marketplaces = runner.json("codex", &["plugin", "marketplace", "list", "--json"])?; + match codex_marketplace(&marketplaces) { + Some(marketplace) if codex_marketplace_is_published(marketplace) => runner.run( + "codex", + &["plugin", "marketplace", "upgrade", CODEX_MARKETPLACE], + )?, + Some(_) => { + runner.run( + "codex", + &["plugin", "marketplace", "remove", CODEX_MARKETPLACE], + )?; + runner.run( + "codex", + &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], + )?; + } + None => runner.run( + "codex", + &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], + )?, + } + + // Adding is idempotent and reconciles the installed cache to the refreshed + // marketplace snapshot. + runner.run("codex", &["plugin", "add", CODEX_PLUGIN]) +} + +fn claude_marketplace(value: &Value) -> Option<&Value> { + value + .as_array()? + .iter() + .find(|item| item.get("name").and_then(Value::as_str) == Some(CLAUDE_MARKETPLACE)) +} + +fn claude_marketplace_is_published(item: &Value) -> bool { + item.get("source").and_then(Value::as_str) == Some("github") + && item + .get("repo") + .and_then(Value::as_str) + .is_some_and(|repo| github_repo_matches(repo, CLAUDE_MARKETPLACE_SOURCE)) +} + +fn claude_plugin(value: &Value) -> Option<&Value> { + value + .as_array()? + .iter() + .find(|item| item.get("id").and_then(Value::as_str) == Some(CLAUDE_PLUGIN)) +} + +fn setup_claude(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + let marketplaces = runner.json("claude", &["plugin", "marketplace", "list", "--json"])?; + let marketplace_replaced = match claude_marketplace(&marketplaces) { + Some(marketplace) if claude_marketplace_is_published(marketplace) => { + runner.run( + "claude", + &["plugin", "marketplace", "update", CLAUDE_MARKETPLACE], + )?; + false + } + Some(_) => { + runner.run( + "claude", + &["plugin", "marketplace", "remove", CLAUDE_MARKETPLACE], + )?; + runner.run( + "claude", + &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], + )?; + true + } + None => { + runner.run( + "claude", + &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], + )?; + false + } + }; + + // Claude removes a marketplace's installed plugins when that marketplace + // is removed, so replacing a stale source requires a fresh installation. + if marketplace_replaced { + return runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]); + } + + let plugins = runner.json("claude", &["plugin", "list", "--json"])?; + match claude_plugin(&plugins) { + None => runner.run("claude", &["plugin", "install", CLAUDE_PLUGIN]), + Some(plugin) => { + runner.run("claude", &["plugin", "update", CLAUDE_PLUGIN])?; + if plugin.get("enabled").and_then(Value::as_bool) == Some(false) { + runner.run("claude", &["plugin", "enable", CLAUDE_PLUGIN])?; + } + Ok(()) + } + } +} + +fn load_object(path: &Path) -> anyhow::Result> { + match std::fs::read(path) { + Ok(raw) => { + let value: Value = serde_json::from_slice(&raw) + .with_context(|| format!("invalid JSON configuration: {}", path.display()))?; + value.as_object().cloned().ok_or_else(|| { + anyhow::anyhow!("configuration must be a JSON object: {}", path.display()) + }) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Map::new()), + Err(error) => { + Err(error).with_context(|| format!("failed to read configuration: {}", path.display())) + } + } +} + +fn write_object_atomic(path: &Path, object: Map) -> anyhow::Result<()> { + let parent = path + .parent() + .ok_or_else(|| anyhow::anyhow!("configuration path has no parent: {}", path.display()))?; + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create configuration directory: {}", + parent.display() + ) + })?; + let mut encoded = serde_json::to_string_pretty(&Value::Object(object))?; + encoded.push('\n'); + let mut temporary = tempfile::NamedTempFile::new_in(parent) + .with_context(|| format!("failed to create temporary file in {}", parent.display()))?; + temporary.write_all(encoded.as_bytes()).with_context(|| { + format!( + "failed to write temporary configuration for {}", + path.display() + ) + })?; + temporary + .persist(path) + .map_err(|error| error.error) + .with_context(|| format!("failed to replace configuration: {}", path.display()))?; + Ok(()) +} + +fn setup_opencode_at(path: &Path) -> anyhow::Result<()> { + let mut config = load_object(path)?; + let plugins = config + .entry("plugin") + .or_insert_with(|| Value::Array(Vec::new())) + .as_array_mut() + .ok_or_else(|| { + anyhow::anyhow!( + "OpenCode `plugin` config must be an array: {}", + path.display() + ) + })?; + plugins.retain(|plugin| { + plugin.as_str().is_none_or(|plugin| { + plugin != "@braintrust/trace-opencode" + && !plugin.starts_with("@braintrust/trace-opencode@") + }) + }); + plugins.push(Value::String(OPENCODE_PLUGIN.into())); + write_object_atomic(path, config) +} + +fn setup_opencode() -> anyhow::Result<()> { + let settings_path = paths::agent_settings_path("opencode", None); + let path = settings_path + .parent() + .unwrap_or_else(|| Path::new(".")) + .join("opencode.json"); + setup_opencode_at(&path) +} + +fn setup_pi(runner: &mut impl CommandRunner) -> anyhow::Result<()> { + runner.run("pi", &["install", PI_PLUGIN]) +} + +fn enable_tracing_at(path: &Path, route: SessionRoute) -> anyhow::Result<()> { + let mut settings = load_object(path)?; + settings.insert("trace_to_braintrust".into(), Value::Bool(true)); + settings.insert("route".into(), serde_json::to_value(route)?); + settings.remove("traceToBraintrust"); + settings.remove("project"); + write_object_atomic(path, settings)?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)) + .with_context(|| format!("failed to protect agent settings: {}", path.display()))?; + } + Ok(()) +} + +fn enable_tracing(source: &str, route: SessionRoute) -> anyhow::Result { + let path = paths::agent_settings_path(source, None); + enable_tracing_at(&path, route)?; + Ok(path) +} + +/// Install or refresh one agent's published tracing adapter and persist its +/// non-secret route selection. +pub fn run_setup(args: SetupArgs, route: SessionRoute) -> anyhow::Result { + let mut runner = SystemCommandRunner; + let (source, display_name) = match args.agent { + SetupAgent::Codex => { + setup_codex(&mut runner)?; + ("codex", "Codex") + } + SetupAgent::Claude => { + setup_claude(&mut runner)?; + ("claude", "Claude Code") + } + SetupAgent::OpenCode => { + setup_opencode()?; + ("opencode", "OpenCode") + } + SetupAgent::Pi => { + setup_pi(&mut runner)?; + ("pi", "Pi") + } + }; + let settings_path = enable_tracing(source, route)?; + Ok(TraceCommandOutput::setup( + source, + display_name, + settings_path, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::wire::{AuthSelection, TraceDestination}; + use std::collections::VecDeque; + + struct FakeRunner { + responses: VecDeque, + calls: Vec, + } + + impl FakeRunner { + fn new(responses: impl IntoIterator) -> Self { + Self { + responses: responses.into_iter().collect(), + calls: Vec::new(), + } + } + + fn called(&self, command: &str) -> bool { + self.calls.iter().any(|call| call == command) + } + } + + impl CommandRunner for FakeRunner { + fn json(&mut self, program: &str, args: &[&str]) -> anyhow::Result { + self.calls.push(format!("{program} {}", args.join(" "))); + self.responses + .pop_front() + .ok_or_else(|| anyhow::anyhow!("missing fake JSON response")) + } + + fn run(&mut self, program: &str, args: &[&str]) -> anyhow::Result<()> { + self.calls.push(format!("{program} {}", args.join(" "))); + Ok(()) + } + } + + #[test] + fn codex_installs_from_the_published_marketplace_when_missing() { + let mut runner = FakeRunner::new([serde_json::json!({"marketplaces": []})]); + + setup_codex(&mut runner).unwrap(); + + assert!( + runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") + ); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + } + + #[test] + fn codex_refreshes_the_published_marketplace_and_plugin() { + let mut runner = FakeRunner::new([serde_json::json!({ + "marketplaces": [{ + "name": CODEX_MARKETPLACE, + "marketplaceSource": { + "sourceType": "github", + "source": CODEX_MARKETPLACE_SOURCE + } + }] + })]); + + setup_codex(&mut runner).unwrap(); + + assert!(runner.called("codex plugin marketplace upgrade braintrust-codex-plugins")); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + assert!(!runner.called("codex plugin marketplace remove braintrust-codex-plugins")); + } + + #[test] + fn codex_replaces_a_same_name_local_marketplace() { + let mut runner = FakeRunner::new([serde_json::json!({ + "marketplaces": [{ + "name": CODEX_MARKETPLACE, + "marketplaceSource": {"sourceType": "local", "source": "/tmp/stale"} + }] + })]); + + setup_codex(&mut runner).unwrap(); + + assert!(runner.called("codex plugin marketplace remove braintrust-codex-plugins")); + assert!( + runner.called("codex plugin marketplace add braintrustdata/braintrust-codex-plugin") + ); + assert!(runner.called("codex plugin add trace-codex@braintrust-codex-plugins")); + } + + #[test] + fn claude_installs_from_the_published_marketplace_when_missing() { + let mut runner = FakeRunner::new([serde_json::json!([]), serde_json::json!([])]); + + setup_claude(&mut runner).unwrap(); + + assert!( + runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") + ); + assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); + } + + #[test] + fn claude_refreshes_the_published_marketplace_and_plugin() { + let mut runner = FakeRunner::new([ + serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "github", + "repo": CLAUDE_MARKETPLACE_SOURCE + }]), + serde_json::json!([{ + "id": CLAUDE_PLUGIN, + "version": "1.4.4", + "enabled": true + }]), + ]); + + setup_claude(&mut runner).unwrap(); + + assert!(runner.called("claude plugin marketplace update braintrust-claude-plugin")); + assert!(runner.called("claude plugin update trace-claude-code@braintrust-claude-plugin")); + assert!(!runner.called("claude plugin marketplace remove braintrust-claude-plugin")); + } + + #[test] + fn claude_replaces_a_same_name_local_marketplace() { + let mut runner = FakeRunner::new([serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "directory", + "path": "/tmp/stale" + }])]); + + setup_claude(&mut runner).unwrap(); + + assert!(runner.called("claude plugin marketplace remove braintrust-claude-plugin")); + assert!( + runner.called("claude plugin marketplace add braintrustdata/braintrust-claude-plugin") + ); + assert!(runner.called("claude plugin install trace-claude-code@braintrust-claude-plugin")); + } + + #[test] + fn claude_updates_then_enables_a_disabled_plugin() { + let mut runner = FakeRunner::new([ + serde_json::json!([{ + "name": CLAUDE_MARKETPLACE, + "source": "github", + "repo": CLAUDE_MARKETPLACE_SOURCE + }]), + serde_json::json!([{"id": CLAUDE_PLUGIN, "enabled": false}]), + ]); + + setup_claude(&mut runner).unwrap(); + + let update = runner + .calls + .iter() + .position(|call| { + call == "claude plugin update trace-claude-code@braintrust-claude-plugin" + }) + .unwrap(); + let enable = runner + .calls + .iter() + .position(|call| { + call == "claude plugin enable trace-claude-code@braintrust-claude-plugin" + }) + .unwrap(); + assert!(update < enable); + } + + #[test] + fn opencode_reconciles_the_published_plugin_and_preserves_config() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("opencode.json"); + std::fs::write( + &path, + r#"{"plugin":["other","@braintrust/trace-opencode@0.9.0"],"model":"test/model"}"#, + ) + .unwrap(); + + setup_opencode_at(&path).unwrap(); + + let config: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert_eq!(config["model"], "test/model"); + assert_eq!( + config["plugin"], + serde_json::json!(["other", "@braintrust/trace-opencode@^1"]) + ); + } + + #[test] + fn pi_installs_the_published_extension_range() { + let mut runner = FakeRunner::new([]); + + setup_pi(&mut runner).unwrap(); + + assert!(runner.called("pi install npm:@braintrust/pi-extension@^1")); + } + + #[test] + fn tracing_settings_preserve_unrelated_fields_and_remove_legacy_keys() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write( + &path, + r#"{"traceToBraintrust":false,"project":"old","auth":{"type":"legacy"}}"#, + ) + .unwrap(); + let route = SessionRoute { + auth: AuthSelection { + profile: Some("work".into()), + org_name: Some("Braintrust SDKs".into()), + }, + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("coding-agents".into()), + }), + ..SessionRoute::default() + }; + + enable_tracing_at(&path, route).unwrap(); + + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!(settings["trace_to_braintrust"], true); + assert_eq!(settings["route"]["auth"]["profile"], "work"); + assert_eq!( + settings["route"]["destination"]["project_name"], + "coding-agents" + ); + assert_eq!(settings["auth"]["type"], "legacy"); + assert!(settings.get("traceToBraintrust").is_none()); + assert!(settings.get("project").is_none()); + } +} diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs new file mode 100644 index 0000000..d41e501 --- /dev/null +++ b/bt-daemon/src/trace_command.rs @@ -0,0 +1,66 @@ +//! Command definitions mounted by the host CLI. +//! +//! Coding-agent command names, aliases, and argument shapes live with the +//! integrations they control. Hosts such as `bt` provide global auth flags and +//! dispatch these commands without duplicating agent-specific CLI knowledge. + +use crate::{HookArgs, ImportArgs, RunArgs, ServeArgs, StatusArgs}; +use clap::{Args, Subcommand}; +use std::path::PathBuf; + +#[derive(Debug, Clone, Args)] +pub struct TraceArgs { + #[command(subcommand)] + pub command: TraceCommand, +} + +#[derive(Debug, Clone, Subcommand)] +// Clap argument structs are parsed once; keeping their natural shapes is +// clearer than boxing individual command variants for stack-size savings. +#[allow(clippy::large_enum_variant)] +pub enum TraceCommand { + /// Install the published Braintrust tracing plugin for a coding agent. + Setup(SetupArgs), + /// Run the tracing daemon (foreground). + #[command(hide = true)] + Daemon(ServeArgs), + /// Forward one coding-agent hook event (read from stdin) to the daemon. + #[command(hide = true)] + Hook(HookArgs), + /// Print daemon/session status. + #[command(hide = true)] + Status(StatusArgs), + /// Gracefully stop the tracing daemon. + #[command(hide = true)] + Stop(StopArgs), + /// Import a past Codex or Claude Code session by its resume id. + Import(ImportArgs), + /// Launch a coding agent with tracing enabled for this invocation. + Run(RunArgs), +} + +#[derive(Debug, Clone, Args)] +pub struct StopArgs { + /// Socket path override (default: see the daemon protocol documentation). + #[arg(long)] + pub socket: Option, +} + +#[derive(Debug, Clone, Args)] +pub struct SetupArgs { + #[command(subcommand)] + pub agent: SetupAgent, +} + +#[derive(Debug, Clone, Copy, Subcommand)] +pub enum SetupAgent { + /// Install the published Codex tracing plugin. + Codex, + /// Install the published Claude Code tracing plugin. + Claude, + /// Configure the published OpenCode tracing plugin. + #[command(name = "opencode", alias = "open-code")] + OpenCode, + /// Install the published Pi tracing extension. + Pi, +} diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs new file mode 100644 index 0000000..474127a --- /dev/null +++ b/bt-daemon/src/trace_runtime.rs @@ -0,0 +1,413 @@ +//! Complete execution of the mounted `bt trace` namespace. +//! +//! The embedding CLI supplies only host services for Braintrust credentials +//! and destination selection. Command dispatch, daemon lifecycle, hook +//! behavior, setup, managed runs, imports, and output contracts stay here with +//! the coding-agent integrations. + +use crate::trace_command::TraceCommand; +use crate::wire::{AuthSelection, SessionConfig, SessionRoute}; +use crate::{ + braintrust_serve_options, paths, run_hook, run_import, run_serve, run_setup, run_status, + run_traced, shutdown_daemon, AuthLease, AuthProvider, AuthResolveReason, BraintrustSinkConfig, + HostInfo, OutputFormat, Registry, RunHookCommand, ServeOptions, StatusArgs, TraceArgs, + TraceCommandOutput, +}; +use async_trait::async_trait; +use std::ffi::OsString; +use std::sync::Arc; + +/// Host-owned services used by the integration runtime. +/// +/// Implementations resolve Braintrust profiles and destination choices but do +/// not dispatch or interpret coding-agent commands. +#[async_trait] +pub trait TraceHostServices: Send + Sync { + /// Resolve the non-secret route selected by the host. Commands such as + /// setup and managed run require a destination; hooks may use the current + /// selection without prompting. + async fn resolve_route(&self, destination_required: bool) -> anyhow::Result; + + /// Resolve a Braintrust credential lease without exposing credentials to + /// plugins, settings files, journals, or command arguments. + async fn resolve_auth( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result; +} + +/// Everything the plugin runtime needs from its embedding CLI. +pub struct TraceHostContext { + pub version: String, + pub output_format: OutputFormat, + pub verbose: bool, + /// Program and arguments that enter the mounted trace namespace. For `bt` + /// this is the current executable followed by `trace`. + pub command: RunHookCommand, + pub services: Arc, +} + +struct HostAuthProvider { + services: Arc, +} + +#[async_trait] +impl AuthProvider for HostAuthProvider { + async fn resolve( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result { + self.services.resolve_auth(selection, reason).await + } +} + +fn serve_options(host: &TraceHostContext) -> ServeOptions { + let cfg = BraintrustSinkConfig { + api_url: None, + app_url: None, + version: host.version.clone(), + }; + let mut options = + braintrust_serve_options(&host.version, cfg, Arc::new(Registry::default_agents())); + options.auth_provider = Some(Arc::new(HostAuthProvider { + services: host.services.clone(), + })); + options +} + +fn child_command(command: &RunHookCommand, child: &str) -> RunHookCommand { + let mut args = command.args.clone(); + args.push(OsString::from(child)); + RunHookCommand { + program: command.program.clone(), + args, + } +} + +fn host_info(host: &TraceHostContext) -> HostInfo { + let command = child_command(&host.command, "daemon"); + let mut serve_argv = Vec::with_capacity(command.args.len() + 1); + serve_argv.push(command.program); + serve_argv.extend(command.args); + HostInfo { + serve_argv, + version: host.version.clone(), + } +} + +fn init_daemon_logging(verbose: bool) { + let fallback = if verbose { "debug" } else { "info" }; + let filter = tracing_subscriber::EnvFilter::new(fallback); + if let Err(error) = tracing_subscriber::fmt() + .with_env_filter(filter) + .with_writer(std::io::stderr) + .try_init() + { + eprintln!("bt trace daemon logging unavailable: {error}"); + } +} + +async fn session_config( + host: &TraceHostContext, + route: &SessionRoute, +) -> anyhow::Result { + let lease = host + .services + .resolve_auth(&route.auth, AuthResolveReason::Initial) + .await?; + Ok(SessionConfig { + auth: lease.auth, + destination: route.destination.clone(), + flush_mode: route.flush_mode, + additional_metadata: route.additional_metadata.clone(), + }) +} + +fn print_output(output: TraceCommandOutput, format: OutputFormat) -> anyhow::Result<()> { + println!("{}", output.render(format)?); + Ok(()) +} + +/// Execute the complete mounted trace command. +pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Result<()> { + match args.command { + TraceCommand::Setup(setup_args) => { + let route = host.services.resolve_route(true).await?; + print_output(run_setup(setup_args, route)?, host.output_format) + } + TraceCommand::Daemon(serve_args) => { + init_daemon_logging(host.verbose); + run_serve(serve_args, serve_options(&host)).await + } + TraceCommand::Hook(hook_args) => { + // A persistent hook must never fail the coding agent's turn. + let result = async { + let route = host.services.resolve_route(false).await?; + run_hook(hook_args, route, host_info(&host)).await + } + .await; + if let Err(error) = result { + eprintln!("bt trace hook (non-fatal): {error}"); + } + Ok(()) + } + TraceCommand::Status(status_args) => print_output( + TraceCommandOutput::status(run_status(status_args).await?), + host.output_format, + ), + TraceCommand::Stop(stop_args) => { + let socket = paths::socket_path(stop_args.socket.as_deref()); + let status_args = StatusArgs { + socket: Some(socket.clone()), + session_id: None, + }; + if run_status(status_args).await?.is_none() { + return print_output(TraceCommandOutput::stop(false, false), host.output_format); + } + shutdown_daemon(&socket).await?; + print_output(TraceCommandOutput::stop(true, true), host.output_format) + } + TraceCommand::Import(import_args) => { + let route = host + .services + .resolve_route(import_args.destination.is_none() && import_args.parent.is_none()) + .await?; + let config = session_config(&host, &route).await?; + run_import(import_args, serve_options(&host), Some(config)).await + } + TraceCommand::Run(run_args) => { + let route = host.services.resolve_route(true).await?; + let hook_command = child_command(&host.command, "hook"); + let status = run_traced(run_args, hook_command, route).await?; + if status.success() { + Ok(()) + } else { + anyhow::bail!("coding agent exited with {status}") + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::trace_command::{SetupAgent, SetupArgs, StopArgs}; + use crate::wire::{BackendAuth, TraceDestination}; + use crate::{ImportArgs, ImportSource, RunArgs, RunSource, StatusArgs, TraceCommand}; + use std::path::PathBuf; + use std::sync::Mutex; + + #[test] + fn mounted_child_commands_preserve_the_host_prefix() { + let mounted = RunHookCommand { + program: OsString::from("/path with spaces/bt"), + args: vec![OsString::from("trace")], + }; + assert_eq!( + child_command(&mounted, "hook"), + RunHookCommand { + program: OsString::from("/path with spaces/bt"), + args: vec![OsString::from("trace"), OsString::from("hook")], + } + ); + let context = TraceHostContext { + version: "test".into(), + output_format: OutputFormat::Human, + verbose: false, + command: mounted, + services: Arc::new(PanicHost), + }; + assert_eq!( + host_info(&context).serve_argv, + ["/path with spaces/bt", "trace", "daemon"] + ); + } + + struct PanicHost; + + #[async_trait] + impl TraceHostServices for PanicHost { + async fn resolve_route(&self, _: bool) -> anyhow::Result { + panic!("host service should not be called") + } + + async fn resolve_auth( + &self, + _: &AuthSelection, + _: AuthResolveReason, + ) -> anyhow::Result { + panic!("host service should not be called") + } + } + + struct RecordingHost { + route_requests: Mutex>, + route_error: Option<&'static str>, + auth_error: Option<&'static str>, + } + + impl RecordingHost { + fn new(route_error: Option<&'static str>, auth_error: Option<&'static str>) -> Self { + Self { + route_requests: Mutex::new(Vec::new()), + route_error, + auth_error, + } + } + } + + #[async_trait] + impl TraceHostServices for RecordingHost { + async fn resolve_route(&self, destination_required: bool) -> anyhow::Result { + self.route_requests + .lock() + .unwrap() + .push(destination_required); + if let Some(error) = self.route_error { + anyhow::bail!(error); + } + Ok(SessionRoute { + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("test-project".into()), + }), + ..SessionRoute::default() + }) + } + + async fn resolve_auth( + &self, + _: &AuthSelection, + _: AuthResolveReason, + ) -> anyhow::Result { + if let Some(error) = self.auth_error { + anyhow::bail!(error); + } + Ok(AuthLease { + profile: "test".into(), + auth: BackendAuth { + token: "secret".into(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + expires_at_ms: None, + }) + } + } + + fn test_host(services: Arc) -> TraceHostContext { + TraceHostContext { + version: "test".into(), + output_format: OutputFormat::Json, + verbose: false, + command: RunHookCommand { + program: OsString::from("bt"), + args: vec![OsString::from("trace")], + }, + services, + } + } + + #[tokio::test] + async fn setup_and_run_require_a_host_resolved_destination() { + for command in [ + TraceCommand::Setup(SetupArgs { + agent: SetupAgent::OpenCode, + }), + TraceCommand::Run(RunArgs { + source: RunSource::Codex, + agent_args: Vec::new(), + }), + ] { + let services = Arc::new(RecordingHost::new(Some("no destination"), None)); + let error = run_trace(TraceArgs { command }, test_host(services.clone())) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "no destination"); + assert_eq!(*services.route_requests.lock().unwrap(), [true]); + } + } + + #[tokio::test] + async fn import_only_requires_a_default_destination_without_an_override() { + for (destination, required) in [ + (None, true), + ( + Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("override".into()), + }), + false, + ), + ] { + let services = Arc::new(RecordingHost::new(None, Some("stop before lookup"))); + let args = ImportArgs { + source: ImportSource::Codex, + session_id: "00000000-0000-0000-0000-000000000000".into(), + destination, + parent: None, + attach: false, + }; + let error = run_trace( + TraceArgs { + command: TraceCommand::Import(args), + }, + test_host(services.clone()), + ) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "stop before lookup"); + assert_eq!(*services.route_requests.lock().unwrap(), [required]); + } + } + + #[tokio::test] + async fn hook_host_failures_are_non_fatal() { + let services = Arc::new(RecordingHost::new(Some("route unavailable"), None)); + let args = crate::HookArgs { + source: "codex".into(), + source_version: None, + socket: None, + session_id_field: "session_id".into(), + event_field: "hook_event_name".into(), + event: None, + no_spawn: false, + flush_on_turn_end: false, + flush_timeout_ms: 10_000, + additional_metadata: None, + managed_run_hook: false, + }; + run_trace( + TraceArgs { + command: TraceCommand::Hook(args), + }, + test_host(services.clone()), + ) + .await + .unwrap(); + assert_eq!(*services.route_requests.lock().unwrap(), [false]); + } + + #[tokio::test] + async fn status_and_absent_stop_do_not_resolve_host_state() { + let temp = tempfile::tempdir().unwrap(); + let missing_socket = temp.path().join("missing.sock"); + for command in [ + TraceCommand::Status(StatusArgs { + socket: Some(missing_socket.clone()), + session_id: None, + }), + TraceCommand::Stop(StopArgs { + socket: Some(PathBuf::from(&missing_socket)), + }), + ] { + run_trace(TraceArgs { command }, test_host(Arc::new(PanicHost))) + .await + .unwrap(); + } + } +}