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 f759738..7f8bf53 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -4,11 +4,10 @@ //! 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. -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/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 459ee8a..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,10 +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, + }); } }; let mut conn = client::Conn::new(stream); @@ -485,6 +492,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..fc697c4 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) @@ -178,16 +195,16 @@ impl Daemon { 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 +222,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 +239,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,39 +345,73 @@ 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) -> bool { self.managed_run_sessions .lock() .unwrap() .entry(managed_run_id.to_string()) .or_default() - .insert(session_id.to_string()); + .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 session_ids = 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, + 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)) @@ -367,6 +469,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); @@ -495,35 +598,44 @@ 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); + 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(()) => { - if let Some(managed_run_id) = managed_run_id { - daemon.record_managed_run_session(&managed_run_id, &session_id); + Ok(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") } 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 +692,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 +763,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..336f845 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. @@ -117,6 +130,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 +139,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, + requirements: RouteRequirements, +) -> anyhow::Result { + let mut route = host.services.resolve_route(requirements).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 +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 = host.services.resolve_route(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) => { @@ -144,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(false).await?; + let route = host + .services + .resolve_route(RouteRequirements::default()) + .await?; run_hook(hook_args, route, host_info(&host)).await } .await; @@ -172,13 +237,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 = host.services.resolve_route(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() { @@ -229,7 +305,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") } @@ -243,9 +319,10 @@ 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>, } impl RecordingHost { @@ -254,17 +331,27 @@ 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, } } } #[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); } @@ -291,7 +378,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, @@ -312,6 +399,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 [ @@ -328,13 +420,35 @@ 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), COMMAND_REQUIREMENTS) + .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), COMMAND_REQUIREMENTS) + .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 [ + for (destination, destination_required) in [ (None, true), ( Some(TraceDestination::ProjectLogs { @@ -361,7 +475,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, + }] + ); } } @@ -389,7 +509,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/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")] diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index dfa9d46..183b00f 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,47 @@ 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 +493,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 +885,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 +1037,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 +1052,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..77d67b9 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -83,6 +83,57 @@ 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();