diff --git a/bt-daemon/src/codex_hooks.rs b/bt-daemon/src/codex_hooks.rs new file mode 100644 index 00000000..b3a7fd41 --- /dev/null +++ b/bt-daemon/src/codex_hooks.rs @@ -0,0 +1,419 @@ +//! Inspect and activate Braintrust hooks through Codex's app-server protocol. +//! Trust writes use the same exact-definition config API as Codex's /hooks UI. + +use anyhow::{anyhow, Context}; +use serde_json::{json, Value}; +use std::path::Path; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::{ + AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, +}; + +const PLUGIN: &str = "trace-codex@braintrust-codex-plugins"; +const PACKAGED_HOOKS: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json" +)); +const TIMEOUT: Duration = Duration::from_secs(5); +const PLUGIN_EVENTS: &[&str] = &[ + "sessionStart", + "userPromptSubmit", + "preToolUse", + "permissionRequest", + "postToolUse", + "preCompact", + "postCompact", + "subagentStart", + "subagentStop", + "stop", + "interrupt", + "sessionEnd", +]; +pub(crate) async fn activate() -> anyhow::Result<()> { + let cwd = std::env::current_dir()?; + let mut command: tokio::process::Command = + crate::subprocess::background_command("codex").into(); + command + .args(["--no-daemon", "app-server", "--stdio"]) + .current_dir(&cwd) + .kill_on_drop(true) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()); + let mut child = command + .spawn() + .context("start Codex to activate tracing hooks")?; + let result = + tokio::time::timeout(TIMEOUT, async { + let mut stdin = child.stdin.take().context("Codex stdin unavailable")?; + let stdout = child.stdout.take().context("Codex stdout unavailable")?; + let mut stdout = BufReader::new(stdout.take(8 * 1024 * 1024)); + request(&mut stdin, &mut stdout, 1, "initialize", json!({ + "clientInfo":{"name":"braintrust_trace_setup","version":env!("CARGO_PKG_VERSION")}, + "capabilities":{"experimentalApi":true} + })).await?; + stdin.write_all(b"{\"method\":\"initialized\"}\n").await?; + let config = request( + &mut stdin, + &mut stdout, + 2, + "config/read", + json!({"includeLayers":false}), + ) + .await?; + let features = &config["config"]["features"]; + anyhow::ensure!( + features["hooks"] != false && features["codex_hooks"] != false, + "Codex hooks are disabled in config; enable hooks and retry setup" + ); + let plugins = request( + &mut stdin, + &mut stdout, + 3, + "plugin/list", + json!({ + "cwds":[cwd], "forceRefetch":false, "marketplaceKinds":["local"] + }), + ) + .await?; + let installed = plugins["marketplaces"] + .as_array() + .into_iter() + .flatten() + .filter_map(|marketplace| marketplace["plugins"].as_array()) + .flatten() + .any(|plugin| { + plugin["id"] == PLUGIN + && plugin["installed"] == true + && plugin["enabled"] == true + }); + anyhow::ensure!( + installed, + "Codex tracing plugin is not installed and enabled; rerun `bt trace enable codex`" + ); + let hooks = request( + &mut stdin, + &mut stdout, + 4, + "hooks/list", + json!({"cwds":[cwd]}), + ) + .await?; + let edits = trust_edits(&hooks)?; + if !edits.is_empty() { + request( + &mut stdin, + &mut stdout, + 5, + "config/batchWrite", + json!({ + "edits":[{"keyPath":"hooks.state", "value":edits,"mergeStrategy":"upsert"}], + "reloadUserConfig":true + }), + ) + .await?; + } + let hooks = request( + &mut stdin, + &mut stdout, + 6, + "hooks/list", + json!({"cwds":[cwd]}), + ) + .await?; + anyhow::ensure!( + trust_edits(&hooks)?.is_empty(), + "Codex did not retain Braintrust hook trust after setup" + ); + Ok(()) + }) + .await; + let _ = child.kill().await; + match result { + Ok(result) => result.context("Codex plugin installed, but hook activation failed"), + Err(_) => Err(anyhow!("Codex hook activation timed out after 5 seconds")), + } +} + +/// Build a command-line-scoped hooks.state value for Codex managed runs. The +/// keys and hashes come from Codex after it has loaded our exact hook overrides. +pub(crate) async fn run_trust_override( + program: &std::ffi::OsStr, + cwd: &Path, + hook_args: &[std::ffi::OsString], + expected_command: &str, +) -> anyhow::Result { + let mut command: tokio::process::Command = + crate::subprocess::background_command(program).into(); + command + .arg("--no-daemon") + .args(hook_args) + .args(["app-server", "--stdio"]) + .current_dir(cwd) + .kill_on_drop(true) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()); + let mut child = command + .spawn() + .context("start Codex to prepare managed hook trust")?; + let result = + tokio::time::timeout(TIMEOUT, async { + let mut stdin = child.stdin.take().context("Codex stdin unavailable")?; + let stdout = child.stdout.take().context("Codex stdout unavailable")?; + let mut stdout = BufReader::new(stdout.take(8 * 1024 * 1024)); + request(&mut stdin, &mut stdout, 1, "initialize", json!({ + "clientInfo":{"name":"braintrust_trace_run","version":env!("CARGO_PKG_VERSION")}, + "capabilities":{"experimentalApi":true} + })).await?; + stdin.write_all(b"{\"method\":\"initialized\"}\n").await?; + let hooks = request( + &mut stdin, + &mut stdout, + 2, + "hooks/list", + json!({"cwds":[cwd]}), + ) + .await?; + managed_run_trust_state(&hooks, expected_command) + }) + .await; + let _ = child.kill().await; + match result { + Ok(result) => result, + Err(_) => Err(anyhow!( + "Codex managed hook trust inspection timed out after 5 seconds" + )), + } +} + +fn managed_run_trust_state( + hooks: &Value, + expected_command: &str, +) -> anyhow::Result { + let expected: Value = serde_json::from_str(PACKAGED_HOOKS)?; + let expected_events: std::collections::HashSet = expected["hooks"] + .as_object() + .context("packaged hooks missing")? + .keys() + .map(|event| format!("{}{}", event[..1].to_ascii_lowercase(), &event[1..])) + .collect(); + let entries = hooks["data"] + .as_array() + .context("hooks/list has no data array")?; + let mut seen = std::collections::HashSet::new(); + let mut trust = std::collections::BTreeMap::new(); + for entry in entries { + anyhow::ensure!( + entry["errors"].as_array().is_some_and(Vec::is_empty), + "Codex reported hook errors during managed-run setup" + ); + for hook in entry["hooks"] + .as_array() + .context("hooks/list entry has no hooks array")? + { + if hook["source"] != "sessionFlags" { + continue; + } + let Some(event) = hook["eventName"].as_str() else { + continue; + }; + if !expected_events.contains(event) { + continue; + } + anyhow::ensure!( + seen.insert(event.to_owned()) + && hook["handlerType"] == "command" + && hook["command"] == expected_command + && hook["enabled"] == true + && hook["isManaged"] == false, + "Codex changed a Braintrust hook while preparing this run" + ); + anyhow::ensure!( + matches!( + hook["trustStatus"].as_str(), + Some("trusted" | "untrusted" | "modified") + ), + "Codex reported an unsupported trust state for a managed Braintrust hook" + ); + let key = hook["key"] + .as_str() + .filter(|value| !value.is_empty()) + .context("Codex omitted a managed Braintrust hook key")?; + let hash = hook["currentHash"] + .as_str() + .filter(|value| !value.is_empty()) + .context("Codex omitted a managed Braintrust hook hash")?; + anyhow::ensure!( + trust.insert(key.to_owned(), hash.to_owned()).is_none(), + "duplicate managed Braintrust hook key" + ); + } + } + anyhow::ensure!( + seen == expected_events, + "Codex did not load the complete Braintrust hook set for this run" + ); + let fields = trust + .iter() + .map(|(key, hash)| { + Ok(format!( + "{}={{trusted_hash={}}}", + serde_json::to_string(key)?, + serde_json::to_string(hash)? + )) + }) + .collect::>>()? + .join(","); + Ok(std::ffi::OsString::from(format!( + "hooks.state={{{fields}}}" + ))) +} + +/// Validate the complete installed file, including the other platform's command, +/// before approving exact hashes returned by Codex. Never infer hashes or keys. +fn trust_edits(hooks: &Value) -> anyhow::Result> { + let expected: Value = serde_json::from_str(PACKAGED_HOOKS)?; + let mut checked_paths = std::collections::HashSet::new(); + let mut events = std::collections::HashSet::new(); + let mut edits = serde_json::Map::new(); + for entry in hooks["data"] + .as_array() + .context("hooks/list has no data array")? + { + for hook in entry["hooks"] + .as_array() + .context("hooks/list entry has no hooks array")? + { + if hook["pluginId"].as_str() != Some(PLUGIN) { + continue; + } + let event = hook["eventName"].as_str().context("hook event missing")?; + anyhow::ensure!( + PLUGIN_EVENTS.contains(&event) + && events.insert(event) + && hook["handlerType"] == "command" + && hook["command"] == "bt trace hook --source codex" + && hook["enabled"] == true + && hook["isManaged"] == false + && hook["source"] == "plugin", + "Unexpected or disabled Braintrust hook definition; inspect /hooks" + ); + let path = hook["sourcePath"] + .as_str() + .context("hook source path missing")?; + anyhow::ensure!( + Path::new(path).is_absolute(), + "hook source path is not absolute" + ); + if checked_paths.insert(path) { + let installed: Value = serde_json::from_slice(&std::fs::read(path)?)?; + anyhow::ensure!(installed == expected, + "Installed Braintrust hooks differ from this bt version; update bt and retry setup"); + } + if hook["trustStatus"] == "trusted" { + continue; + } + anyhow::ensure!( + hook["trustStatus"] == "untrusted" || hook["trustStatus"] == "modified", + "Unknown Braintrust hook trust state" + ); + let key = hook["key"] + .as_str() + .filter(|s| !s.is_empty()) + .context("hook key missing")?; + let hash = hook["currentHash"] + .as_str() + .filter(|s| !s.is_empty()) + .context("hook hash missing")?; + anyhow::ensure!( + edits + .insert(key.to_owned(), json!({"trusted_hash":hash})) + .is_none(), + "Duplicate Braintrust hook key" + ); + } + } + anyhow::ensure!( + events.len() == PLUGIN_EVENTS.len(), + "Braintrust lifecycle hooks are missing" + ); + Ok(edits) +} + +async fn request( + stdin: &mut (impl AsyncWrite + Unpin), + stdout: &mut (impl AsyncBufRead + Unpin), + id: u64, + method: &str, + params: Value, +) -> anyhow::Result { + let mut encoded = serde_json::to_vec(&json!({"id":id,"method":method,"params":params}))?; + encoded.push(b'\n'); + stdin.write_all(&encoded).await?; + stdin.flush().await?; + loop { + let mut line = String::new(); + anyhow::ensure!( + stdout.read_line(&mut line).await? != 0, + "Codex closed the probe stream" + ); + let response: Value = serde_json::from_str(&line).context("invalid Codex response")?; + if response.get("id") != Some(&json!(id)) { + continue; + } + if let Some(error) = response.get("error") { + return Err(anyhow!( + "{method}: {}", + error + .get("message") + .and_then(Value::as_str) + .unwrap_or("request failed") + )); + } + return response + .get("result") + .cloned() + .context("Codex response has no result"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn required_events_match_the_packaged_plugin() { + let hooks: Value = serde_json::from_str(PACKAGED_HOOKS).unwrap(); + let events: std::collections::BTreeSet<_> = hooks["hooks"] + .as_object() + .unwrap() + .keys() + .map(|key| format!("{}{}", key[..1].to_ascii_lowercase(), &key[1..])) + .collect(); + assert_eq!( + events, + PLUGIN_EVENTS + .iter() + .map(|event| (*event).to_owned()) + .collect() + ); + } + + #[tokio::test] + async fn rpc_ignores_notifications_and_reports_unsupported_methods() { + let mut input = Vec::new(); + let mut output = std::io::Cursor::new( + b"{\"method\":\"notice\"}\n{\"id\":3,\"error\":{\"message\":\"Method not found\"}}\n", + ); + let error = request(&mut input, &mut output, 3, "hooks/list", json!({})) + .await + .unwrap_err(); + assert!(error.to_string().contains("Method not found")); + assert_eq!( + serde_json::from_slice::(&input).unwrap()["method"], + "hooks/list" + ); + } +} diff --git a/bt-daemon/src/command_output.rs b/bt-daemon/src/command_output.rs index 704e7f09..00b57da1 100644 --- a/bt-daemon/src/command_output.rs +++ b/bt-daemon/src/command_output.rs @@ -186,8 +186,9 @@ impl TraceCommandOutput { display_name: impl Into, settings_path: impl Into, ) -> Self { + let source = source.into(); Self::Enable(SetupCommandOutput { - source: source.into(), + source, display_name: display_name.into(), settings_path: settings_path.into(), restart_required: true, @@ -216,8 +217,9 @@ impl TraceCommandOutput { } pub fn update(source: impl Into, display_name: impl Into) -> Self { + let source = source.into(); Self::Update(UpdateCommandOutput { - source: source.into(), + source, display_name: display_name.into(), restart_required: true, }) @@ -427,6 +429,26 @@ mod tests { .contains("removed for Google Antigravity")); } + #[test] + fn codex_setup_and_updates_do_not_request_manual_trust() { + for output in [ + TraceCommandOutput::setup("codex", "Codex", PathBuf::from("braintrust.json")), + TraceCommandOutput::update("codex", "Codex"), + ] { + let human = output.render(OutputFormat::Human).unwrap(); + assert!(!human.contains("Open /hooks")); + let json: serde_json::Value = + serde_json::from_str(&output.render(OutputFormat::Json).unwrap()).unwrap(); + assert!(json.get("next_steps").is_none()); + } + let disabled = + TraceCommandOutput::disable("codex", "Codex", PathBuf::from("braintrust.json")); + assert!(!disabled + .render(OutputFormat::Human) + .unwrap() + .contains("Open /hooks")); + } + #[test] fn stop_json_reports_idempotent_and_successful_shutdowns() { let absent: serde_json::Value = serde_json::from_str( diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 61a3f58f..a532d1b1 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -15,6 +15,7 @@ pub mod paths; mod agents; mod args; mod client; +mod codex_hooks; mod command_output; mod correlation; mod delivery_ledger; diff --git a/bt-daemon/src/managed_run/codex.rs b/bt-daemon/src/managed_run/codex.rs index aeadca29..8c0dab1a 100644 --- a/bt-daemon/src/managed_run/codex.rs +++ b/bt-daemon/src/managed_run/codex.rs @@ -6,36 +6,80 @@ use super::{managed_hook_shell_command, Injection, ManagedRun, Terminal}; use crate::agents::Codex; use crate::args::RunHookCommand; -pub(super) const CODEX_RUN_HOOK_EVENTS: &[&str] = &[ - "SessionStart", - "UserPromptSubmit", - "PreToolUse", - "PermissionRequest", - "PostToolUse", - "PreCompact", - "PostCompact", - "SubagentStart", - "SubagentStop", - "Stop", - "SessionEnd", -]; +pub(super) const PACKAGED_HOOKS: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json" +)); + +pub(super) fn managed_commands(hook_command: &RunHookCommand) -> anyhow::Result<(String, String)> { + Ok(( + managed_hook_shell_command(hook_command, "codex", false)?, + managed_hook_shell_command(hook_command, "codex", true)?, + )) +} fn codex_managed_run_args( unix_command: &str, windows_command: &str, ) -> anyhow::Result> { - let unix_command = serde_json::to_string(unix_command)?; - let windows_command = serde_json::to_string(windows_command)?; + let manifest: serde_json::Value = serde_json::from_str(PACKAGED_HOOKS)?; let mut args = vec![OsString::from("--enable"), OsString::from("hooks")]; - for event in CODEX_RUN_HOOK_EVENTS { + for (event, groups) in manifest["hooks"] + .as_object() + .ok_or_else(|| anyhow::anyhow!("Codex hook manifest has no hooks map"))? + { + let mut groups = groups.clone(); + for group in groups + .as_array_mut() + .ok_or_else(|| anyhow::anyhow!("Codex hook groups must be an array"))? + { + for hook in group["hooks"] + .as_array_mut() + .ok_or_else(|| anyhow::anyhow!("Codex hook group has no hooks array"))? + { + anyhow::ensure!(hook["type"] == "command", "Unsupported Codex hook type"); + hook["command"] = serde_json::Value::String(unix_command.to_owned()); + hook["commandWindows"] = serde_json::Value::String(windows_command.to_owned()); + } + } args.push(OsString::from("-c")); args.push(OsString::from(format!( - "hooks.{event}=[{{hooks=[{{type=\"command\",command={unix_command},commandWindows={windows_command}}}]}}]" + "hooks.{event}={}", + toml_literal(&groups)? ))); } Ok(args) } +fn toml_literal(value: &serde_json::Value) -> anyhow::Result { + Ok(match value { + serde_json::Value::Null => anyhow::bail!("null is unsupported in Codex hook config"), + serde_json::Value::Bool(value) => value.to_string(), + serde_json::Value::Number(value) => value.to_string(), + serde_json::Value::String(value) => serde_json::to_string(value)?, + serde_json::Value::Array(values) => format!( + "[{}]", + values + .iter() + .map(toml_literal) + .collect::>>()? + .join(",") + ), + serde_json::Value::Object(values) => format!( + "{{{}}}", + values + .iter() + .map(|(key, value)| Ok(format!( + "{}={}", + serde_json::to_string(key)?, + toml_literal(value)? + ))) + .collect::>>()? + .join(",") + ), + }) +} + impl ManagedRun for Codex { fn executable(&self) -> (&'static str, &'static str) { ("CODEX_BIN", "codex") @@ -59,8 +103,7 @@ impl ManagedRun for Codex { hook_command: &RunHookCommand, _managed_run_id: &str, ) -> anyhow::Result { - let unix_command = managed_hook_shell_command(hook_command, "codex", false)?; - let windows_command = managed_hook_shell_command(hook_command, "codex", true)?; + let (unix_command, windows_command) = managed_commands(hook_command)?; Ok(Injection { args: codex_managed_run_args(&unix_command, &windows_command)?, ..Injection::default() diff --git a/bt-daemon/src/managed_run/mod.rs b/bt-daemon/src/managed_run/mod.rs index 1f3716f4..f51b9703 100644 --- a/bt-daemon/src/managed_run/mod.rs +++ b/bt-daemon/src/managed_run/mod.rs @@ -12,7 +12,7 @@ use std::io::IsTerminal; use std::path::PathBuf; use crate::agents::{registrar, Agent}; -use crate::args::{RunArgs, RunHookCommand}; +use crate::args::{RunArgs, RunHookCommand, RunSource}; use crate::client::{flush_managed_run_in, shutdown_daemon}; use crate::hook::MANAGED_RUN_ID_ENV; use crate::route::resolve_span_plugin_paths; @@ -88,7 +88,28 @@ pub async fn run_traced( let executable = std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable)); let managed_run_id = uuid::Uuid::new_v4().to_string(); - let injection = agent.inject(&hook_command, &managed_run_id)?; + let mut injection = agent.inject(&hook_command, &managed_run_id)?; + if args.source == RunSource::Codex { + let executable_for_trust = + std::env::var_os("CODEX_BIN").unwrap_or_else(|| OsString::from("codex")); + let cwd = std::env::current_dir()?; + let (unix_command, windows_command) = + crate::managed_run::codex::managed_commands(&hook_command)?; + let expected_command = if cfg!(windows) { + windows_command + } else { + unix_command + }; + let trust_override = crate::codex_hooks::run_trust_override( + &executable_for_trust, + &cwd, + &injection.args, + &expected_command, + ) + .await?; + 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 @@ -231,7 +252,6 @@ pub(super) fn quote_windows_command_arg(arg: &str) -> String { #[cfg(test)] mod tests { - use super::codex::CODEX_RUN_HOOK_EVENTS; use super::*; use crate::agents::{Codex, Pi}; use crate::args::RunSource; @@ -328,10 +348,21 @@ mod tests { assert!(!args .iter() .any(|arg| arg == "--dangerously-bypass-hook-trust")); - assert_eq!( - args.iter().filter(|arg| *arg == "-c").count(), - CODEX_RUN_HOOK_EVENTS.len() - ); + let hook_overrides: Vec<_> = args + .iter() + .filter_map(|arg| arg.to_str()) + .filter(|arg| arg.starts_with("hooks.")) + .collect(); + let manifest: serde_json::Value = + serde_json::from_str(super::codex::PACKAGED_HOOKS).unwrap(); + let expected_events = manifest["hooks"].as_object().unwrap(); + assert_eq!(hook_overrides.len(), expected_events.len()); + assert!(hook_overrides + .iter() + .any(|arg| arg.starts_with("hooks.Interrupt="))); + assert!(hook_overrides + .iter() + .any(|arg| arg.starts_with("hooks.SessionEnd="))); let config = args .iter() .find_map(|arg| { diff --git a/bt-daemon/src/setup/codex.rs b/bt-daemon/src/setup/codex.rs index 1dcfe1df..f963b6f1 100644 --- a/bt-daemon/src/setup/codex.rs +++ b/bt-daemon/src/setup/codex.rs @@ -5,7 +5,6 @@ use super::common::{ }; use crate::agents::Codex; use crate::setup::Setup; -use anyhow::bail; use serde_json::Value; const MARKETPLACE_NAME: &str = "braintrust-codex-plugins"; @@ -63,7 +62,9 @@ impl Setup for Codex { fn update(&self, runner: &mut dyn CommandRunner) -> anyhow::Result<()> { if installed_plugin(&list_plugins(runner)?).is_none() { - bail!("Codex tracing plugin is not installed; run `bt trace enable codex`"); + // Setup owns the desired state: repair a removed plugin the same + // way enable does, then let the caller reapply and verify trust. + return self.enable(runner); } MARKETPLACE.refresh_published(runner, "Codex", "codex")?; runner.run("codex", &["plugin", "add", PLUGIN]) @@ -138,14 +139,18 @@ mod tests { } #[test] - fn update_refuses_to_install_a_missing_plugin() { - let mut runner = FakeRunner::new([serde_json::json!({"installed": []})]); + fn update_repairs_a_missing_plugin() { + let mut runner = FakeRunner::new([ + serde_json::json!({"installed": []}), + serde_json::json!({"marketplaces": []}), + ]); - assert!(Codex.update(&mut runner).is_err()); - assert!(!runner - .calls - .iter() - .any(|call| call.starts_with("codex plugin add "))); + Codex.update(&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] diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 684c97b6..a3cc1c63 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -6,7 +6,7 @@ //! the coding-agent integrations. use crate::route::{apply_additional_metadata, apply_tags, resolve_span_plugin_paths}; -use crate::trace_command::{DoctorAgent, DoctorArgs, TraceCommand}; +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, @@ -495,12 +495,14 @@ fn daemon_warnings( async fn doctor_output(host: &TraceHostContext, args: DoctorArgs) -> DoctorCommandOutput { let source = args.agent.identity().id; + let settings = crate::settings::AgentSettings::load(source); + let socket = paths::socket_path(None); doctor_output_at( host, args, - &paths::socket_path(None), + &socket, paths::agent_settings_path(source, None), - crate::settings::AgentSettings::load(source), + settings, ) .await } @@ -653,13 +655,23 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul if !enable_args.plugin.is_empty() { route.span_plugins = resolve_span_plugin_paths(&enable_args.plugin)?; } - print_output(run_enable(enable_args, route)?, host.output_format) + let codex = matches!(enable_args.agent, SetupAgent::Codex); + let output = run_enable(enable_args, route)?; + if codex { + crate::codex_hooks::activate().await?; + } + print_output(output, host.output_format) } TraceCommand::Disable(disable_args) => { print_output(run_disable(disable_args.agent)?, host.output_format) } TraceCommand::Update(update_args) => { - print_output(crate::run_update(update_args.agent)?, host.output_format) + let codex = matches!(update_args.agent, SetupAgent::Codex); + let output = crate::run_update(update_args.agent)?; + if codex { + crate::codex_hooks::activate().await?; + } + print_output(output, host.output_format) } TraceCommand::Daemon(serve_args) => { init_daemon_logging(host.verbose); @@ -764,7 +776,7 @@ async fn run_hook_command( #[cfg(test)] mod tests { use super::*; - use crate::trace_command::{SetupAgent, SetupArgs, StopArgs}; + use crate::trace_command::{SetupArgs, StopArgs}; use crate::wire::{BackendAuth, TraceDestination}; use crate::{ImportArgs, ImportSource, RunArgs, RunSource, StatusArgs, TraceCommand}; use std::path::PathBuf; diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 30d7751a..84937d8f 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -284,6 +284,60 @@ fn codex_happy_path_builds_session_turn_llm_tool_tree() { ); } +#[test] +fn interrupt_hook_closes_aborted_work_without_a_later_hook_or_flush() { + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("rollout.jsonl"); + let path = transcript.to_str().unwrap(); + for record in [ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"s","cwd":"/x/app"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"t1"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Working"}]}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"response_item","payload":{"type":"function_call","call_id":"c1","name":"shell","arguments":"{}","metadata":{"turn_id":"t1"}}}), + ] { + append(&transcript, record); + } + let mut translator = Registry::default_agents().create("codex", "s"); + let ctx = SessionCtx { + session_id: "s".into(), + config: None, + }; + let mut ops = translator + .handle(&envelope("s", "SessionStart", path, json!({})), &ctx) + .unwrap(); + let before = reduce(ops.clone()); + assert_eq!(find(&before, SpanType::Task, "turn: t1").end_ms, None); + assert_eq!(find(&before, SpanType::Tool, "shell").end_ms, None); + append( + &transcript, + json!({"timestamp":"2026-01-01T00:00:05Z","type":"event_msg","payload":{"type":"turn_aborted","turn_id":"t1","reason":"interrupted"}}), + ); + ops.extend( + translator + .handle( + &envelope("s", "Interrupt", path, json!({"turn_id":"t1"})), + &ctx, + ) + .unwrap(), + ); + let after = reduce(ops); + for row in [ + find(&after, SpanType::Task, "turn: t1"), + find(&after, SpanType::Tool, "shell"), + ] { + assert_eq!(row.end_ms, Some(1_767_225_605_000)); + assert_eq!(row.metadata.as_ref().unwrap()["outcome"], "interrupted"); + assert!(row.error.is_none()); + } + let turn = find(&after, SpanType::Task, "turn: t1"); + let llm = after + .values() + .find(|row| row.span_type == SpanType::Llm && row.parent_span_ids.contains(&turn.span_id)) + .unwrap(); + assert!(llm.end_ms.is_some()); + assert_eq!(llm.metadata.as_ref().unwrap()["outcome"], "interrupted"); +} + #[test] fn codex_turn_aborted_closes_open_work_as_interrupted() { let tmp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 8a38f31d..314cfd56 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -2714,6 +2714,28 @@ async fn managed_run_flushes_and_preserves_agent_exit_status() { std::fs::write( &agent, r#"#!/bin/sh +for argument in "$@"; do + if [ "$argument" = app-server ]; then + while IFS= read -r request; do + case "$request" in + *'"method":"initialize"'*) + printf '{"id":1,"result":{}}\n' + ;; + *'"method":"hooks/list"'*) + hooks= + index=0 + for event in sessionStart userPromptSubmit preToolUse permissionRequest postToolUse preCompact postCompact subagentStart subagentStop stop interrupt sessionEnd; do + index=$((index + 1)) + [ -z "$hooks" ] || hooks="$hooks," + hooks="$hooks{\"source\":\"sessionFlags\",\"eventName\":\"$event\",\"handlerType\":\"command\",\"command\":\"'unused-hook-command' '--source' 'codex' '--managed-run-hook'\",\"enabled\":true,\"isManaged\":false,\"trustStatus\":\"untrusted\",\"key\":\"hook-$index\",\"currentHash\":\"hash-$index\"}" + done + printf '{"id":2,"result":{"data":[{"errors":[],"hooks":[%s]}]}}\n' "$hooks" + ;; + esac + done + exit 0 + fi +done previous= last= for argument in "$@"; do diff --git a/src/plugins/codex/content/plugins/trace-codex/README.md b/src/plugins/codex/content/plugins/trace-codex/README.md index efc5553e..389c3333 100644 --- a/src/plugins/codex/content/plugins/trace-codex/README.md +++ b/src/plugins/codex/content/plugins/trace-codex/README.md @@ -21,9 +21,13 @@ Install and configure the published plugin with the Braintrust CLI: bt trace enable codex --project my-coding-agent ``` -Use `--profile` or `--org` when needed. Restart Codex after setup so it loads -the plugin. Codex will apply its normal hook-review flow; approve the stable -Braintrust hook definition through `/hooks` when prompted. +Use `--profile` or `--org` when needed. Setup activates Braintrust tracing in +Codex. Restart Codex after setup so it loads the plugin. `bt trace update codex` +refreshes the plugin and reapplies its setup. + +If setup reports an error, update Codex and `bt`, then retry. Use +`bt trace doctor codex` to inspect route, authentication, daemon, and plugin +diagnostics. To inspect tracing configuration and daemon status: diff --git a/src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json b/src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json index 768e5ca5..5359fe17 100644 --- a/src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json +++ b/src/plugins/codex/content/plugins/trace-codex/hooks/hooks.json @@ -119,6 +119,30 @@ } ] } + ], + "Interrupt": [ + { + "hooks": [ + { + "type": "command", + "command": "bt trace hook --source codex", + "commandWindows": "bt trace hook --source codex", + "statusMessage": "Braintrust tracing" + } + ] + } + ], + "SessionEnd": [ + { + "hooks": [ + { + "type": "command", + "command": "bt trace hook --source codex", + "commandWindows": "bt trace hook --source codex", + "statusMessage": "Braintrust tracing" + } + ] + } ] } }