From 40974a55ea84f11fc03bd4bbff60a57a76913a89 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 14 Aug 2026 00:55:13 +0800 Subject: [PATCH 1/3] Improve hook routing handling This allows sessions to having multiple routings, which is useful for scenarios like importing to multiple projects. --- bt-daemon/src/dispatch.rs | 99 +++++------ bt-daemon/src/lib.rs | 6 + bt-daemon/src/server.rs | 289 +++++++++++++++++++++++---------- bt-daemon/src/trace_runtime.rs | 78 ++++++++- bt-daemon/src/wire/methods.rs | 4 + 5 files changed, 344 insertions(+), 132 deletions(-) diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index f759738..b56fd45 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -8,7 +8,6 @@ //! Braintrust happens later in the actor; a downstream error never fails the //! caller's turn. -use crate::journal::JournalWriter; use crate::sink::SinkFactory; use crate::translate::{Registry, SessionCtx}; use crate::wire::Envelope; @@ -33,7 +32,6 @@ enum SessionMsg { pub struct Session { pub source: String, tx: mpsc::UnboundedSender, - journal: tokio::sync::Mutex, pub counters: Arc, pub last_error: Arc>>, pub permalink: Arc>>, @@ -45,8 +43,8 @@ impl Session { session_id: String, source: String, plugin_version: Option, - journal: JournalWriter, replay: Vec, + config: crate::wire::SessionConfig, translators: Arc, sink_factory: Arc, ) -> Arc { @@ -65,26 +63,21 @@ impl Session { last_error: last_error.clone(), permalink: permalink.clone(), replay, + config, }; tokio::spawn(actor.run(rx)); Arc::new(Session { source, tx, - journal: tokio::sync::Mutex::new(journal), counters, last_error, permalink, }) } - /// Journal (redacted) then enqueue. Both complete before the caller acks. - pub async fn append_and_enqueue(&self, mut env: Envelope) -> anyhow::Result<()> { - hydrate_transcript_snapshot(&mut env).await; - { - let mut j = self.journal.lock().await; - j.append(&env).await?; - } + /// Enqueue an event after the daemon has journaled it. + pub fn enqueue(&self, env: Envelope) -> anyhow::Result<()> { self.counters.queued.fetch_add(1, Ordering::Relaxed); self.tx .send(SessionMsg::Event(Box::new(env))) @@ -130,7 +123,7 @@ impl Session { /// journal at lifecycle boundaries so recovery/replay does not depend on a /// path that Claude may later rewrite or delete. Fail open: a missing file is /// handled by the translator exactly as before. -async fn hydrate_transcript_snapshot(env: &mut Envelope) { +pub(crate) async fn hydrate_transcript_snapshot(env: &mut Envelope) { if env.source != "claude-code" || !matches!( env.event.as_str(), @@ -173,6 +166,7 @@ struct SessionActor { last_error: Arc>>, permalink: Arc>>, replay: Vec, + config: crate::wire::SessionConfig, } impl SessionActor { @@ -189,15 +183,20 @@ impl SessionActor { // Still drain the queue so the daemon's counters settle and // callers waiting on flush don't hang. while let Some(msg) = rx.recv().await { - if let SessionMsg::Event(_) = msg { - self.counters.queued.fetch_sub(1, Ordering::Relaxed); - } else if let SessionMsg::Configure(_, r) = msg { - let _ = r.send(()); - } else if let SessionMsg::Flush(r) = msg { - let _ = r.send(0); - } else if let SessionMsg::Shutdown(r) = msg { - let _ = r.send(()); - break; + match msg { + SessionMsg::Event(_) => { + self.counters.queued.fetch_sub(1, Ordering::Relaxed); + } + SessionMsg::Configure(_, r) => { + let _ = r.send(()); + } + SessionMsg::Flush(r) => { + let _ = r.send(0); + } + SessionMsg::Shutdown(r) => { + let _ = r.send(()); + break; + } } } return; @@ -205,23 +204,15 @@ impl SessionActor { }; let mut ctx = SessionCtx { session_id: self.session_id.clone(), - config: None, + config: Some(self.config.clone()), }; - // Rebuild translator state before accepting the first new event. Keep - // the deterministic replay ops buffered until live credentials arrive; - // then re-emitting them repairs any rows lost by a prior crash. The - // stable span ids ensure these target existing rows rather than create - // duplicate spans. - let mut replay_ops = Vec::new(); - for env in &self.replay { - if let Some(cfg) = &env.config { - ctx.config = Some(cfg.clone()); - } - match translator.handle(env, &ctx) { - Ok(mut ops) => replay_ops.append(&mut ops), - Err(e) => self.set_error(format!("journal replay failed: {e}")), - } - } + sink.configure(&self.config); + self.refresh_permalink(sink.as_ref()); + // Rebuild translator state before accepting the first new event. + // Stable span ids make this both crash recovery and a complete copy + // when an existing source session is sent to another destination. + self.replay_into(&mut translator, &mut sink, &ctx, &self.replay) + .await; while let Some(msg) = rx.recv().await { match msg { @@ -231,15 +222,6 @@ impl SessionActor { ctx.config = Some(cfg.clone()); self.refresh_permalink(sink.as_ref()); } - if !replay_ops.is_empty() { - match sink.emit(&replay_ops).await { - Ok(n) => { - self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); - replay_ops.clear(); - } - Err(e) => self.set_error(format!("sink replay emit failed: {e}")), - } - } match translator.handle(&env, &ctx) { Ok(ops) => match sink.emit(&ops).await { Ok(n) => { @@ -270,6 +252,31 @@ impl SessionActor { } } + async fn replay_into( + &self, + translator: &mut Box, + sink: &mut Box, + ctx: &SessionCtx, + replay: &[Envelope], + ) { + let mut replay_ops = Vec::new(); + for env in replay { + match translator.handle(env, ctx) { + Ok(mut ops) => replay_ops.append(&mut ops), + Err(e) => self.set_error(format!("journal replay failed: {e}")), + } + } + if replay_ops.is_empty() { + return; + } + match sink.emit(&replay_ops).await { + Ok(n) => { + self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); + } + Err(e) => self.set_error(format!("sink replay emit failed: {e}")), + } + } + async fn drain_flush( &self, translator: &mut Box, diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 459ee8a..5888776 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -355,6 +355,7 @@ pub async fn flush_managed_run( return Ok(wire::FlushResult { flushed: true, pending: 0, + accepted_sessions: 0, }) } }; @@ -485,6 +486,11 @@ pub async fn run_traced( }; let socket = paths::socket_path(None); match flush_managed_run(&managed_run_id, &socket, MANAGED_RUN_FLUSH_TIMEOUT_MS).await { + Ok(result) if result.accepted_sessions == 0 => { + anyhow::bail!( + "managed run produced no accepted trace events; verify hook output and `bt trace status`" + ) + } Ok(result) if result.flushed => {} Ok(result) => tracing::warn!( managed_run_id, diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index 52d0bb3..2a940b2 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -2,7 +2,7 @@ //! serves JSON-RPC connections, and shuts down gracefully (idle timeout, //! `daemon.shutdown`, or SIGINT/SIGTERM). -use crate::dispatch::Session; +use crate::dispatch::{hydrate_transcript_snapshot, Session}; use crate::journal::{self, JournalWriter}; use crate::sink::SinkFactory; use crate::translate::Registry; @@ -70,16 +70,35 @@ struct SessionAuthState { lease: AuthLease, } +/// One independent delivery pipeline for a source session and the exact route +/// carried by its hook or import envelope. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct DeliveryKey { + session_id: String, + route: String, +} + +impl DeliveryKey { + fn new(session_id: &str, route: &SessionRoute) -> anyhow::Result { + Ok(Self { + session_id: session_id.to_string(), + route: serde_json::to_string(route)?, + }) + } +} + pub struct Daemon { version: String, data_dir: PathBuf, translators: Arc, sink_factory: Arc, auth_provider: Option>, - session_auth: tokio::sync::Mutex>, - managed_run_sessions: Mutex>>, - auth_errors: Mutex>, - sessions: Mutex>>, + session_auth: tokio::sync::Mutex>, + session_locks: Mutex>>>, + journals: Mutex>>>, + managed_run_sessions: Mutex>>, + auth_errors: Mutex>, + sessions: Mutex>>, started: Instant, last_activity: Mutex, shutting_down: AtomicBool, @@ -95,6 +114,8 @@ impl Daemon { sink_factory: opts.sink_factory, auth_provider: opts.auth_provider, session_auth: tokio::sync::Mutex::new(HashMap::new()), + session_locks: Mutex::new(HashMap::new()), + journals: Mutex::new(HashMap::new()), managed_run_sessions: Mutex::new(HashMap::new()), auth_errors: Mutex::new(HashMap::new()), sessions: Mutex::new(HashMap::new()), @@ -105,7 +126,7 @@ impl Daemon { }) } - async fn configure_event(&self, env: &mut Envelope) -> anyhow::Result<()> { + async fn configure_event(&self, env: &mut Envelope) -> anyhow::Result { let Some(provider) = &self.auth_provider else { anyhow::bail!("daemon host has no Braintrust auth provider"); }; @@ -118,18 +139,14 @@ impl Daemon { "session route is missing its trace destination; select a project or destination during `bt trace setup` or `bt trace run`" ); } + let key = DeliveryKey::new(&env.session_id, &route)?; let (selection, reason, expected_profile) = { let states = self.session_auth.lock().await; - match states.get(&env.session_id) { + match states.get(&key) { Some(state) => { - if !state.route.same_route(&route) { - anyhow::bail!( - "session route changed after initialization; start a new agent session to change profile, organization, or destination" - ); - } if !lease_is_expiring(&state.lease) { env.config = Some(state.route.with_auth(state.lease.auth.clone())); - return Ok(()); + return Ok(key); } ( AuthSelection { @@ -150,7 +167,7 @@ impl Daemon { env.source ); self.auth_errors.lock().unwrap().insert( - env.session_id.clone(), + key.clone(), (env.source.clone(), message.clone()), ); anyhow::anyhow!(message) @@ -174,20 +191,39 @@ impl Daemon { } } + let resolved_org = lease + .auth + .org_name + .as_deref() + .filter(|org| !org.trim().is_empty()) + .ok_or_else(|| { + let message = format!( + "selected Braintrust profile {:?} did not resolve an organization; pass --org or select an organization during setup", + lease.profile + ); + self.auth_errors.lock().unwrap().insert( + key.clone(), + (env.source.clone(), message.clone()), + ); + anyhow::anyhow!(message) + })? + .to_string(); + let _ = resolved_org; + env.config = Some(route.with_auth(lease.auth.clone())); self.session_auth .lock() .await - .insert(env.session_id.clone(), SessionAuthState { route, lease }); - self.auth_errors.lock().unwrap().remove(&env.session_id); - Ok(()) + .insert(key.clone(), SessionAuthState { route, lease }); + self.auth_errors.lock().unwrap().remove(&key); + Ok(key) } - async fn refresh_session_before_flush(&self, session_id: &str) -> anyhow::Result<()> { + async fn refresh_session_before_flush(&self, key: &DeliveryKey) -> anyhow::Result<()> { let Some(provider) = &self.auth_provider else { return Ok(()); }; - let Some(state) = self.session_auth.lock().await.get(session_id).cloned() else { + let Some(state) = self.session_auth.lock().await.get(key).cloned() else { return Ok(()); }; if !lease_is_expiring(&state.lease) { @@ -205,13 +241,13 @@ impl Daemon { } let config = state.route.with_auth(lease.auth.clone()); self.session_auth.lock().await.insert( - session_id.to_string(), + key.clone(), SessionAuthState { route: state.route, lease, }, ); - let session = { self.sessions.lock().unwrap().get(session_id).cloned() }; + let session = { self.sessions.lock().unwrap().get(key).cloned() }; if let Some(session) = session { session.configure(config).await?; } @@ -222,52 +258,103 @@ impl Daemon { *self.last_activity.lock().unwrap() = Instant::now(); } - async fn session_for(&self, env: &Envelope) -> anyhow::Result> { + async fn replay_for( + &self, + session_id: &str, + route: &SessionRoute, + ) -> anyhow::Result> { + match journal::read_journal(&journal::journal_path(&self.data_dir, session_id)).await { + Ok(entries) => Ok(entries + .into_iter() + .filter(|entry| { + entry + .route + .as_ref() + .is_some_and(|candidate| candidate.same_route(route)) + }) + .map(journal::envelope_from_redacted) + .collect()), + Err(error) + if error + .downcast_ref::() + .is_some_and(|io| io.kind() == std::io::ErrorKind::NotFound) => + { + Ok(Vec::new()) + } + Err(error) => Err(error), + } + } + + async fn session_for(&self, env: &Envelope, key: &DeliveryKey) -> anyhow::Result> { { let map = self.sessions.lock().unwrap(); - if let Some(s) = map.get(&env.session_id) { - return Ok(s.clone()); + if let Some(session) = map.get(key) { + return Ok(session.clone()); } } // Open the journal outside the lock (async I/O), then insert under it, // resolving a race where two connections create the same session. - let replay = - match journal::read_journal(&journal::journal_path(&self.data_dir, &env.session_id)) - .await - { - Ok(entries) => entries - .into_iter() - .map(journal::envelope_from_redacted) - .collect(), - Err(e) - if e.downcast_ref::() - .is_some_and(|io| io.kind() == std::io::ErrorKind::NotFound) => - { - Vec::new() - } - Err(e) => { - tracing::warn!(session_id = %env.session_id, "journal replay skipped: {e}"); - Vec::new() - } - }; - let journal = JournalWriter::open(&self.data_dir, &env.session_id).await?; + let route = env + .route + .as_ref() + .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; + let replay = match self.replay_for(&env.session_id, route).await { + Ok(replay) => replay, + Err(error) => { + tracing::warn!(session_id = %env.session_id, "journal replay skipped: {error}"); + Vec::new() + } + }; + let config = env.config.clone().ok_or_else(|| { + anyhow::anyhow!("resolved route is missing its session configuration") + })?; let mut map = self.sessions.lock().unwrap(); - if let Some(s) = map.get(&env.session_id) { + if let Some(s) = map.get(key) { return Ok(s.clone()); } let session = Session::spawn( env.session_id.clone(), env.source.clone(), env.plugin_version.clone(), - journal, replay, + config, self.translators.clone(), self.sink_factory.clone(), ); - map.insert(env.session_id.clone(), session.clone()); + map.insert(key.clone(), session.clone()); Ok(session) } + async fn append_to_journal(&self, env: &mut Envelope) -> anyhow::Result<()> { + hydrate_transcript_snapshot(env).await; + let existing = { self.journals.lock().unwrap().get(&env.session_id).cloned() }; + let writer = match existing { + Some(writer) => writer, + None => { + let writer = Arc::new(tokio::sync::Mutex::new( + JournalWriter::open(&self.data_dir, &env.session_id).await?, + )); + self.journals + .lock() + .unwrap() + .entry(env.session_id.clone()) + .or_insert_with(|| writer.clone()) + .clone() + } + }; + let result = writer.lock().await.append(env).await; + result + } + + fn session_lock(&self, session_id: &str) -> Arc> { + self.session_locks + .lock() + .unwrap() + .entry(session_id.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + fn total_queued(&self) -> u64 { self.sessions .lock() @@ -277,17 +364,17 @@ impl Daemon { .sum() } - fn record_managed_run_session(&self, managed_run_id: &str, session_id: &str) { + fn record_managed_run_session(&self, managed_run_id: &str, key: &DeliveryKey) { self.managed_run_sessions .lock() .unwrap() .entry(managed_run_id.to_string()) .or_default() - .insert(session_id.to_string()); + .insert(key.clone()); } async fn flush_managed_run(&self, params: ManagedRunFlushParams) -> FlushResult { - let session_ids = self + let delivery_keys = self .managed_run_sessions .lock() .unwrap() @@ -297,19 +384,20 @@ impl Daemon { let mut result = FlushResult { flushed: true, pending: 0, + accepted_sessions: delivery_keys.len() as u64, }; - for session_id in session_ids { - if let Err(error) = self.refresh_session_before_flush(&session_id).await { + for key in delivery_keys { + if let Err(error) = self.refresh_session_before_flush(&key).await { tracing::warn!( managed_run_id = %params.managed_run_id, - session_id, + session_id = %key.session_id, %error, "managed run session auth refresh failed" ); result.flushed = false; continue; } - let session = { self.sessions.lock().unwrap().get(&session_id).cloned() }; + let session = { self.sessions.lock().unwrap().get(&key).cloned() }; if let Some(session) = session { let (flushed, pending) = session .flush(Duration::from_millis(params.timeout_ms)) @@ -497,33 +585,39 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str let managed_run_id = env.managed_run_id.clone(); tracing::info!(source, event, session_id, "event received"); daemon.touch(); + let session_lock = daemon.session_lock(&session_id); + let _session_guard = session_lock.lock().await; let result = async { - daemon + let delivery_key = daemon .configure_event(&mut env) .await .map_err(|error| format!("session auth failed: {error}"))?; let session = daemon - .session_for(&env) + .session_for(&env, &delivery_key) .await .map_err(|error| format!("session init failed: {error}"))?; - session - .append_and_enqueue(env) + daemon + .append_to_journal(&mut env) .await - .map_err(|error| format!("enqueue failed: {error}")) + .map_err(|error| format!("journal failed: {error}"))?; + session + .enqueue(env) + .map_err(|error| format!("enqueue failed: {error}"))?; + Ok(delivery_key) } .await; match &result { - Ok(()) => { + Ok(delivery_key) => { if let Some(managed_run_id) = managed_run_id { - daemon.record_managed_run_session(&managed_run_id, &session_id); + daemon.record_managed_run_session(&managed_run_id, delivery_key); } tracing::info!(source, event, session_id, "event accepted") } Err(error) => tracing::warn!(source, event, session_id, error, "event rejected"), } - result + result.map(|_| ()) } async fn handle_request(daemon: &Arc, req: Request) -> Response { @@ -580,23 +674,43 @@ async fn handle_request(daemon: &Arc, req: Request) -> Response { } method::SESSION_FLUSH => { let p = parse!(FlushParams); - if let Err(error) = daemon.refresh_session_before_flush(&p.session_id).await { - return Response::err( - id, - RpcError::new( - error_code::INTERNAL, - format!("session auth refresh failed: {error}"), - ), - ); + let delivery_keys: Vec<_> = daemon + .sessions + .lock() + .unwrap() + .keys() + .filter(|key| key.session_id == p.session_id) + .cloned() + .collect(); + let accepted_sessions = delivery_keys.len() as u64; + let mut flushed = true; + let mut pending = 0u64; + for key in delivery_keys { + if let Err(error) = daemon.refresh_session_before_flush(&key).await { + return Response::err( + id, + RpcError::new( + error_code::INTERNAL, + format!("session auth refresh failed: {error}"), + ), + ); + } + let session = { daemon.sessions.lock().unwrap().get(&key).cloned() }; + if let Some(session) = session { + let (route_flushed, route_pending) = + session.flush(Duration::from_millis(p.timeout_ms)).await; + flushed &= route_flushed; + pending = pending.saturating_add(route_pending); + } } - let session = { daemon.sessions.lock().unwrap().get(&p.session_id).cloned() }; - let (flushed, pending) = match session { - Some(s) => s.flush(Duration::from_millis(p.timeout_ms)).await, - None => (true, 0), - }; Response::ok( id, - serde_json::to_value(FlushResult { flushed, pending }).unwrap(), + serde_json::to_value(FlushResult { + flushed, + pending, + accepted_sessions, + }) + .unwrap(), ) } method::MANAGED_RUN_FLUSH => { @@ -631,26 +745,35 @@ impl Daemon { let map = self.sessions.lock().unwrap(); let mut sessions: Vec<_> = map .iter() - .filter(|(sid, _)| p.session_id.as_ref().is_none_or(|want| *want == **sid)) - .map(|(sid, s)| SessionStatus { - session_id: sid.clone(), + .filter(|(key, _)| { + p.session_id + .as_ref() + .is_none_or(|want| want == &key.session_id) + }) + .map(|(key, s)| SessionStatus { + session_id: key.session_id.clone(), source: s.source.clone(), + route: serde_json::from_str(&key.route).ok(), queued: s.counters.queued.load(Ordering::Relaxed), spans_emitted: s.counters.spans_emitted.load(Ordering::Relaxed), permalink: s.permalink.lock().unwrap().clone(), last_error: s.last_error.lock().unwrap().clone(), }) .collect(); - for (session_id, (source, error)) in self.auth_errors.lock().unwrap().iter() { - if p.session_id.as_ref().is_some_and(|want| want != session_id) { + for (key, (source, error)) in self.auth_errors.lock().unwrap().iter() { + if p.session_id + .as_ref() + .is_some_and(|want| want != &key.session_id) + { continue; } - if map.contains_key(session_id) { + if map.contains_key(key) { continue; } sessions.push(SessionStatus { - session_id: session_id.clone(), + session_id: key.session_id.clone(), source: source.clone(), + route: serde_json::from_str(&key.route).ok(), queued: 0, spans_emitted: 0, permalink: None, diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 474127a..9bc8d61 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -117,6 +117,7 @@ async fn session_config( .services .resolve_auth(&route.auth, AuthResolveReason::Initial) .await?; + require_resolved_org(route, &lease)?; Ok(SessionConfig { auth: lease.auth, destination: route.destination.clone(), @@ -125,6 +126,47 @@ async fn session_config( }) } +fn require_resolved_org(route: &SessionRoute, lease: &AuthLease) -> anyhow::Result { + let org_name = lease + .auth + .org_name + .as_deref() + .filter(|org| !org.trim().is_empty()) + .ok_or_else(|| { + anyhow::anyhow!( + "organization choice required for tracing; pass --org or select an organization during setup" + ) + })?; + if let Some(expected) = route + .auth + .org_name + .as_deref() + .filter(|org| !org.trim().is_empty()) + { + if expected != org_name { + anyhow::bail!( + "selected profile resolved organization {org_name:?}, expected {expected:?}" + ); + } + } + Ok(org_name.to_string()) +} + +async fn resolve_command_route( + host: &TraceHostContext, + destination_required: bool, +) -> anyhow::Result { + let mut route = host.services.resolve_route(destination_required).await?; + let lease = host + .services + .resolve_auth(&route.auth, AuthResolveReason::Initial) + .await?; + let org_name = require_resolved_org(&route, &lease)?; + route.auth.profile = Some(lease.profile); + route.auth.org_name = Some(org_name); + Ok(route) +} + fn print_output(output: TraceCommandOutput, format: OutputFormat) -> anyhow::Result<()> { println!("{}", output.render(format)?); Ok(()) @@ -134,7 +176,7 @@ fn print_output(output: TraceCommandOutput, format: OutputFormat) -> anyhow::Res pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Result<()> { match args.command { TraceCommand::Setup(setup_args) => { - let route = host.services.resolve_route(true).await?; + let route = resolve_command_route(&host, true).await?; print_output(run_setup(setup_args, route)?, host.output_format) } TraceCommand::Daemon(serve_args) => { @@ -178,7 +220,7 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul run_import(import_args, serve_options(&host), Some(config)).await } TraceCommand::Run(run_args) => { - let route = host.services.resolve_route(true).await?; + let route = resolve_command_route(&host, true).await?; let hook_command = child_command(&host.command, "hook"); let status = run_traced(run_args, hook_command, route).await?; if status.success() { @@ -246,6 +288,7 @@ mod tests { route_requests: Mutex>, route_error: Option<&'static str>, auth_error: Option<&'static str>, + resolved_org: Option<&'static str>, } impl RecordingHost { @@ -254,6 +297,16 @@ mod tests { route_requests: Mutex::new(Vec::new()), route_error, auth_error, + resolved_org: Some("test-org"), + } + } + + fn without_org() -> Self { + Self { + route_requests: Mutex::new(Vec::new()), + route_error: None, + auth_error: None, + resolved_org: None, } } } @@ -291,7 +344,7 @@ mod tests { token: "secret".into(), api_url: None, app_url: None, - org_name: None, + org_name: self.resolved_org.map(str::to_string), org_id: None, }, expires_at_ms: None, @@ -332,6 +385,25 @@ mod tests { } } + #[tokio::test] + async fn command_routes_persist_the_resolved_profile_and_organization() { + let services = Arc::new(RecordingHost::new(None, None)); + let route = resolve_command_route(&test_host(services), true) + .await + .unwrap(); + assert_eq!(route.auth.profile.as_deref(), Some("test")); + assert_eq!(route.auth.org_name.as_deref(), Some("test-org")); + } + + #[tokio::test] + async fn command_routes_reject_an_unresolved_organization() { + let services = Arc::new(RecordingHost::without_org()); + let error = resolve_command_route(&test_host(services), true) + .await + .unwrap_err(); + assert!(error.to_string().contains("organization choice required")); + } + #[tokio::test] async fn import_only_requires_a_default_destination_without_an_override() { for (destination, required) in [ diff --git a/bt-daemon/src/wire/methods.rs b/bt-daemon/src/wire/methods.rs index 9ca17cc..143b1ea 100644 --- a/bt-daemon/src/wire/methods.rs +++ b/bt-daemon/src/wire/methods.rs @@ -68,6 +68,8 @@ fn default_flush_timeout_ms() -> u64 { pub struct FlushResult { pub flushed: bool, pub pending: u64, + #[serde(default)] + pub accepted_sessions: u64, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -88,6 +90,8 @@ pub struct StatusResult { pub struct SessionStatus { pub session_id: String, pub source: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub route: Option, pub queued: u64, pub spans_emitted: u64, #[serde(default, skip_serializing_if = "Option::is_none")] From 3d290e67b47ac8daec358297a65e6376b2639113 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 14 Aug 2026 05:48:51 +0800 Subject: [PATCH 2/3] Persist managed-run acceptance and allow interactive auth resolution Managed-run flush accounting now survives a daemon restart or idle-exit by recording each accepted {session_id, route} pair alongside the journal, since the in-memory map alone was lost across daemon generations. TraceHostServices::resolve_route now takes RouteRequirements{destination_required, interactive_auth} so a host can tell hooks (never prompt) apart from setup/run/import (safe to ask for login or list organizations/projects when no default is resolved). Adds a regression test for concurrent `bt trace import` of the same session to different destinations, and corrects docs/protocol.md, which still described one route per session. Co-Authored-By: Claude Sonnet 5 --- bt-daemon/docs/protocol.md | 69 ++++++--- bt-daemon/src/dispatch.rs | 2 +- bt-daemon/src/journal.rs | 92 +++++++++++- bt-daemon/src/lib.rs | 14 +- bt-daemon/src/server.rs | 47 +++++- bt-daemon/src/trace_runtime.rs | 89 ++++++++--- bt-daemon/tests/pipeline.rs | 262 +++++++++++++++++++++++++++++++-- bt-daemon/tests/replay.rs | 39 +++++ 8 files changed, 550 insertions(+), 64 deletions(-) diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index 937a091..6faf8e2 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -111,7 +111,9 @@ those are handled asynchronously and surfaced via `status.get`. ### `session.flush` (request) -Block until the session's spans are delivered, or `timeout_ms` elapses. +Block until every route's spans for the session are delivered, or `timeout_ms` +elapses. A session_id may have opened more than one route (see "Multiple +routes per session" below); this flushes all of them. Params: ```json @@ -119,10 +121,12 @@ Params: ``` Result: ```json -{ "flushed": true, "pending": 0 } +{ "flushed": true, "pending": 0, "accepted_sessions": 1 } ``` `flushed: false` with `pending > 0` means the timeout was hit with work -outstanding. Used by session-end hooks and flush-on-turn-end mode. +outstanding across one or more routes. `accepted_sessions` counts how many +independent routes this session_id has open. Used by session-end hooks and +flush-on-turn-end mode. ### `managed_run.flush` (request) @@ -147,6 +151,10 @@ Result: { "session_id": "…", "source": "codex", + "route": { + "auth": { "profile": "work", "org_name": "acme" }, + "destination": { "type": "project_logs", "project_name": "codex" } + }, "queued": 0, "spans_emitted": 42, "permalink": "https://www.braintrust.dev/app/…", @@ -155,7 +163,10 @@ Result: ] } ``` -Powers a `status` CLI and pi's trace-link widget. +`sessions` lists one entry **per route**, not per session_id: a session_id +reporting to two destinations appears twice, each entry carrying its own +`route`, counters, and permalink. Powers a `status` CLI and pi's trace-link +widget. ### `daemon.shutdown` (request) @@ -194,9 +205,11 @@ Field notes: - **`source`** selects the daemon-side translator. `debug` is a built-in pass-through translator used by the prototype and tests. -- **`session_id`** is the per-session queue + state key. The shim extracts it - from the payload (default JSON field `session_id`, overridable with - `--session-id-field`); both Claude Code and Codex use `session_id`. +- **`session_id`** identifies the source agent session. Combined with `route` + it forms the queue + state key (see "Multiple routes per session" below). + The shim extracts it from the payload (default JSON field `session_id`, + overridable with `--session-id-field`); both Claude Code and Codex use + `session_id`. - **`event`** is the agent-native hook name (not normalized). Extracted from the payload (default field `hook_event_name`, overridable with `--event`). - **`ts_ms`** is stamped by the shim **at capture time** (epoch millis), @@ -211,13 +224,22 @@ Field notes: is optional and resolves through `bt`'s default profile when absent; `org_name` optionally constrains organization selection. The daemon resolves the live credential, pins the returned canonical profile for the lifetime - of the session, and refreshes an expiring lease without changing that route. - A route cannot change after a session's first accepted event. `destination` - is required so setup/run must make project or parent selection explicit. - `flush_mode` ∈ `fire_and_forget` | `flush_on_turn_end`. - New front-ends set the typed `destination`: `project_logs` accepts a project - id and/or name, `experiment` accepts an experiment id, and `parent_span` - carries the complete exported `SpanComponents` object. + of that route's pipeline, and refreshes an expiring lease without changing + the route. `destination` is required so setup/run must make project or + parent selection explicit. `flush_mode` ∈ `fire_and_forget` | + `flush_on_turn_end`. New front-ends set the typed `destination`: + `project_logs` accepts a project id and/or name, `experiment` accepts an + experiment id, and `parent_span` carries the complete exported + `SpanComponents` object. +- **Multiple routes per session.** `session_id` plus the exact `route` forms + one independent delivery pipeline: its own auth resolution, translator, + sink, and queue. A session_id is not pinned to a single route — events for + the same session_id but a different route open a second, fully independent + pipeline rather than replacing or rejecting the first. This lets one source + session report concurrently to multiple destinations, including multiple + organizations (e.g. two `bt trace import` runs, or an active hook capture + alongside a concurrent import, targeting different projects or orgs for the + same underlying session). ### Redaction @@ -277,15 +299,24 @@ continue using setup settings, and concurrent managed runs can select distinct profiles, organizations, and destinations while sharing one daemon. - **Journal (WAL).** Every accepted event is appended (auth-redacted) to - `/journal/.ndjson` before/at enqueue. `data_dir` - defaults to `$XDG_STATE_HOME/braintrust/bt-daemon` or + `/journal/.ndjson` before/at enqueue — one journal + file per session_id, shared across every route that session_id has opened. + `data_dir` defaults to `$XDG_STATE_HOME/braintrust/bt-daemon` or `$HOME/.braintrust/state/bt-daemon` on Unix, and `%LOCALAPPDATA%\Braintrust\bt-daemon` on Windows. On restart the daemon - rebuilds a session's unfinished correlation state by applying its journal to - a fresh translator. The resulting rows may be resubmitted to repair delivery + 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 interrupted by a crash, but their deterministic ids target the same backend - rows and must never create duplicate spans. + rows and must never create duplicate spans, and a route never receives + another route's rows. Journals are GC'd after 7 days. +- **Managed-run acceptance records.** Alongside the journal, each accepted + event that carries a `managed_run_id` also appends `{session_id, route}` to + `/managed-runs/.ndjson`. `managed_run.flush` reads + this record when the daemon that accepted the events has since restarted or + idle-exited, so flush accounting for a child process tree survives a daemon + generation change. GC'd after 7 days like the journal. - **Deterministic span ids.** Translators derive span ids as UUIDv5 over stable keys (`session_id`, `turn_id`, `call_id`, …) so a replayed re-emit merges server-side (`_is_merge`) instead of duplicating. diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index b56fd45..7f8bf53 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -4,7 +4,7 @@ //! concurrently. //! //! Ack semantics: `event.log` is acked once the event is journaled and handed -//! to the session's queue (see [`Session::append_and_enqueue`]). Delivery to +//! to the session's queue (see [`Session::enqueue`]). Delivery to //! Braintrust happens later in the actor; a downstream error never fails the //! caller's turn. diff --git a/bt-daemon/src/journal.rs b/bt-daemon/src/journal.rs index a220f2d..69da587 100644 --- a/bt-daemon/src/journal.rs +++ b/bt-daemon/src/journal.rs @@ -4,8 +4,14 @@ //! //! Format: one [`RedactedEnvelope`] JSON value per line in //! `/journal/.ndjson`. +//! +//! Managed-run acceptance records live alongside the journals so a flush can +//! still tell which delivery pipelines a managed child produced after the +//! daemon that accepted them has restarted or idle-exited. -use crate::wire::{BackendAuth, Envelope, RedactedEnvelope}; +use crate::wire::{BackendAuth, Envelope, RedactedEnvelope, SessionRoute}; +use serde::{Deserialize, Serialize}; +use std::collections::HashSet; use std::path::{Path, PathBuf}; use tokio::io::AsyncWriteExt; @@ -29,6 +35,90 @@ pub fn journal_path(data_dir: &Path, session_id: &str) -> PathBuf { journal_dir(data_dir).join(format!("{}.ndjson", sanitize(session_id))) } +pub fn managed_run_dir(data_dir: &Path) -> PathBuf { + data_dir.join("managed-runs") +} + +pub fn managed_run_path(data_dir: &Path, managed_run_id: &str) -> PathBuf { + managed_run_dir(data_dir).join(format!("{}.ndjson", sanitize(managed_run_id))) +} + +/// One delivery pipeline accepted from a managed child process tree. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagedRunKey { + pub session_id: String, + pub route: SessionRoute, +} + +/// Append one accepted delivery pipeline to the managed run's record. +pub async fn append_managed_run_key( + data_dir: &Path, + managed_run_id: &str, + key: &ManagedRunKey, +) -> anyhow::Result<()> { + let dir = managed_run_dir(data_dir); + tokio::fs::create_dir_all(&dir).await?; + let mut line = serde_json::to_vec(key)?; + line.push(b'\n'); + let mut file = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(managed_run_path(data_dir, managed_run_id)) + .await?; + file.write_all(&line).await?; + file.flush().await?; + Ok(()) +} + +/// Read a managed run's accepted delivery pipelines back, deduplicated. +/// A missing record means the run never produced an accepted event. +pub async fn read_managed_run_keys(data_dir: &Path, managed_run_id: &str) -> Vec { + let Ok(data) = tokio::fs::read_to_string(managed_run_path(data_dir, managed_run_id)).await + else { + return Vec::new(); + }; + let mut keys = Vec::new(); + let mut seen = HashSet::new(); + for line in data.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + if seen.insert(line.to_string()) { + if let Ok(key) = serde_json::from_str::(line) { + keys.push(key); + } + } + } + keys +} + +/// Best-effort age-based collection of managed-run records, mirroring journal +/// GC. +pub async fn gc_old_managed_runs(data_dir: &Path, max_age: std::time::Duration) { + let dir = managed_run_dir(data_dir); + let Ok(mut entries) = tokio::fs::read_dir(&dir).await else { + return; + }; + let now = std::time::SystemTime::now(); + while let Ok(Some(entry)) = entries.next_entry().await { + let path = entry.path(); + if path.extension().and_then(|v| v.to_str()) != Some("ndjson") { + continue; + } + let old = entry + .metadata() + .await + .ok() + .and_then(|m| m.modified().ok()) + .and_then(|modified| now.duration_since(modified).ok()) + .is_some_and(|age| age > max_age); + if old { + let _ = tokio::fs::remove_file(&path).await; + } + } +} + /// Append-only journal writer for one session. pub struct JournalWriter { file: tokio::fs::File, diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 5888776..382c720 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -36,7 +36,7 @@ pub use server::{AuthLease, AuthProvider, AuthResolveReason, ServeOptions}; pub use setup::run_setup; pub use sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, Sink, SinkFactory}; pub use trace_command::{SetupAgent, SetupArgs, StopArgs, TraceArgs, TraceCommand}; -pub use trace_runtime::{run_trace, TraceHostContext, TraceHostServices}; +pub use trace_runtime::{run_trace, RouteRequirements, TraceHostContext, TraceHostServices}; pub use translate::{ AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, }; @@ -343,7 +343,9 @@ pub async fn flush_session( } /// Flush every daemon session accepted from one managed child process tree. -/// A missing daemon means the child emitted no accepted trace events. +/// A missing daemon is not itself a failure: the daemon may have idle-exited +/// after draining every session, so acceptance is checked against the +/// persisted managed-run record instead. pub async fn flush_managed_run( managed_run_id: &str, socket: &std::path::Path, @@ -352,11 +354,15 @@ pub async fn flush_managed_run( let stream = match client::connect(socket).await { Ok(stream) => stream, Err(_) => { + let accepted_sessions = + journal::read_managed_run_keys(&paths::data_dir(None), managed_run_id) + .await + .len() as u64; return Ok(wire::FlushResult { flushed: true, pending: 0, - accepted_sessions: 0, - }) + accepted_sessions, + }); } }; let mut conn = client::Conn::new(stream); diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index 2a940b2..b3d4a25 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -364,23 +364,56 @@ impl Daemon { .sum() } - fn record_managed_run_session(&self, managed_run_id: &str, key: &DeliveryKey) { + fn record_managed_run_session(&self, managed_run_id: &str, key: &DeliveryKey) -> bool { self.managed_run_sessions .lock() .unwrap() .entry(managed_run_id.to_string()) .or_default() - .insert(key.clone()); + .insert(key.clone()) + } + + async fn persist_managed_run_session( + &self, + managed_run_id: &str, + key: &DeliveryKey, + route: &SessionRoute, + ) { + if !self.record_managed_run_session(managed_run_id, key) { + return; + } + let record = journal::ManagedRunKey { + session_id: key.session_id.clone(), + route: route.clone(), + }; + if let Err(error) = journal::append_managed_run_key(&self.data_dir, managed_run_id, &record) + .await + { + tracing::warn!( + managed_run_id, + session_id = %key.session_id, + %error, + "failed to record managed run delivery pipeline" + ); + } } async fn flush_managed_run(&self, params: ManagedRunFlushParams) -> FlushResult { - let delivery_keys = self + let mut delivery_keys: HashSet = self .managed_run_sessions .lock() .unwrap() .get(¶ms.managed_run_id) .cloned() .unwrap_or_default(); + // A daemon restart or idle exit loses the in-memory mapping; the + // persisted record keeps flush accounting accurate for runs whose + // events were accepted by an earlier daemon generation. + for record in journal::read_managed_run_keys(&self.data_dir, ¶ms.managed_run_id).await { + if let Ok(key) = DeliveryKey::new(&record.session_id, &record.route) { + delivery_keys.insert(key); + } + } let mut result = FlushResult { flushed: true, pending: 0, @@ -455,6 +488,7 @@ pub async fn run(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { let daemon = Daemon::new(opts, data_dir); journal::gc_old_journals(&daemon.data_dir, Duration::from_secs(7 * 24 * 60 * 60)).await; + journal::gc_old_managed_runs(&daemon.data_dir, Duration::from_secs(7 * 24 * 60 * 60)).await; let idle_timeout = Duration::from_secs(args.idle_timeout_secs); spawn_idle_watchdog(daemon.clone(), idle_timeout); @@ -583,6 +617,7 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str let event = env.event.clone(); let session_id = env.session_id.clone(); let managed_run_id = env.managed_run_id.clone(); + let route = env.route.clone(); tracing::info!(source, event, session_id, "event received"); daemon.touch(); let session_lock = daemon.session_lock(&session_id); @@ -610,8 +645,10 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str match &result { Ok(delivery_key) => { - if let Some(managed_run_id) = managed_run_id { - daemon.record_managed_run_session(&managed_run_id, delivery_key); + if let (Some(managed_run_id), Some(route)) = (managed_run_id, route) { + daemon + .persist_managed_run_session(&managed_run_id, delivery_key, &route) + .await; } tracing::info!(source, event, session_id, "event accepted") } diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 9bc8d61..06c4a9a 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -17,6 +17,19 @@ use async_trait::async_trait; use std::ffi::OsString; use std::sync::Arc; +/// What a command still needs the host to resolve before tracing can start. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct RouteRequirements { + /// The command cannot proceed without a trace destination. Hosts should + /// resolve one from flags, stored defaults, or an interactive selection. + pub destination_required: bool, + /// The command may prompt to complete missing auth selections: asking the + /// user to log in when no profile exists, or listing available + /// organizations when the profile has no default. Hooks leave this false + /// so an agent's turn never blocks on interactive input. + pub interactive_auth: bool, +} + /// Host-owned services used by the integration runtime. /// /// Implementations resolve Braintrust profiles and destination choices but do @@ -24,9 +37,9 @@ use std::sync::Arc; #[async_trait] pub trait TraceHostServices: Send + Sync { /// Resolve the non-secret route selected by the host. Commands such as - /// setup and managed run require a destination; hooks may use the current - /// selection without prompting. - async fn resolve_route(&self, destination_required: bool) -> anyhow::Result; + /// setup and managed run require a destination and allow interactive auth + /// resolution; hooks may use the current selection without prompting. + async fn resolve_route(&self, requirements: RouteRequirements) -> anyhow::Result; /// Resolve a Braintrust credential lease without exposing credentials to /// plugins, settings files, journals, or command arguments. @@ -154,9 +167,9 @@ fn require_resolved_org(route: &SessionRoute, lease: &AuthLease) -> anyhow::Resu async fn resolve_command_route( host: &TraceHostContext, - destination_required: bool, + requirements: RouteRequirements, ) -> anyhow::Result { - let mut route = host.services.resolve_route(destination_required).await?; + let mut route = host.services.resolve_route(requirements).await?; let lease = host .services .resolve_auth(&route.auth, AuthResolveReason::Initial) @@ -176,7 +189,14 @@ fn print_output(output: TraceCommandOutput, format: OutputFormat) -> anyhow::Res pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Result<()> { match args.command { TraceCommand::Setup(setup_args) => { - let route = resolve_command_route(&host, true).await?; + let route = resolve_command_route( + &host, + RouteRequirements { + destination_required: true, + interactive_auth: true, + }, + ) + .await?; print_output(run_setup(setup_args, route)?, host.output_format) } TraceCommand::Daemon(serve_args) => { @@ -186,7 +206,7 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul TraceCommand::Hook(hook_args) => { // A persistent hook must never fail the coding agent's turn. let result = async { - let route = host.services.resolve_route(false).await?; + let route = host.services.resolve_route(RouteRequirements::default()).await?; run_hook(hook_args, route, host_info(&host)).await } .await; @@ -214,13 +234,24 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul TraceCommand::Import(import_args) => { let route = host .services - .resolve_route(import_args.destination.is_none() && import_args.parent.is_none()) + .resolve_route(RouteRequirements { + destination_required: import_args.destination.is_none() + && import_args.parent.is_none(), + interactive_auth: true, + }) .await?; let config = session_config(&host, &route).await?; run_import(import_args, serve_options(&host), Some(config)).await } TraceCommand::Run(run_args) => { - let route = resolve_command_route(&host, true).await?; + let route = resolve_command_route( + &host, + RouteRequirements { + destination_required: true, + interactive_auth: true, + }, + ) + .await?; let hook_command = child_command(&host.command, "hook"); let status = run_traced(run_args, hook_command, route).await?; if status.success() { @@ -271,7 +302,7 @@ mod tests { #[async_trait] impl TraceHostServices for PanicHost { - async fn resolve_route(&self, _: bool) -> anyhow::Result { + async fn resolve_route(&self, _: RouteRequirements) -> anyhow::Result { panic!("host service should not be called") } @@ -285,7 +316,7 @@ mod tests { } struct RecordingHost { - route_requests: Mutex>, + route_requests: Mutex>, route_error: Option<&'static str>, auth_error: Option<&'static str>, resolved_org: Option<&'static str>, @@ -313,11 +344,11 @@ mod tests { #[async_trait] impl TraceHostServices for RecordingHost { - async fn resolve_route(&self, destination_required: bool) -> anyhow::Result { - self.route_requests - .lock() - .unwrap() - .push(destination_required); + async fn resolve_route( + &self, + requirements: RouteRequirements, + ) -> anyhow::Result { + self.route_requests.lock().unwrap().push(requirements); if let Some(error) = self.route_error { anyhow::bail!(error); } @@ -365,6 +396,11 @@ mod tests { } } + const COMMAND_REQUIREMENTS: RouteRequirements = RouteRequirements { + destination_required: true, + interactive_auth: true, + }; + #[tokio::test] async fn setup_and_run_require_a_host_resolved_destination() { for command in [ @@ -381,14 +417,14 @@ mod tests { .await .unwrap_err(); assert_eq!(error.to_string(), "no destination"); - assert_eq!(*services.route_requests.lock().unwrap(), [true]); + assert_eq!(*services.route_requests.lock().unwrap(), [COMMAND_REQUIREMENTS]); } } #[tokio::test] async fn command_routes_persist_the_resolved_profile_and_organization() { let services = Arc::new(RecordingHost::new(None, None)); - let route = resolve_command_route(&test_host(services), true) + let route = resolve_command_route(&test_host(services), COMMAND_REQUIREMENTS) .await .unwrap(); assert_eq!(route.auth.profile.as_deref(), Some("test")); @@ -398,7 +434,7 @@ mod tests { #[tokio::test] async fn command_routes_reject_an_unresolved_organization() { let services = Arc::new(RecordingHost::without_org()); - let error = resolve_command_route(&test_host(services), true) + let error = resolve_command_route(&test_host(services), COMMAND_REQUIREMENTS) .await .unwrap_err(); assert!(error.to_string().contains("organization choice required")); @@ -406,7 +442,7 @@ mod tests { #[tokio::test] async fn import_only_requires_a_default_destination_without_an_override() { - for (destination, required) in [ + for (destination, destination_required) in [ (None, true), ( Some(TraceDestination::ProjectLogs { @@ -433,7 +469,13 @@ mod tests { .await .unwrap_err(); assert_eq!(error.to_string(), "stop before lookup"); - assert_eq!(*services.route_requests.lock().unwrap(), [required]); + assert_eq!( + *services.route_requests.lock().unwrap(), + [RouteRequirements { + destination_required, + interactive_auth: true, + }] + ); } } @@ -461,7 +503,10 @@ mod tests { ) .await .unwrap(); - assert_eq!(*services.route_requests.lock().unwrap(), [false]); + assert_eq!( + *services.route_requests.lock().unwrap(), + [RouteRequirements::default()] + ); } #[tokio::test] diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index dfa9d46..243db31 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -11,7 +11,7 @@ use bt_daemon::{ }; #[cfg(all(feature = "cli", unix))] use bt_daemon::{run_traced, RunArgs, RunHookCommand, RunSource}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::ffi::OsString; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -60,6 +60,60 @@ impl Sink for TrackingSink { } } +/// One observed delivery pipeline: the resolved route configuration plus the +/// span activity delivered through it. +#[derive(Default)] +struct RouteSinkRecord { + org: Mutex>, + destination: Mutex>, + emitted: std::sync::atomic::AtomicU64, + flushes: std::sync::atomic::AtomicU64, +} + +#[derive(Default)] +struct RouteRecordingSinkFactory { + sinks: Mutex>>, +} + +impl SinkFactory for RouteRecordingSinkFactory { + fn create( + &self, + _session_id: &str, + _source: &str, + _plugin_version: Option<&str>, + ) -> anyhow::Result> { + let record = Arc::new(RouteSinkRecord::default()); + self.sinks.lock().unwrap().push(record.clone()); + Ok(Box::new(RouteRecordingSink { record })) + } +} + +struct RouteRecordingSink { + record: Arc, +} + +#[async_trait] +impl Sink for RouteRecordingSink { + fn configure(&mut self, config: &SessionConfig) { + *self.record.org.lock().unwrap() = config.auth.org_name.clone(); + *self.record.destination.lock().unwrap() = config.destination.clone(); + } + + async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { + self.record + .emitted + .fetch_add(ops.len() as u64, std::sync::atomic::Ordering::Relaxed); + Ok(ops.len() as u64) + } + + async fn flush(&mut self) -> anyhow::Result<()> { + self.record + .flushes + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Ok(()) + } +} + fn dummy_host() -> HostInfo { // The daemon is started in-process, so the client never spawns; serve_argv // is unused but must be non-empty. @@ -136,7 +190,12 @@ impl AuthProvider for TestAuthProvider { token: format!("secret-{profile}-{call_index}"), api_url: Some(format!("https://{profile}.example.test")), app_url: None, - org_name: selection.org_name.clone(), + // Model a profile with a default organization for selections + // that do not constrain one. + org_name: selection + .org_name + .clone() + .or_else(|| Some(format!("{profile}-org"))), org_id: Some(format!("org-{profile}")), }, expires_at_ms: (self.first_lease_expired && call_index == 1).then_some(0), @@ -181,6 +240,46 @@ fn test_endpoint(tmp: &Path) -> PathBuf { } } +async fn start_daemon_with( + provider: Arc, + sink_factory: Arc, +) -> ( + PathBuf, + PathBuf, + tokio::task::JoinHandle<()>, + tempfile::TempDir, +) { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let handle = start_daemon_at_with(data_dir.clone(), socket.clone(), provider, sink_factory).await; + (data_dir, socket, handle, tmp) +} + +async fn start_daemon_at_with( + data_dir: PathBuf, + socket: PathBuf, + provider: Arc, + sink_factory: Arc, +) -> tokio::task::JoinHandle<()> { + let args = ServeArgs { + socket: Some(socket.clone()), + data_dir: Some(data_dir), + idle_timeout_secs: 0, + }; + let opts = ServeOptions { + version: "test".to_string(), + translators: Arc::new(Registry::default_agents()), + sink_factory, + auth_provider: Some(provider), + }; + let handle = tokio::spawn(async move { + let _ = run_serve(args, opts).await; + }); + wait_for(&socket).await; + handle +} + async fn wait_for(endpoint: &Path) { for _ in 0..200 { if let Ok(Some(_)) = run_status(StatusArgs { @@ -393,35 +492,173 @@ async fn expiring_profile_lease_is_refreshed_for_the_pinned_profile() { } #[tokio::test] -async fn active_session_rejects_route_changes() { +async fn one_session_reports_to_multiple_routes_and_orgs_concurrently() { let provider = Arc::new(TestAuthProvider { calls: Mutex::new(Vec::new()), fail: false, first_lease_expired: false, }); - let (_data_dir, socket, handle, _tmp) = start_routed_daemon(provider).await; + let recording = Arc::new(RouteRecordingSinkFactory::default()); + let (_data_dir, socket, handle, _tmp) = + start_daemon_with(provider.clone(), recording.clone()).await; let host = dummy_host(); forward_envelope( - &routed_envelope("pinned", "work", "work-org", "SessionStart"), + &routed_envelope("shared", "work", "work-org", "SessionStart"), &socket, &host, false, ) .await .unwrap(); - let error = forward_envelope( - &routed_envelope("pinned", "personal", "personal-org", "Stop"), + forward_envelope( + &routed_envelope("shared", "personal", "personal-org", "Stop"), &socket, &host, false, ) .await - .unwrap_err(); - assert!(error.to_string().contains("session route changed")); + .unwrap(); + + let flushed = flush_session("shared", &socket, 5000).await.unwrap(); + assert!(flushed.flushed, "flush did not complete: {flushed:?}"); + assert_eq!(flushed.pending, 0); + assert_eq!( + flushed.accepted_sessions, 2, + "each route is an independent delivery pipeline" + ); + + let calls = provider.calls.lock().unwrap().clone(); + assert_eq!(calls.len(), 2, "each route resolves its own credentials"); + assert_eq!(calls[0].0.org_name.as_deref(), Some("work-org")); + assert_eq!(calls[1].0.org_name.as_deref(), Some("personal-org")); + + let sinks = recording.sinks.lock().unwrap().clone(); + assert_eq!(sinks.len(), 2, "one sink per route, not per session"); + let orgs: HashSet<_> = sinks + .iter() + .map(|sink| sink.org.lock().unwrap().clone().unwrap()) + .collect(); + assert_eq!( + orgs, + HashSet::from(["work-org".to_string(), "personal-org".to_string()]) + ); + for sink in &sinks { + assert!( + sink.emitted.load(std::sync::atomic::Ordering::Relaxed) > 0, + "every destination must receive spans" + ); + assert_eq!(sink.flushes.load(std::sync::atomic::Ordering::Relaxed), 1); + } + let status = run_status(StatusArgs { + socket: Some(socket.clone()), + session_id: Some("shared".into()), + }) + .await + .unwrap() + .unwrap(); + assert_eq!(status.sessions.len(), 2); + let status_orgs: HashSet<_> = status + .sessions + .iter() + .map(|session| { + session + .route + .as_ref() + .and_then(|route| route.auth.org_name.clone()) + .unwrap() + }) + .collect(); + assert_eq!( + status_orgs, + HashSet::from(["work-org".to_string(), "personal-org".to_string()]) + ); + + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[tokio::test] +async fn restarted_daemon_replays_each_route_only_to_its_own_destination() { + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + }); + let recording = Arc::new(RouteRecordingSinkFactory::default()); + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let handle = start_daemon_at_with(data_dir.clone(), socket.clone(), provider, recording).await; + let host = dummy_host(); + + forward_envelope( + &routed_envelope("rerouted", "work", "work-org", "SessionStart"), + &socket, + &host, + false, + ) + .await + .unwrap(); + forward_envelope( + &routed_envelope("rerouted", "personal", "personal-org", "PostToolUse"), + &socket, + &host, + false, + ) + .await + .unwrap(); + flush_session("rerouted", &socket, 5000).await.unwrap(); shutdown(&socket).await; handle.await.unwrap(); + + // A fresh daemon generation must rebuild each route's pipeline from the + // shared journal without cross-delivering one route's events to the other + // route's destination. + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + }); + let recording = Arc::new(RouteRecordingSinkFactory::default()); + let second = start_daemon_at_with( + data_dir.clone(), + socket.clone(), + provider, + recording.clone(), + ) + .await; + forward_envelope( + &routed_envelope("rerouted", "work", "work-org", "Stop"), + &socket, + &host, + false, + ) + .await + .unwrap(); + let flushed = flush_session("rerouted", &socket, 5000).await.unwrap(); + assert!(flushed.flushed, "flush did not complete: {flushed:?}"); + assert_eq!(flushed.accepted_sessions, 1); + + let sinks = recording.sinks.lock().unwrap().clone(); + assert_eq!( + sinks.len(), + 1, + "only the route with new events is rebuilt eagerly" + ); + assert_eq!( + sinks[0].org.lock().unwrap().clone().as_deref(), + Some("work-org") + ); + assert_eq!( + sinks[0].emitted.load(std::sync::atomic::Ordering::Relaxed), + 3, + "replayed SessionStart (root + span) plus the new Stop span" + ); + + shutdown(&socket).await; + second.await.unwrap(); } #[tokio::test] @@ -647,6 +884,7 @@ async fn spawn_on_demand_runs_the_real_standalone_daemon() { let tmp = tempfile::tempdir().unwrap(); let data_dir = tmp.path().join("spawned-data"); let socket = test_endpoint(tmp.path()); + let _org = EnvVarGuard::set("BRAINTRUST_ORG_NAME", "spawn-org"); let host = HostInfo { serve_argv: vec![ OsString::from(env!("CARGO_BIN_EXE_bt-daemon")), @@ -798,13 +1036,13 @@ async fn managed_run_flush_is_scoped_to_its_accepted_sessions() { handle.await.unwrap(); } -#[cfg(all(feature = "cli", unix))] +#[cfg(feature = "cli")] struct EnvVarGuard { key: &'static str, previous: Option, } -#[cfg(all(feature = "cli", unix))] +#[cfg(feature = "cli")] impl EnvVarGuard { fn set(key: &'static str, value: impl AsRef) -> Self { let previous = std::env::var_os(key); @@ -813,7 +1051,7 @@ impl EnvVarGuard { } } -#[cfg(all(feature = "cli", unix))] +#[cfg(feature = "cli")] impl Drop for EnvVarGuard { fn drop(&mut self) { match &self.previous { diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 5ac4c02..fcab8f1 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -83,6 +83,45 @@ async fn imports_native_codex_rollout_through_codex_translator() { assert!(turn_ids.contains(&"turn-2")); } +/// Two `bt trace import` invocations for the same session id, each headed to +/// a different destination, must not interfere: every field the daemon's +/// per-route dispatch protects (accepted events, span counts) belongs to an +/// in-process `ImportProcessor` with no shared state, so running them +/// concurrently must produce two complete, independent copies rather than one +/// import starving or corrupting the other. +#[tokio::test] +async fn concurrent_imports_of_the_same_session_reach_independent_destinations() { + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("rollout.jsonl"); + write_jsonl( + &transcript, + &[ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"codex-concurrent","cwd":"/tmp/demo","cli_version":"1.2.3"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"turn_context","payload":{"model":"gpt-test"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}}), + json!({"timestamp":"2026-01-01T00:00:04Z","type":"event_msg","payload":{"type":"user_message","message":"list files"}}), + json!({"timestamp":"2026-01-01T00:00:05Z","type":"response_item","payload":{"type":"function_call","call_id":"call-1","name":"shell","arguments":"{\"command\":\"ls\"}"}}), + json!({"timestamp":"2026-01-01T00:00:06Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call-1","output":"README.md"}}), + json!({"timestamp":"2026-01-01T00:00:07Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"Done"}}), + ], + ); + + let project_a = tmp.path().join("spans-a"); + let project_b = tmp.path().join("spans-b"); + let (result_a, result_b) = tokio::join!( + import_transcript(&transcript, ImportSource::Codex, options(&project_a), None, false), + import_transcript(&transcript, ImportSource::Codex, options(&project_b), None, false), + ); + result_a.unwrap(); + result_b.unwrap(); + + for project in [&project_a, &project_b] { + let rows = rows(&project.join("codex-concurrent.ndjson")); + assert_eq!(inserted(&rows, "task"), 2, "session and one turn"); + assert_eq!(inserted(&rows, "tool"), 1); + } +} + #[tokio::test] async fn imports_native_claude_transcript_with_multiple_turns_and_tools() { let tmp = tempfile::tempdir().unwrap(); From 427d61b87cec51d1b902edf9c56175dfd384fde8 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 14 Aug 2026 06:44:15 +0800 Subject: [PATCH 3/3] Fix CI: don't require a resolved org for every hook event configure_event() was requiring a non-empty resolved org for every accepted event, not just for setup/run/import. That's stricter than before "Improve hook routing handling" (which introduced the check) and broke ordinary hook delivery in any environment where the auth provider doesn't default an org (e.g. the standalone bt-daemon binary's EnvironmentAuthProvider, used by the agent_integration e2e suite), causing every event to be rejected and no traces to be captured. Interactive org resolution is already enforced in the right place for setup/run/import: resolve_command_route()/require_resolved_org in trace_runtime.rs, gated by RouteRequirements. Also run cargo fmt, which CI checks and this branch's prior commit had drifted from. Co-Authored-By: Claude Sonnet 5 --- bt-daemon/src/server.rs | 23 ++--------------------- bt-daemon/src/trace_runtime.rs | 10 ++++++++-- bt-daemon/tests/pipeline.rs | 3 ++- bt-daemon/tests/replay.rs | 16 ++++++++++++++-- 4 files changed, 26 insertions(+), 26 deletions(-) diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index b3d4a25..fc697c4 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -191,25 +191,6 @@ impl Daemon { } } - let resolved_org = lease - .auth - .org_name - .as_deref() - .filter(|org| !org.trim().is_empty()) - .ok_or_else(|| { - let message = format!( - "selected Braintrust profile {:?} did not resolve an organization; pass --org or select an organization during setup", - lease.profile - ); - self.auth_errors.lock().unwrap().insert( - key.clone(), - (env.source.clone(), message.clone()), - ); - anyhow::anyhow!(message) - })? - .to_string(); - let _ = resolved_org; - env.config = Some(route.with_auth(lease.auth.clone())); self.session_auth .lock() @@ -386,8 +367,8 @@ impl Daemon { session_id: key.session_id.clone(), route: route.clone(), }; - if let Err(error) = journal::append_managed_run_key(&self.data_dir, managed_run_id, &record) - .await + if let Err(error) = + journal::append_managed_run_key(&self.data_dir, managed_run_id, &record).await { tracing::warn!( managed_run_id, diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 06c4a9a..336f845 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -206,7 +206,10 @@ pub async fn run_trace(args: TraceArgs, host: TraceHostContext) -> anyhow::Resul TraceCommand::Hook(hook_args) => { // A persistent hook must never fail the coding agent's turn. let result = async { - let route = host.services.resolve_route(RouteRequirements::default()).await?; + let route = host + .services + .resolve_route(RouteRequirements::default()) + .await?; run_hook(hook_args, route, host_info(&host)).await } .await; @@ -417,7 +420,10 @@ mod tests { .await .unwrap_err(); assert_eq!(error.to_string(), "no destination"); - assert_eq!(*services.route_requests.lock().unwrap(), [COMMAND_REQUIREMENTS]); + assert_eq!( + *services.route_requests.lock().unwrap(), + [COMMAND_REQUIREMENTS] + ); } } diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 243db31..183b00f 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -252,7 +252,8 @@ async fn start_daemon_with( let tmp = tempfile::tempdir().unwrap(); let data_dir = tmp.path().join("data"); let socket = test_endpoint(tmp.path()); - let handle = start_daemon_at_with(data_dir.clone(), socket.clone(), provider, sink_factory).await; + let handle = + start_daemon_at_with(data_dir.clone(), socket.clone(), provider, sink_factory).await; (data_dir, socket, handle, tmp) } diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index fcab8f1..77d67b9 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -109,8 +109,20 @@ async fn concurrent_imports_of_the_same_session_reach_independent_destinations() let project_a = tmp.path().join("spans-a"); let project_b = tmp.path().join("spans-b"); let (result_a, result_b) = tokio::join!( - import_transcript(&transcript, ImportSource::Codex, options(&project_a), None, false), - import_transcript(&transcript, ImportSource::Codex, options(&project_b), None, false), + import_transcript( + &transcript, + ImportSource::Codex, + options(&project_a), + None, + false + ), + import_transcript( + &transcript, + ImportSource::Codex, + options(&project_b), + None, + false + ), ); result_a.unwrap(); result_b.unwrap();