diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index dfc514e..7b482aa 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -11,6 +11,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "android_system_properties" version = "0.1.5" @@ -271,6 +277,8 @@ dependencies = [ "clap", "regex", "reqwest", + "rquickjs", + "rquickjs-serde", "serde", "serde_json", "sha2", @@ -598,6 +606,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -756,6 +770,11 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] [[package]] name = "heck" @@ -1407,6 +1426,15 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "relative-path" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bca40a312222d8ba74837cb474edef44b37f561da5f773981007a10bbaa992b0" +dependencies = [ + "serde", +] + [[package]] name = "reqwest" version = "0.12.28" @@ -1463,6 +1491,45 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rquickjs" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e04e4eedfb060b503b5f0a2644abb890b0b3620d3fb674f9455f230014964e4" +dependencies = [ + "rquickjs-core", +] + +[[package]] +name = "rquickjs-core" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16e4f499ac5b943d97ee6dbc44f23c2c10426f420f7d2f1793d6318911b6608c" +dependencies = [ + "hashbrown", + "relative-path", + "rquickjs-sys", +] + +[[package]] +name = "rquickjs-serde" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04cf0aa631f8d0c5051db35f9f59899c34074d7b6be280c03fd4ce6165d0ed35" +dependencies = [ + "rquickjs", + "serde", +] + +[[package]] +name = "rquickjs-sys" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13ac243b86a74120814ef7e9e30ad5a2c1199b7b9963b1cf7c84e4cdc1cad99" +dependencies = [ + "cc", +] + [[package]] name = "rustc-hash" version = "2.1.3" diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index d404440..0718e0f 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -24,6 +24,8 @@ async-trait = "0.1" chrono = "0.4" clap = { version = "4", features = ["derive", "env"] } regex = "1" +rquickjs = "0.12.2" +rquickjs-serde = "0.6.1" serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.10" diff --git a/bt-daemon/README.md b/bt-daemon/README.md index e5feabb..d408288 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -59,6 +59,101 @@ the default `bt` profile. Credentials and backend URLs are never stored here; production resolves and refreshes them through `bt`. `bt trace run` supplies a process-local settings overlay and never changes any of these files. +### JavaScript span plugins + +`--plugin PATH` registers a synchronous ES module that transforms each +sink-neutral span row after translation and immediately before delivery. Repeat +the flag to compose plugins from left to right. `enable` persists its ordered list +for ordinary agent sessions. Managed runs and imports are isolated from that +list and use only the `--plugin` flags passed to their command. Each path is +canonicalized to an absolute path before it is validated or stored. + +Each module must default-export a synchronous function. It receives a span and +`{ operation, source, session_id, env }`, and must return a JSON-compatible span +object. Span, root, and parent identities cannot be changed: + +```js +// redact.mjs +function redact(value) { + if (typeof value === "string") { + return value.replace(/sk-[A-Za-z0-9_-]+/g, "[REDACTED]"); + } + if (Array.isArray(value)) return value.map(redact); + if (value && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([key, child]) => [key, redact(child)]), + ); + } + return value; +} + +export default function redactSpan(span) { + const next = { ...span }; + for (const field of ["input", "output", "error"]) { + if (field in next) next[field] = redact(next[field]); + } + return next; +} +``` + +The context can drive a second transform without changing the first one: + +```js +// tag-ci.mjs +export default function tagCi(span, context) { + if (!context.env.CI) return span; + + return { + ...span, + tags: [...new Set([...(span.tags ?? []), "ci"])], + metadata: { + ...(span.metadata ?? {}), + deployment: context.env.DEPLOYMENT_ENV ?? "unknown", + trace_source: context.source, + }, + }; +} +``` + +Register both transforms persistently for ordinary Codex sessions. The +redactor runs first and its returned span becomes the tagger's input: + +```bash +bt trace enable codex --plugin ./redact.mjs --plugin ./tag-ci.mjs +``` + +`run` and `import` plugins apply only to that command. They replace, rather than +merge with, plugins saved by `enable`: + +```bash +# Only local.mjs runs; redact.mjs and tag-ci.mjs remain global enable behavior. +bt trace run --plugin ./local.mjs codex -- "summarize this change" + +# Only sanitize-history.mjs transforms spans produced by this import. +bt trace import codex SESSION_ID --plugin ./sanitize-history.mjs +``` + +The journal stores raw input events, not transformed spans. After daemon +recovery, replayed events therefore pass through the resumed session's current +route: ordinary sessions use the current globally configured plugins, while a +managed session continues using only that run's isolated plugins. + +`context.operation` is `"insert"` or `"merge"`; `context.source` and +`context.session_id` identify the translated event stream; and `context.env` +contains the daemon process environment. Environment variable names are +uppercased on Windows so common lookups such as `context.env.PATH` remain +portable. + +The environment map is captured from the daemon process when each worker-local +span processor is constructed. Plugins execute in bounded, thread-local +QuickJS runtimes with no filesystem or network host APIs. Modules must be +self-contained and transforms must be stateless: module globals belong to a +worker thread, not a session. If a plugin fails, that worker reports and skips +only that plugin on subsequent spans; the remaining plugins continue to run. +Plugins are trusted local code: although they have no host APIs, they can copy +environment values into spans that are delivered to Braintrust. Read only the +specific variables needed by the transform; never attach `context.env` itself. + ### Additional root metadata `additional_metadata` is a JSON object merged into each traced session's root diff --git a/bt-daemon/config.json.example b/bt-daemon/config.json.example index 78355d2..ed1363b 100644 --- a/bt-daemon/config.json.example +++ b/bt-daemon/config.json.example @@ -14,6 +14,9 @@ "additional_metadata": { "team": "platform", "environment": "development" - } + }, + "span_plugins": [ + "/absolute/path/to/redact.mjs" + ] } } diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index 6bc832f..b7e8c94 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -199,7 +199,8 @@ Used for version handover and by tests. "project_name": "codex" }, "flush_mode": "fire_and_forget", - "additional_metadata": { "…": "…" } + "additional_metadata": { "…": "…" }, + "span_plugins": ["/absolute/path/redact.mjs"] } } ``` @@ -248,7 +249,9 @@ Field notes: Live credentials returned by the host provider are **never** written to the journal, logs, status, or RPC response. Envelopes journal only their non-secret -`route`, allowing restart recovery to resolve a fresh lease. +`route`, allowing restart recovery to resolve a fresh lease. Span plugins read +an environment snapshot captured inside their daemon worker process; it is not +part of the envelope or journal schema. ## Daemon lifecycle @@ -315,8 +318,10 @@ profiles, organizations, and destinations while sharing one daemon. `$HOME/.braintrust/state/bt-daemon` on Unix, and `%LOCALAPPDATA%\Braintrust\bt-daemon` on Windows. On restart the daemon rebuilds each route's unfinished correlation state independently, replaying - only the journal entries whose `route` matches that pipeline into a fresh - translator. The resulting rows may be resubmitted to repair delivery + only the journal entries whose delivery route matches that pipeline into a + fresh translator. Span plugin paths are ignored for this comparison so raw + events can be replayed through the current plugin chain. The resulting rows + may be resubmitted to repair delivery interrupted by a crash, but their deterministic ids target the same backend rows and must never create duplicate spans, and a route never receives another route's rows. Replay streams the journal and is bounded to the diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index d065898..83aae51 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -318,7 +318,36 @@ impl SessionActor { }; while let Some(ops) = next { if !ops.is_empty() { - match sink.emit(&ops).await { + let plugin_paths = ctx + .config + .as_ref() + .map(|config| config.span_plugins.as_slice()) + .unwrap_or_default(); + let mut processed = Vec::with_capacity(ops.len()); + for op in &ops { + match crate::span_processor::process( + plugin_paths, + op, + &self.source, + &self.session_id, + ) { + Ok(result) => { + for failure in result.failures { + self.set_error(format!( + "span plugin {} failed; disabled on this worker: {}", + failure.path.display(), + failure.message + )); + } + processed.push(result.op); + } + Err(error) => { + self.set_error(format!("span plugin processor failed: {error}")); + processed.push(op.clone()); + } + } + } + match sink.emit(&processed).await { Ok(n) => { self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); } @@ -370,7 +399,7 @@ impl SessionActor { if !entry .route .as_ref() - .is_some_and(|candidate| candidate.same_route(&plan.route)) + .is_some_and(|candidate| candidate.same_replay_route(&plan.route)) { continue; } diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 762a8ad..c36a31c 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -21,6 +21,7 @@ mod server; mod settings; mod setup; mod sink; +mod span_processor; mod trace_command; mod trace_runtime; mod transcript_import; @@ -163,6 +164,10 @@ pub struct ImportArgs { /// JSON object merged into every imported root span's metadata. #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] pub additional_metadata: Option, + /// JavaScript span transform for this import. Repeat to compose an isolated + /// transform chain; persistent setup plugins are not included. + #[arg(long, value_name = "PATH")] + pub plugin: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] @@ -182,6 +187,10 @@ pub struct RunArgs { /// JSON object merged into root-span metadata for this invocation. #[arg(long, env = "BRAINTRUST_ADDITIONAL_METADATA")] pub additional_metadata: Option, + /// JavaScript span transform for this invocation. Repeat to compose an + /// isolated transform chain; persistent setup plugins are not included. + #[arg(long, value_name = "PATH")] + pub plugin: Vec, /// Arguments forwarded verbatim to the coding agent. #[arg(allow_hyphen_values = true)] pub agent_args: Vec, @@ -465,6 +474,7 @@ pub async fn run_import( mut config: Option, ) -> anyhow::Result<()> { validate_import_selection(&args)?; + apply_import_span_plugins(&mut config, &args.plugin)?; let destination = args .parent .map(|components| wire::TraceDestination::ParentSpan { components }) @@ -477,6 +487,31 @@ pub async fn run_import( import_transcripts(&files, args.source, opts, config).await } +fn apply_import_span_plugins( + config: &mut Option, + plugins: &[PathBuf], +) -> anyhow::Result<()> { + let plugins = resolve_span_plugin_paths(plugins)?; + if let Some(config) = config.as_mut() { + config.span_plugins = plugins; + } else if !plugins.is_empty() { + *config = Some(SessionConfig { + auth: wire::BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + destination: None, + flush_mode: wire::FlushMode::FireAndForget, + additional_metadata: None, + span_plugins: plugins, + }); + } + Ok(()) +} + fn validate_import_selection(args: &ImportArgs) -> anyhow::Result<()> { if args.all != args.session_ids.is_empty() { anyhow::bail!("provide explicit session ids or use --all, but not both"); @@ -493,8 +528,9 @@ fn validate_import_selection(args: &ImportArgs) -> anyhow::Result<()> { pub async fn run_traced( args: RunArgs, hook_command: RunHookCommand, - route: SessionRoute, + mut route: SessionRoute, ) -> anyhow::Result { + apply_run_span_plugins(&mut route, &args.plugin)?; if route.destination.is_none() { anyhow::bail!( "managed run requires a trace destination; select a project, object destination, or parent span" @@ -562,6 +598,24 @@ pub async fn run_traced( status } +fn apply_run_span_plugins(route: &mut SessionRoute, plugins: &[PathBuf]) -> anyhow::Result<()> { + route.span_plugins = resolve_span_plugin_paths(plugins)?; + Ok(()) +} + +pub(crate) fn resolve_span_plugin_paths(paths: &[PathBuf]) -> anyhow::Result> { + let paths: Vec<_> = paths + .iter() + .map(|path| { + path.canonicalize().map_err(|error| { + anyhow::anyhow!("could not resolve span plugin {}: {error}", path.display()) + }) + }) + .collect::>()?; + crate::span_processor::validate(&paths)?; + Ok(paths) +} + fn managed_run_args( source: RunSource, hook_command: &RunHookCommand, @@ -800,6 +854,7 @@ pub async fn import_transcripts( } struct ImportLive { + source: String, translator: Box, sink: Box, ctx: SessionCtx, @@ -837,6 +892,7 @@ impl ImportProcessor { self.sessions.insert( sid.clone(), ImportLive { + source: env.source.clone(), translator, sink, ctx: SessionCtx { @@ -870,7 +926,42 @@ impl ImportProcessor { // for every native turn boundary. const FLUSH_OPS: usize = 500; for chunk in ops.chunks(FLUSH_OPS) { - live.sink.emit(chunk).await?; + let plugins = live + .ctx + .config + .as_ref() + .map(|config| config.span_plugins.as_slice()) + .unwrap_or_default(); + let mut transformed = Vec::with_capacity(chunk.len()); + for op in chunk { + match crate::span_processor::process( + plugins, + op, + &live.source, + &live.ctx.session_id, + ) { + Ok(result) => { + for failure in result.failures { + tracing::warn!( + session_id = %live.ctx.session_id, + plugin = %failure.path.display(), + error = %failure.message, + "span plugin failed during import; disabled on this worker" + ); + } + transformed.push(result.op); + } + Err(error) => { + tracing::warn!( + session_id = %live.ctx.session_id, + %error, + "span plugin processor failed during import" + ); + transformed.push(op.clone()); + } + } + } + live.sink.emit(&transformed).await?; live.pending_ops += chunk.len(); if live.pending_ops >= FLUSH_OPS { live.sink.flush().await?; @@ -981,6 +1072,22 @@ mod tests { assert_eq!(args.session_idle_timeout_secs, 30); } + #[test] + fn span_plugin_paths_are_canonicalized_before_use() { + let dir = tempfile::Builder::new() + .prefix("span-plugin-path-") + .tempdir_in(".") + .unwrap(); + let plugin = dir.path().join("plugin.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let relative = PathBuf::from(dir.path().file_name().unwrap()).join("plugin.mjs"); + + let resolved = resolve_span_plugin_paths(&[relative]).unwrap(); + + assert_eq!(resolved, [plugin.canonicalize().unwrap()]); + assert!(resolved[0].is_absolute()); + } + #[test] fn additional_metadata_overrides_a_route_only_with_a_json_object() { let mut route = SessionRoute { @@ -1001,6 +1108,52 @@ mod tests { .contains("invalid --additional-metadata JSON")); } + #[test] + fn managed_run_plugins_replace_inherited_plugins() { + let temp = tempfile::tempdir().unwrap(); + let plugin = temp.path().join("run.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let mut route = SessionRoute { + span_plugins: vec![PathBuf::from("persisted.mjs")], + ..SessionRoute::default() + }; + + apply_run_span_plugins(&mut route, &[]).unwrap(); + assert!(route.span_plugins.is_empty()); + + apply_run_span_plugins(&mut route, std::slice::from_ref(&plugin)).unwrap(); + assert_eq!(route.span_plugins, [plugin.canonicalize().unwrap()]); + } + + #[test] + fn import_plugins_replace_inherited_plugins() { + let temp = tempfile::tempdir().unwrap(); + let plugin = temp.path().join("import.mjs"); + std::fs::write(&plugin, "export default span => span").unwrap(); + let mut config = Some(SessionConfig { + auth: wire::BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + destination: None, + flush_mode: wire::FlushMode::FireAndForget, + additional_metadata: None, + span_plugins: vec![PathBuf::from("persisted.mjs")], + }); + + apply_import_span_plugins(&mut config, &[]).unwrap(); + assert!(config.as_ref().unwrap().span_plugins.is_empty()); + + apply_import_span_plugins(&mut config, std::slice::from_ref(&plugin)).unwrap(); + assert_eq!( + config.unwrap().span_plugins, + [plugin.canonicalize().unwrap()] + ); + } + #[test] fn import_args_accept_multiple_sessions_or_all() { let explicit = ImportCli::try_parse_from([ @@ -1037,6 +1190,7 @@ mod tests { parent: None, attach: true, additional_metadata: None, + plugin: Vec::new(), }; assert!(validate_import_selection(&args) .unwrap_err() @@ -1095,6 +1249,7 @@ mod tests { RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: Vec::new(), }, test_run_hook_command(), diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index 4792f78..4971767 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -341,6 +341,13 @@ fn enable_tracing_at(path: &Path, mut route: SessionRoute) -> anyhow::Result<()> .filter(|metadata| metadata.is_object()) .cloned(); } + if route.span_plugins.is_empty() { + route.span_plugins = settings + .get("route") + .and_then(|route| route.get("span_plugins")) + .and_then(|plugins| serde_json::from_value(plugins.clone()).ok()) + .unwrap_or_default(); + } settings.insert("trace_to_braintrust".into(), Value::Bool(true)); settings.insert("route".into(), serde_json::to_value(route)?); settings.remove("traceToBraintrust"); @@ -739,4 +746,32 @@ mod tests { assert_eq!(config["plugin"], serde_json::json!(["other"])); assert_eq!(config["model"], "test/model"); } + + #[test] + fn tracing_settings_preserve_plugins_until_setup_explicitly_replaces_them() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("braintrust.json"); + std::fs::write(&path, r#"{"route":{"span_plugins":["old.mjs"]}}"#).unwrap(); + + enable_tracing_at(&path, SessionRoute::default()).unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["span_plugins"], + serde_json::json!(["old.mjs"]) + ); + + enable_tracing_at( + &path, + SessionRoute { + span_plugins: vec![PathBuf::from("first.mjs"), PathBuf::from("second.mjs")], + ..SessionRoute::default() + }, + ) + .unwrap(); + let settings: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + settings["route"]["span_plugins"], + serde_json::json!(["first.mjs", "second.mjs"]) + ); + } } diff --git a/bt-daemon/src/span_processor.rs b/bt-daemon/src/span_processor.rs new file mode 100644 index 0000000..2fc8aa9 --- /dev/null +++ b/bt-daemon/src/span_processor.rs @@ -0,0 +1,403 @@ +//! Synchronous JavaScript span transforms. +//! +//! Session actors already execute on Tokio's worker pool. Each worker thread +//! lazily owns one QuickJS runtime and module cache, so JavaScript values never +//! cross threads and unrelated workers can transform spans concurrently. + +use crate::translate::{SpanOp, SpanRow}; +use rquickjs::{Context, Function, Module, Persistent, Runtime}; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use std::cell::RefCell; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime}; + +const MEMORY_LIMIT_BYTES: usize = 64 * 1024 * 1024; +const STACK_LIMIT_BYTES: usize = 512 * 1024; +const CALL_TIMEOUT: Duration = Duration::from_millis(50); + +thread_local! { + static ENGINE: RefCell> = const { RefCell::new(None) }; +} + +#[derive(Clone, Copy, Serialize)] +#[serde(rename_all = "snake_case")] +enum Operation { + Insert, + Merge, +} + +#[derive(Serialize)] +struct PluginContext<'a> { + operation: Operation, + source: &'a str, + session_id: &'a str, + env: &'a BTreeMap, +} + +struct Engine { + modules: HashMap, + failed_plugins: HashSet, + env: BTreeMap, + context: Context, + started: Instant, + deadline_ms: Arc, + // Must drop after every Context and persistent JavaScript value. + _runtime: Runtime, +} + +struct CachedModule { + modified: Option, + len: u64, + function: Persistent>, +} + +impl Engine { + fn new() -> anyhow::Result { + let runtime = Runtime::new()?; + runtime.set_memory_limit(MEMORY_LIMIT_BYTES); + runtime.set_max_stack_size(STACK_LIMIT_BYTES); + let started = Instant::now(); + let deadline_ms = Arc::new(AtomicU64::new(0)); + let interrupt_deadline = deadline_ms.clone(); + let interrupt_started = started; + runtime.set_interrupt_handler(Some(Box::new(move || { + let deadline = interrupt_deadline.load(Ordering::Relaxed); + deadline != 0 && interrupt_started.elapsed().as_millis() as u64 >= deadline + }))); + let context = Context::full(&runtime)?; + Ok(Self { + modules: HashMap::new(), + failed_plugins: HashSet::new(), + env: environment(), + context, + started, + deadline_ms, + _runtime: runtime, + }) + } + + fn load(&mut self, path: &Path) -> anyhow::Result>> { + let metadata = std::fs::metadata(path) + .map_err(|error| anyhow::anyhow!("failed to inspect {}: {error}", path.display()))?; + let modified = metadata.modified().ok(); + if let Some(module) = self.modules.get(path) { + if module.modified == modified && module.len == metadata.len() { + return Ok(module.function.clone()); + } + } + let source = std::fs::read(path) + .map_err(|error| anyhow::anyhow!("failed to read {}: {error}", path.display()))?; + let digest = Sha256::digest(&source); + let name = format!("bt-span-plugin:{digest:x}"); + self.arm_deadline(); + let function = self.context.with(|ctx| -> anyhow::Result<_> { + let (module, promise) = Module::declare(ctx.clone(), name, source)?.eval()?; + promise.finish::<()>()?; + let function: Function<'_> = module + .get("default") + .map_err(|error| anyhow::anyhow!("default export is not a function: {error}"))?; + Ok(Persistent::save(&ctx, function)) + }); + self.deadline_ms.store(0, Ordering::Relaxed); + let function = function?; + self.modules.insert( + path.to_path_buf(), + CachedModule { + modified, + len: metadata.len(), + function: function.clone(), + }, + ); + Ok(function) + } + + fn call( + &mut self, + path: &Path, + row: &SpanRow, + operation: Operation, + source: &str, + session_id: &str, + ) -> anyhow::Result { + let function = self.load(path)?; + let context = PluginContext { + operation, + source, + session_id, + env: &self.env, + }; + self.arm_deadline(); + let result = self.context.with(|ctx| -> anyhow::Result { + let function = function.restore(&ctx)?; + let row = rquickjs_serde::to_value(ctx.clone(), row)?; + let context = rquickjs_serde::to_value(ctx.clone(), &context)?; + let result = function.call::<_, rquickjs::Value<'_>>((row, context))?; + if result.as_promise().is_some() { + anyhow::bail!("plugin returned a Promise; span plugins must be synchronous"); + } + Ok(rquickjs_serde::from_value_strict(result)?) + }); + self.deadline_ms.store(0, Ordering::Relaxed); + result + } + + fn arm_deadline(&self) { + let deadline = self + .started + .elapsed() + .saturating_add(CALL_TIMEOUT) + .as_millis() as u64; + self.deadline_ms.store(deadline.max(1), Ordering::Relaxed); + } +} + +#[derive(Debug)] +pub struct PluginFailure { + pub path: PathBuf, + pub message: String, +} + +pub struct ProcessResult { + pub op: SpanOp, + pub failures: Vec, +} + +/// Apply an ordered plugin chain on the worker thread currently executing the +/// session actor. A failing plugin is skipped on subsequent calls handled by +/// this worker, while the rest of the ordered chain continues to run. +pub fn process( + plugins: &[PathBuf], + op: &SpanOp, + source: &str, + session_id: &str, +) -> anyhow::Result { + if plugins.is_empty() { + return Ok(ProcessResult { + op: op.clone(), + failures: Vec::new(), + }); + } + let (operation, mut row) = match op { + SpanOp::Insert(row) => (Operation::Insert, row.clone()), + SpanOp::Merge(row) => (Operation::Merge, row.clone()), + }; + let original_ids = ( + row.span_id.clone(), + row.root_span_id.clone(), + row.parent_span_ids.clone(), + ); + let mut failures = Vec::new(); + ENGINE.with_borrow_mut(|slot| -> anyhow::Result<()> { + if slot.is_none() { + *slot = Some(Engine::new()?); + } + let engine = slot.as_mut().expect("engine initialized"); + for plugin in plugins { + if engine.failed_plugins.contains(plugin) { + continue; + } + let candidate = engine.call(plugin, &row, operation, source, session_id); + let candidate = match candidate { + Ok(candidate) + if ( + candidate.span_id.as_str(), + candidate.root_span_id.as_str(), + &candidate.parent_span_ids, + ) == ( + original_ids.0.as_str(), + original_ids.1.as_str(), + &original_ids.2, + ) => + { + candidate + } + Ok(_) => { + let message = "changed immutable span identity fields".to_owned(); + engine.failed_plugins.insert(plugin.clone()); + failures.push(PluginFailure { + path: plugin.clone(), + message, + }); + continue; + } + Err(error) => { + engine.failed_plugins.insert(plugin.clone()); + failures.push(PluginFailure { + path: plugin.clone(), + message: error.to_string(), + }); + continue; + } + }; + row = candidate; + } + Ok(()) + })?; + Ok(ProcessResult { + op: match op { + SpanOp::Insert(_) => SpanOp::Insert(row), + SpanOp::Merge(_) => SpanOp::Merge(row), + }, + failures, + }) +} + +/// Compile each module and verify that it default-exports a function. Explicit +/// CLI commands call this before persisting or launching with a plugin chain. +pub fn validate(plugins: &[PathBuf]) -> anyhow::Result<()> { + ENGINE.with_borrow_mut(|slot| -> anyhow::Result<()> { + if slot.is_none() { + *slot = Some(Engine::new()?); + } + let engine = slot.as_mut().expect("engine initialized"); + for plugin in plugins { + engine.load(plugin)?; + } + Ok(()) + }) +} + +fn environment() -> BTreeMap { + std::env::vars_os() + .filter_map(|(key, value)| { + let key = key.into_string().ok()?; + let value = value.into_string().ok()?; + // Windows environment variable names are case-insensitive, while + // JavaScript object properties are not. Use a stable casing there + // so portable plugins can read conventional names such as PATH. + #[cfg(windows)] + let key = key.to_ascii_uppercase(); + Some((key, value)) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn row() -> SpanRow { + SpanRow { + span_id: "span".into(), + root_span_id: "root".into(), + name: "original".into(), + ..SpanRow::default() + } + } + + #[test] + fn composes_plugins_and_exposes_context_env() { + let dir = tempfile::tempdir().unwrap(); + let first = dir.path().join("first.mjs"); + let second = dir.path().join("second.mjs"); + std::fs::write( + &first, + "export default (span, context) => ({...span, name: `${context.source}:${context.env.PATH}:${span.name}`})", + ) + .unwrap(); + std::fs::write( + &second, + "export default span => ({...span, metadata: {second: true}})", + ) + .unwrap(); + let result = process(&[first, second], &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert!(result.failures.is_empty()); + let SpanOp::Insert(processed) = result.op else { + panic!("expected insert") + }; + assert_eq!( + processed.name, + format!("codex:{}:original", std::env::var("PATH").unwrap()) + ); + assert_eq!(processed.metadata.unwrap()["second"], true); + } + + #[test] + fn rejects_identity_changes_without_dropping_the_span() { + let dir = tempfile::tempdir().unwrap(); + let plugin = dir.path().join("identity.mjs"); + std::fs::write( + &plugin, + "export default span => ({...span, span_id: 'different'})", + ) + .unwrap(); + let result = process(&[plugin], &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert_eq!(result.failures.len(), 1); + assert!(result.failures[0] + .message + .contains("immutable span identity")); + let SpanOp::Insert(processed) = result.op else { + panic!("expected insert") + }; + assert_eq!(processed.span_id, "span"); + } + + #[test] + fn interrupts_runaway_plugins_and_rejects_promises() { + let dir = tempfile::tempdir().unwrap(); + let runaway = dir.path().join("runaway.mjs"); + std::fs::write(&runaway, "export default span => { while (true) {} }").unwrap(); + let started = Instant::now(); + let result = process(&[runaway], &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert_eq!(result.failures.len(), 1); + assert!(started.elapsed() < Duration::from_secs(2)); + + let asynchronous = dir.path().join("async.mjs"); + std::fs::write( + &asynchronous, + "export default async span => ({...span, name: 'later'})", + ) + .unwrap(); + let result = process(&[asynchronous], &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert_eq!(result.failures.len(), 1); + assert!(result.failures[0].message.contains("must be synchronous")); + + let non_json = dir.path().join("non-json.mjs"); + std::fs::write(&non_json, "export default () => Symbol('not-json')").unwrap(); + let result = process(&[non_json], &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert_eq!(result.failures.len(), 1); + } + + #[test] + fn skips_only_the_failed_plugin_and_continues_the_chain() { + let dir = tempfile::tempdir().unwrap(); + let first = dir.path().join("first.mjs"); + let broken = dir.path().join("broken.mjs"); + let last = dir.path().join("last.mjs"); + std::fs::write( + &first, + "export default span => ({...span, name: `first:${span.name}`})", + ) + .unwrap(); + std::fs::write( + &broken, + "export default () => { throw new Error('broken') }", + ) + .unwrap(); + std::fs::write( + &last, + "export default span => ({...span, name: `last:${span.name}`})", + ) + .unwrap(); + let plugins = [first, broken.clone(), last]; + + let first_result = process(&plugins, &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert_eq!(first_result.failures.len(), 1); + assert_eq!(first_result.failures[0].path, broken); + let SpanOp::Insert(first_row) = first_result.op else { + panic!("expected insert") + }; + assert_eq!(first_row.name, "last:first:original"); + + let second_result = process(&plugins, &SpanOp::Insert(row()), "codex", "session").unwrap(); + assert!(second_result.failures.is_empty()); + let SpanOp::Insert(second_row) = second_result.op else { + panic!("expected insert") + }; + assert_eq!(second_row.name, "last:first:original"); + } +} diff --git a/bt-daemon/src/trace_command.rs b/bt-daemon/src/trace_command.rs index 235ed44..5bed608 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -56,6 +56,10 @@ pub struct EnableArgs { /// JSON object persisted in this agent's tracing route and merged into root-span metadata. #[arg(long, global = true, env = "BRAINTRUST_ADDITIONAL_METADATA")] pub additional_metadata: Option, + /// JavaScript span transform to persist for this agent. Repeat to compose + /// transforms in order. + #[arg(long, global = true, value_name = "PATH")] + pub plugin: Vec, } /// Backwards-compatible API name for hosts that mounted the former setup command. @@ -106,6 +110,7 @@ mod tests { TraceCommand::Setup(SetupArgs { agent: SetupAgent::Claude, additional_metadata: Some(ref value), + .. }) if value == r#"{"setup":true}"# )); @@ -168,4 +173,50 @@ mod tests { }) if value == r#"{"import":true}"# )); } + + #[test] + fn public_commands_preserve_repeated_plugin_order() { + for args in [ + vec![ + "bt", + "setup", + "codex", + "--plugin", + "first.mjs", + "--plugin", + "second.mjs", + ], + vec![ + "bt", + "run", + "--plugin", + "first.mjs", + "--plugin", + "second.mjs", + "codex", + ], + vec![ + "bt", + "import", + "codex", + "session", + "--plugin", + "first.mjs", + "--plugin", + "second.mjs", + ], + ] { + let parsed = Cli::try_parse_from(args).unwrap(); + let plugins = match parsed.trace.command { + TraceCommand::Setup(args) => args.plugin, + TraceCommand::Run(args) => args.plugin, + TraceCommand::Import(args) => args.plugin, + _ => unreachable!(), + }; + assert_eq!( + plugins, + [PathBuf::from("first.mjs"), PathBuf::from("second.mjs")] + ); + } + } } diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 39a8731..499dcfc 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -136,6 +136,7 @@ async fn session_config( destination: route.destination.clone(), flush_mode: route.flush_mode, additional_metadata: route.additional_metadata.clone(), + span_plugins: route.span_plugins.clone(), }) } @@ -198,6 +199,9 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul ) .await?; apply_additional_metadata(&mut route, enable_args.additional_metadata.as_deref())?; + if !enable_args.plugin.is_empty() { + route.span_plugins = crate::resolve_span_plugin_paths(&enable_args.plugin)?; + } print_output(run_enable(enable_args, route)?, host.output_format) } TraceCommand::Disable(disable_args) => { @@ -416,10 +420,12 @@ mod tests { TraceCommand::Setup(SetupArgs { agent: SetupAgent::OpenCode, additional_metadata: None, + plugin: Vec::new(), }), TraceCommand::Run(RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: Vec::new(), }), ] { @@ -489,6 +495,7 @@ mod tests { parent: None, attach: false, additional_metadata: None, + plugin: Vec::new(), }; let error = run_trace( TraceArgs { diff --git a/bt-daemon/src/wire/envelope.rs b/bt-daemon/src/wire/envelope.rs index f61cc4f..384ad3a 100644 --- a/bt-daemon/src/wire/envelope.rs +++ b/bt-daemon/src/wire/envelope.rs @@ -3,6 +3,7 @@ use braintrust_sdk_rust::SpanComponents; use serde::{Deserialize, Serialize}; +use std::path::PathBuf; /// One captured hook event, forwarded from a shim to the daemon. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -63,6 +64,10 @@ pub struct SessionRoute { pub flush_mode: FlushMode, #[serde(default, skip_serializing_if = "Option::is_none")] pub additional_metadata: Option, + /// Ordered JavaScript span transforms. Paths are resolved by explicit + /// setup, run, and import commands before entering a session route. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub span_plugins: Vec, } impl SessionRoute { @@ -72,12 +77,23 @@ impl SessionRoute { destination: self.destination.clone(), flush_mode: self.flush_mode, additional_metadata: self.additional_metadata.clone(), + span_plugins: self.span_plugins.clone(), } } pub fn same_route(&self, other: &Self) -> bool { serde_json::to_value(self).ok() == serde_json::to_value(other).ok() } + + /// Raw journal entries can be replayed through a newer plugin chain as + /// long as their Braintrust delivery route is otherwise unchanged. + pub fn same_replay_route(&self, other: &Self) -> bool { + let mut left = self.clone(); + let mut right = other.clone(); + left.span_plugins.clear(); + right.span_plugins.clear(); + left.same_route(&right) + } } /// Trace settings and backend credentials resolved by the shim. @@ -91,6 +107,8 @@ pub struct SessionConfig { pub flush_mode: FlushMode, #[serde(default, skip_serializing_if = "Option::is_none")] pub additional_metadata: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub span_plugins: Vec, } /// Where a session's root span should be logged. @@ -264,6 +282,7 @@ mod tests { destination: None, flush_mode: FlushMode::FireAndForget, additional_metadata: None, + span_plugins: Vec::new(), }), } } diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 4cc622a..79aa4d7 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -26,6 +26,7 @@ fn session_config(base: &str) -> SessionConfig { }), flush_mode: FlushMode::FireAndForget, additional_metadata: None, + span_plugins: Vec::new(), } } diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 955d574..5f5a12d 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -493,6 +493,7 @@ fn configured_ctx(session_id: &str, additional_metadata: Value) -> SessionCtx { }), flush_mode: FlushMode::FireAndForget, additional_metadata: Some(additional_metadata), + span_plugins: Vec::new(), }), } } diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 44bbeae..ea5c604 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -980,6 +980,134 @@ async fn restart_replays_journal_with_stable_span_ids_before_new_events() { second.await.unwrap(); } +#[tokio::test] +async fn span_plugins_transform_live_and_replayed_rows_with_daemon_environment() { + let (data_dir, socket, first, tmp) = start_daemon().await; + let host = dummy_host(); + let first_plugin = tmp.path().join("first.mjs"); + let second_plugin = tmp.path().join("second.mjs"); + std::fs::write( + &first_plugin, + "export default (span, context) => ({...span, name: `${context.env.PATH}:${span.name}`})", + ) + .unwrap(); + std::fs::write( + &second_plugin, + "export default (span, context) => ({...span, name: `${context.operation}:current:${span.name}`})", + ) + .unwrap(); + + let mut start = envelope("plugin-replay", "SessionStart", 1); + start + .route + .as_mut() + .unwrap() + .span_plugins + .push(first_plugin); + forward_envelope(&start, &socket, &host, false) + .await + .unwrap(); + flush_session("plugin-replay", &socket, 5000).await.unwrap(); + shutdown(&socket).await; + first.await.unwrap(); + + let second = start_daemon_at(data_dir.clone(), socket.clone()).await; + let mut stop = envelope("plugin-replay", "Stop", 2); + stop.route + .as_mut() + .unwrap() + .span_plugins + .push(second_plugin); + forward_envelope(&stop, &socket, &host, false) + .await + .unwrap(); + flush_session("plugin-replay", &socket, 5000).await.unwrap(); + + let spans = std::fs::read_to_string(data_dir.join("spans/plugin-replay.ndjson")).unwrap(); + let names: Vec<_> = spans + .lines() + .filter_map(|line| { + let value: serde_json::Value = serde_json::from_str(line).unwrap(); + value + .get("Insert") + .or_else(|| value.get("Merge")) + .and_then(|row| row.get("name")) + .and_then(serde_json::Value::as_str) + .map(str::to_owned) + }) + .collect(); + let path_prefix = format!("{}:", std::env::var("PATH").unwrap()); + assert!(names.iter().any(|name| name.starts_with(&path_prefix))); + assert!( + names.iter().any(|name| name.contains(":current:")), + "the new plugin chain should process replayed and live rows: {names:?}" + ); + assert!( + !names + .iter() + .any(|name| name.contains(&format!(":current:{path_prefix}"))), + "recovery should replace the journal's old plugin chain with the resumed route: {names:?}" + ); + + shutdown(&socket).await; + second.await.unwrap(); +} + +#[tokio::test] +async fn a_failing_span_plugin_is_reported_and_fails_open() { + let (data_dir, socket, handle, tmp) = start_daemon().await; + let plugin = tmp.path().join("bad.mjs"); + let later_plugin = tmp.path().join("later.mjs"); + std::fs::write( + &plugin, + "export default span => ({...span, span_id: 'corrupt'})", + ) + .unwrap(); + std::fs::write( + &later_plugin, + "export default span => ({...span, name: `after-failure:${span.name}`})", + ) + .unwrap(); + let mut env = envelope("plugin-failure", "SessionStart", 1); + env.route + .as_mut() + .unwrap() + .span_plugins + .extend([plugin, later_plugin]); + forward_envelope(&env, &socket, &dummy_host(), false) + .await + .unwrap(); + flush_session("plugin-failure", &socket, 5000) + .await + .unwrap(); + + let status = run_status(StatusArgs { + socket: Some(socket.clone()), + session_id: Some("plugin-failure".into()), + }) + .await + .unwrap() + .unwrap(); + assert!( + status.sessions[0] + .last_error + .as_deref() + .is_some_and(|error| error.contains("failed; disabled on this worker")), + "unexpected plugin status: {:?}", + status.sessions[0].last_error + ); + let spans = std::fs::read_to_string(data_dir.join("spans/plugin-failure.ndjson")).unwrap(); + assert!(!spans.contains("corrupt")); + assert!(spans.contains("after-failure:")); + assert!( + !spans.is_empty(), + "the original rows should still be delivered" + ); + + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn claude_boundary_journal_references_a_self_contained_transcript_mirror() { let (data_dir, socket, handle, tmp) = start_daemon().await; @@ -1310,6 +1438,7 @@ esac RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: vec![session_id.into(), mode.into()], }, RunHookCommand { diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 4ac4120..d7d70f3 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -1,3 +1,4 @@ +use bt_daemon::wire::{BackendAuth, FlushMode, SessionConfig}; use bt_daemon::{ import_transcript, import_transcripts, DebugSinkFactory, ImportSource, Registry, ServeOptions, }; @@ -69,6 +70,56 @@ fn inserted(rows: &[Value], span_type: &str) -> usize { .count() } +#[tokio::test] +async fn import_uses_the_same_span_plugin_stage() { + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("plugin-import.jsonl"); + write_jsonl( + &transcript, + &[ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"plugin-import","cwd":"/tmp/demo"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"done"}}), + ], + ); + let plugin = tmp.path().join("import.mjs"); + std::fs::write( + &plugin, + "export default span => ({...span, name: `imported:${span.name}`})", + ) + .unwrap(); + let output = tmp.path().join("spans"); + import_transcript( + &transcript, + ImportSource::Codex, + options(&output), + Some(SessionConfig { + auth: BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: None, + org_id: None, + }, + destination: None, + flush_mode: FlushMode::FireAndForget, + additional_metadata: None, + span_plugins: vec![plugin], + }), + false, + ) + .await + .unwrap(); + + let output = rows(&output.join("plugin-import.ndjson")); + assert!(output + .iter() + .filter_map(|op| op.get("Insert")) + .all(|row| row["name"] + .as_str() + .is_some_and(|name| name.starts_with("imported:")))); +} + #[tokio::test] async fn imports_multiple_transcripts_in_one_invocation() { let tmp = tempfile::tempdir().unwrap(); diff --git a/src/runtime/js-daemon-client/src/index.ts b/src/runtime/js-daemon-client/src/index.ts index 1e6c6b9..e99f214 100644 --- a/src/runtime/js-daemon-client/src/index.ts +++ b/src/runtime/js-daemon-client/src/index.ts @@ -14,6 +14,7 @@ export interface DaemonSessionRoute { destination: unknown flush_mode?: "fire_and_forget" | "flush_on_turn_end" additional_metadata?: Record + span_plugins?: string[] } export interface DaemonTraceSettings { diff --git a/src/runtime/js-daemon-client/tests/client.test.ts b/src/runtime/js-daemon-client/tests/client.test.ts index 53eff79..1e9ec75 100644 --- a/src/runtime/js-daemon-client/tests/client.test.ts +++ b/src/runtime/js-daemon-client/tests/client.test.ts @@ -110,6 +110,7 @@ test("serializes initialize, events, flush, and status over one connection", asy event: name, ts_ms: Date.now(), payload: {}, + route: { destination: {}, span_plugins: ["plugin.mjs"] }, }) assert.deepEqual(await Promise.all([client.log(envelope("one")), client.log(envelope("two"))]), [ true,