diff --git a/src/agents.rs b/src/agents.rs index 9482f6c..3c4d6d0 100644 --- a/src/agents.rs +++ b/src/agents.rs @@ -13,7 +13,7 @@ use std::sync::Arc; use anyhow::{bail, Context}; use clap::{Args, Subcommand}; -use serde_json::{Map, Value}; +use serde_json::{json, Map, Value}; use bt_daemon::wire::{BackendAuth, FlushMode, SessionConfig}; use bt_daemon::{ @@ -71,12 +71,57 @@ enum SetupAgent { Claude, } -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 GENERATED_MARKETPLACE: &str = "braintrust-bt-trace"; +const GENERATED_CODEX_PLUGIN: &str = "trace-codex@braintrust-bt-trace"; +const GENERATED_CLAUDE_PLUGIN: &str = "trace-claude-code@braintrust-bt-trace"; +const LEGACY_CODEX_PLUGIN: &str = "trace-codex@braintrust-codex-plugins"; +const LEGACY_CLAUDE_PLUGIN: &str = "trace-claude-code@braintrust-claude-plugin"; + +const CODEX_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreCompact", + "PostCompact", + "SubagentStart", + "SubagentStop", + "Stop", +]; + +const CLAUDE_HOOK_EVENTS: &[&str] = &[ + "ConfigChange", + "CwdChanged", + "Elicitation", + "ElicitationResult", + "FileChanged", + "InstructionsLoaded", + "MessageDisplay", + "Notification", + "PermissionDenied", + "PermissionRequest", + "PostCompact", + "PostToolBatch", + "PostToolUse", + "PostToolUseFailure", + "PreCompact", + "PreToolUse", + "SessionEnd", + "SessionStart", + "Setup", + "Stop", + "StopFailure", + "SubagentStart", + "SubagentStop", + "TaskCompleted", + "TaskCreated", + "TeammateIdle", + "UserPromptExpansion", + "UserPromptSubmit", + "WorktreeCreate", + "WorktreeRemove", +]; /// How the shim (re)launches the daemon: `bt trace daemon` from this same /// binary. @@ -140,76 +185,194 @@ fn run_command(program: &str, args: &[&str]) -> anyhow::Result<()> { Ok(()) } -fn codex_marketplace_installed(value: &Value) -> bool { +fn codex_marketplace_installed(value: &Value, name: &str) -> bool { value .get("marketplaces") .and_then(Value::as_array) .is_some_and(|items| { items .iter() - .any(|item| item.get("name").and_then(Value::as_str) == Some(CODEX_MARKETPLACE)) + .any(|item| item.get("name").and_then(Value::as_str) == Some(name)) }) } -fn codex_plugin_installed(value: &Value) -> bool { +fn codex_plugin_installed(value: &Value, id: &str) -> bool { value .get("installed") .and_then(Value::as_array) .is_some_and(|items| { items .iter() - .any(|item| item.get("pluginId").and_then(Value::as_str) == Some(CODEX_PLUGIN)) + .any(|item| item.get("pluginId").and_then(Value::as_str) == Some(id)) }) } -fn claude_marketplace_installed(value: &Value) -> bool { +fn claude_marketplace_installed(value: &Value, name: &str) -> bool { value.as_array().is_some_and(|items| { items .iter() - .any(|item| item.get("name").and_then(Value::as_str) == Some(CLAUDE_MARKETPLACE)) + .any(|item| item.get("name").and_then(Value::as_str) == Some(name)) }) } -fn claude_plugin(value: &Value) -> Option<&Value> { +fn claude_plugin<'a>(value: &'a Value, id: &str) -> Option<&'a Value> { value .as_array()? .iter() - .find(|item| item.get("id").and_then(Value::as_str) == Some(CLAUDE_PLUGIN)) + .find(|item| item.get("id").and_then(Value::as_str) == Some(id)) +} + +fn write_json(path: &Path, value: &Value) -> anyhow::Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("create plugin directory {}", parent.display()))?; + } + let mut encoded = serde_json::to_string_pretty(value)?; + encoded.push('\n'); + crate::utils::write_text_atomic(path, &encoded) +} + +fn hook_config(source: &str, events: &[&str], codex: bool) -> Value { + let command = format!( + "bt trace hook --source {source} --source-version {}", + crate::CLI_VERSION + ); + let hooks = events + .iter() + .map(|event| { + let hook = if codex { + json!({ + "type": "command", + "command": command, + "commandWindows": command, + "statusMessage": "Braintrust tracing" + }) + } else { + json!({ + "type": "command", + "command": command, + "async": false + }) + }; + ((*event).to_owned(), json!([{ "hooks": [hook] }])) + }) + .collect::>(); + json!({ "hooks": hooks }) +} + +fn generate_codex_plugin() -> anyhow::Result { + let root = crate::config::global_config_dir()?.join("trace-plugins/codex"); + write_json( + &root.join(".agents/plugins/marketplace.json"), + &json!({ + "name": GENERATED_MARKETPLACE, + "plugins": [{ + "name": "trace-codex", + "source": { + "source": "local", + "path": "./plugins/trace-codex" + } + }] + }), + )?; + let plugin = root.join("plugins/trace-codex"); + write_json( + &plugin.join(".codex-plugin/plugin.json"), + &json!({ + "name": "trace-codex", + "version": crate::CLI_VERSION, + "description": "Trace Codex sessions through the Braintrust bt daemon", + "hooks": "./hooks/hooks.json" + }), + )?; + write_json( + &plugin.join("hooks/hooks.json"), + &hook_config("codex", CODEX_HOOK_EVENTS, true), + )?; + Ok(root) +} + +fn generate_claude_plugin() -> anyhow::Result { + let root = crate::config::global_config_dir()?.join("trace-plugins/claude"); + write_json( + &root.join(".claude-plugin/marketplace.json"), + &json!({ + "$schema": "https://anthropic.com/claude-code/marketplace.schema.json", + "name": GENERATED_MARKETPLACE, + "version": crate::CLI_VERSION, + "description": "Braintrust tracing through the bt daemon", + "owner": { "name": "Braintrust" }, + "plugins": [{ + "name": "trace-claude-code", + "description": "Trace Claude Code sessions through the Braintrust bt daemon", + "source": "./plugins/trace-claude-code", + "category": "observability" + }] + }), + )?; + let plugin = root.join("plugins/trace-claude-code"); + write_json( + &plugin.join(".claude-plugin/plugin.json"), + &json!({ + "name": "trace-claude-code", + "version": crate::CLI_VERSION, + "description": "Trace Claude Code sessions through the Braintrust bt daemon" + }), + )?; + write_json( + &plugin.join("hooks/hooks.json"), + &hook_config("claude-code", CLAUDE_HOOK_EVENTS, false), + )?; + Ok(root) } fn setup_codex() -> anyhow::Result<()> { + let marketplace = generate_codex_plugin()?; + let marketplace = marketplace + .to_str() + .context("generated Codex plugin path is not UTF-8")?; let marketplaces = command_json("codex", &["plugin", "marketplace", "list", "--json"])?; - if !codex_marketplace_installed(&marketplaces) { + + let plugins = command_json("codex", &["plugin", "list", "--json"])?; + for plugin in [LEGACY_CODEX_PLUGIN, GENERATED_CODEX_PLUGIN] { + if codex_plugin_installed(&plugins, plugin) { + run_command("codex", &["plugin", "remove", plugin])?; + } + } + if codex_marketplace_installed(&marketplaces, GENERATED_MARKETPLACE) { run_command( "codex", - &["plugin", "marketplace", "add", CODEX_MARKETPLACE_SOURCE], + &["plugin", "marketplace", "remove", GENERATED_MARKETPLACE], )?; } - - let plugins = command_json("codex", &["plugin", "list", "--json"])?; - if !codex_plugin_installed(&plugins) { - run_command("codex", &["plugin", "add", CODEX_PLUGIN])?; - } + run_command("codex", &["plugin", "marketplace", "add", marketplace])?; + run_command("codex", &["plugin", "add", GENERATED_CODEX_PLUGIN])?; Ok(()) } fn setup_claude() -> anyhow::Result<()> { + let marketplace = generate_claude_plugin()?; + let marketplace = marketplace + .to_str() + .context("generated Claude plugin path is not UTF-8")?; let marketplaces = command_json("claude", &["plugin", "marketplace", "list", "--json"])?; - if !claude_marketplace_installed(&marketplaces) { + let plugins = command_json("claude", &["plugin", "list", "--json"])?; + if claude_plugin(&plugins, LEGACY_CLAUDE_PLUGIN) + .is_some_and(|plugin| plugin.get("enabled").and_then(Value::as_bool) != Some(false)) + { + run_command("claude", &["plugin", "disable", LEGACY_CLAUDE_PLUGIN])?; + } + if claude_plugin(&plugins, GENERATED_CLAUDE_PLUGIN).is_some() { + run_command("claude", &["plugin", "uninstall", GENERATED_CLAUDE_PLUGIN])?; + } + if claude_marketplace_installed(&marketplaces, GENERATED_MARKETPLACE) { run_command( "claude", - &["plugin", "marketplace", "add", CLAUDE_MARKETPLACE_SOURCE], + &["plugin", "marketplace", "remove", GENERATED_MARKETPLACE], )?; } - - let plugins = command_json("claude", &["plugin", "list", "--json"])?; - match claude_plugin(&plugins) { - None => run_command("claude", &["plugin", "install", CLAUDE_PLUGIN])?, - Some(plugin) if plugin.get("enabled").and_then(Value::as_bool) == Some(false) => { - run_command("claude", &["plugin", "enable", CLAUDE_PLUGIN])?; - } - Some(_) => {} - } + run_command("claude", &["plugin", "marketplace", "add", marketplace])?; + run_command("claude", &["plugin", "install", GENERATED_CLAUDE_PLUGIN])?; Ok(()) } diff --git a/tests/cli.rs b/tests/cli.rs index aa13a8d..9683dc8 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -297,6 +297,7 @@ fn trace_setup_codex_installs_plugin_and_preserves_existing_settings() { bt_command() .env("HOME", home.path()) + .env("XDG_CONFIG_HOME", state_dir.path()) .env("PATH", bin_dir.path()) .env("AGENT_SETUP_LOG", &log) .env("BT_DAEMON_CONFIG", &config) @@ -308,8 +309,25 @@ fn trace_setup_codex_installs_plugin_and_preserves_existing_settings() { )); let calls = fs::read_to_string(log).expect("read fake CLI calls"); - assert!(calls.contains("plugin marketplace add braintrustdata/braintrust-codex-plugin")); - assert!(calls.contains("plugin add trace-codex@braintrust-codex-plugins")); + let marketplace = state_dir.path().join("bt/trace-plugins/codex"); + assert!(calls.contains(&format!("plugin marketplace add {}", marketplace.display()))); + assert!(calls.contains("plugin add trace-codex@braintrust-bt-trace")); + assert!(!calls.contains("braintrustdata/braintrust-codex-plugin")); + + let hooks: serde_json::Value = serde_json::from_slice( + &fs::read(marketplace.join("plugins/trace-codex/hooks/hooks.json")) + .expect("read Codex hooks"), + ) + .expect("parse Codex hooks"); + let codex_command = hooks["hooks"]["SessionStart"][0]["hooks"][0]["command"] + .as_str() + .expect("Codex hook command"); + assert!(codex_command.starts_with("bt trace hook --source codex --source-version ")); + assert!(!codex_command.ends_with("--source-version ")); + assert_eq!( + hooks["hooks"]["Stop"][0]["hooks"][0]["commandWindows"], + codex_command + ); let settings: serde_json::Value = serde_json::from_slice(&fs::read(config).expect("read config")).expect("parse config"); @@ -334,6 +352,7 @@ fn trace_setup_claude_installs_plugin_and_creates_default_settings() { bt_command() .env("HOME", home.path()) + .env("XDG_CONFIG_HOME", state_dir.path()) .env("PATH", bin_dir.path()) .env("AGENT_SETUP_LOG", &log) .env("BT_DAEMON_CONFIG", &config) @@ -345,8 +364,22 @@ fn trace_setup_claude_installs_plugin_and_creates_default_settings() { )); let calls = fs::read_to_string(log).expect("read fake CLI calls"); - assert!(calls.contains("plugin marketplace add braintrustdata/braintrust-claude-plugin")); - assert!(calls.contains("plugin install trace-claude-code@braintrust-claude-plugin")); + let marketplace = state_dir.path().join("bt/trace-plugins/claude"); + assert!(calls.contains(&format!("plugin marketplace add {}", marketplace.display()))); + assert!(calls.contains("plugin install trace-claude-code@braintrust-bt-trace")); + assert!(!calls.contains("braintrustdata/braintrust-claude-plugin")); + + let hooks: serde_json::Value = serde_json::from_slice( + &fs::read(marketplace.join("plugins/trace-claude-code/hooks/hooks.json")) + .expect("read Claude hooks"), + ) + .expect("parse Claude hooks"); + let claude_command = hooks["hooks"]["SessionStart"][0]["hooks"][0]["command"] + .as_str() + .expect("Claude hook command"); + assert!(claude_command.starts_with("bt trace hook --source claude-code --source-version ")); + assert!(!claude_command.ends_with("--source-version ")); + assert_eq!(hooks["hooks"]["SessionEnd"][0]["hooks"][0]["async"], false); let settings: serde_json::Value = serde_json::from_slice(&fs::read(config).expect("read config")).expect("parse config"); @@ -356,19 +389,20 @@ fn trace_setup_claude_installs_plugin_and_creates_default_settings() { #[cfg(unix)] #[test] -fn trace_setup_claude_enables_an_existing_disabled_plugin() { +fn trace_setup_claude_replaces_generated_plugin_and_disables_legacy_tracing() { let home = tempfile::tempdir().expect("home tempdir"); let bin_dir = tempfile::tempdir().expect("bin tempdir"); let state_dir = tempfile::tempdir().expect("state tempdir"); let log = state_dir.path().join("claude.log"); write_agent_cli( &bin_dir.path().join("claude"), - r#"[{"name":"braintrust-claude-plugin"}]"#, - r#"[{"id":"trace-claude-code@braintrust-claude-plugin","enabled":false}]"#, + r#"[{"name":"braintrust-claude-plugin"},{"name":"braintrust-bt-trace"}]"#, + r#"[{"id":"trace-claude-code@braintrust-claude-plugin","enabled":true},{"id":"trace-claude-code@braintrust-bt-trace","enabled":true}]"#, ); bt_command() .env("HOME", home.path()) + .env("XDG_CONFIG_HOME", state_dir.path()) .env("PATH", bin_dir.path()) .env("AGENT_SETUP_LOG", &log) .args(["trace", "setup", "claude"]) @@ -376,9 +410,41 @@ fn trace_setup_claude_enables_an_existing_disabled_plugin() { .success(); let calls = fs::read_to_string(log).expect("read fake CLI calls"); - assert!(calls.contains("plugin enable trace-claude-code@braintrust-claude-plugin")); - assert!(!calls.contains("plugin marketplace add")); - assert!(!calls.contains("plugin install")); + assert!(calls.contains("plugin disable trace-claude-code@braintrust-claude-plugin")); + assert!(calls.contains("plugin uninstall trace-claude-code@braintrust-bt-trace")); + assert!(calls.contains("plugin marketplace remove braintrust-bt-trace")); + assert!(calls.contains("plugin marketplace add")); + assert!(calls.contains("plugin install trace-claude-code@braintrust-bt-trace")); +} + +#[cfg(unix)] +#[test] +fn trace_setup_codex_replaces_generated_plugin_and_removes_legacy_tracing() { + let home = tempfile::tempdir().expect("home tempdir"); + let bin_dir = tempfile::tempdir().expect("bin tempdir"); + let state_dir = tempfile::tempdir().expect("state tempdir"); + let log = state_dir.path().join("codex.log"); + write_agent_cli( + &bin_dir.path().join("codex"), + r#"{"marketplaces":[{"name":"braintrust-codex-plugins"},{"name":"braintrust-bt-trace"}]}"#, + r#"{"installed":[{"pluginId":"trace-codex@braintrust-codex-plugins"},{"pluginId":"trace-codex@braintrust-bt-trace"}]}"#, + ); + + bt_command() + .env("HOME", home.path()) + .env("XDG_CONFIG_HOME", state_dir.path()) + .env("PATH", bin_dir.path()) + .env("AGENT_SETUP_LOG", &log) + .args(["trace", "setup", "codex"]) + .assert() + .success(); + + let calls = fs::read_to_string(log).expect("read fake CLI calls"); + assert!(calls.contains("plugin remove trace-codex@braintrust-codex-plugins")); + assert!(calls.contains("plugin remove trace-codex@braintrust-bt-trace")); + assert!(calls.contains("plugin marketplace remove braintrust-bt-trace")); + assert!(calls.contains("plugin marketplace add")); + assert!(calls.contains("plugin add trace-codex@braintrust-bt-trace")); } #[test]