diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index 225091e..48d32b1 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", @@ -599,6 +607,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" @@ -757,6 +771,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" @@ -1417,6 +1436,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" @@ -1473,6 +1501,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 65784c7..47b4872 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 57acb25..f3323a4 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -205,7 +205,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"] } } ``` @@ -270,7 +271,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 @@ -337,8 +340,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 fdfad18..13d52dc 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -457,7 +457,36 @@ impl SessionActor { &ops, ); } - 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); } @@ -510,7 +539,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 64a50c3..d7f400b 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -23,6 +23,7 @@ mod server; mod settings; mod setup; mod sink; +mod span_processor; mod trace_command; mod trace_runtime; mod transcript_import; @@ -213,6 +214,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)] @@ -257,6 +262,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, @@ -553,6 +562,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 = import_parent_components(&args)? .map(|components| wire::TraceDestination::ParentSpan { components }) .or_else(|| args.destination.clone()); @@ -564,6 +574,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"); @@ -657,8 +692,9 @@ fn import_parent_components(args: &ImportArgs) -> anyhow::Result 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" @@ -786,6 +822,24 @@ impl ManagedRunRuntime { } } +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, @@ -1026,6 +1080,7 @@ pub async fn import_transcripts( } struct ImportLive { + source: String, translator: Box, sink: Box, ctx: SessionCtx, @@ -1066,6 +1121,7 @@ impl ImportProcessor { self.sessions.insert( sid.clone(), ImportLive { + source: env.source.clone(), translator, sink, ctx: SessionCtx { @@ -1117,7 +1173,42 @@ impl ImportProcessor { live.root_span_id = Some(row.root_span_id.clone()); } } - 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?; @@ -1283,6 +1374,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 { @@ -1303,6 +1410,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([ @@ -1404,6 +1557,7 @@ mod tests { parent_project: None, attach: true, additional_metadata: None, + plugin: Vec::new(), }; assert!(validate_import_selection(&args) .unwrap_err() @@ -1482,6 +1636,7 @@ mod tests { RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: Vec::new(), }, test_run_hook_command(), @@ -1499,6 +1654,7 @@ mod tests { RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: vec![OsString::from("--dangerously-bypass-hook-trust")], }, test_run_hook_command(), diff --git a/bt-daemon/src/setup.rs b/bt-daemon/src/setup.rs index 8258eaf..50d1466 100644 --- a/bt-daemon/src/setup.rs +++ b/bt-daemon/src/setup.rs @@ -475,6 +475,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)?); for key in [ @@ -1077,4 +1084,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 f0ae928..fca9500 100644 --- a/bt-daemon/src/trace_command.rs +++ b/bt-daemon/src/trace_command.rs @@ -98,6 +98,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. @@ -150,6 +154,7 @@ mod tests { TraceCommand::Setup(SetupArgs { agent: SetupAgent::Claude, additional_metadata: Some(ref value), + .. }) if value == r#"{"setup":true}"# )); @@ -260,4 +265,50 @@ mod tests { assert!(Cli::try_parse_from(["bt", "setup", "antigravity", "--disable"]).is_err()); } + + #[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 1438ec2..4eda398 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -178,6 +178,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(), }) } @@ -319,6 +320,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) => { @@ -549,6 +553,7 @@ mod tests { TraceCommand::Setup(SetupArgs { agent: SetupAgent::OpenCode, additional_metadata: None, + plugin: Vec::new(), }), true, ), @@ -556,6 +561,7 @@ mod tests { TraceCommand::Run(RunArgs { source: RunSource::Codex, additional_metadata: None, + plugin: Vec::new(), agent_args: Vec::new(), }), false, @@ -635,6 +641,7 @@ mod tests { parent_project: 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 7969dc1..8fb3cc1 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 operating-system process observed while capturing an event. /// @@ -138,6 +139,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 { @@ -147,12 +152,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. @@ -166,6 +182,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. @@ -349,6 +367,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 b1d054d..6bf3b85 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 896cd88..a17b140 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -507,6 +507,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 7d358f8..367aba0 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -1151,6 +1151,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; @@ -1534,6 +1662,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 c286a58..0555b61 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -107,6 +107,7 @@ async fn attached_import_summary_reports_the_effective_parent_root() { }), flush_mode: FlushMode::FireAndForget, additional_metadata: None, + span_plugins: Vec::new(), }; let summaries = import_transcript( @@ -123,6 +124,56 @@ async fn attached_import_summary_reports_the_effective_parent_root() { assert_eq!(summaries[0].root_span_id.as_deref(), Some(parent_root)); } +#[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 8a180a8..788e9e7 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 9036c67..e776cfa 100644 --- a/src/runtime/js-daemon-client/tests/client.test.ts +++ b/src/runtime/js-daemon-client/tests/client.test.ts @@ -127,6 +127,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,