diff --git a/bt-daemon/README.md b/bt-daemon/README.md index d633658..b1b414d 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -22,16 +22,19 @@ One self-contained Cargo crate, liftable to its own repo by copying the `cli` feature for isolated testing/development. Env/flag static-token auth only; not an end-user artifact. -## Dual consumption +## Dual consumption and authentication -The daemon core is credential-passive — it only ever *receives* a resolved -`BackendAuth` with each session's config — so two front-ends share all core -behavior. The `cli` feature only enables the standalone binary and its logging -subscriber: +Hooks send only a non-secret `SessionRoute`: an optional profile and +organization selection plus the trace destination. The long-lived daemon asks +its host's `AuthProvider` for a credential lease, pins the returned canonical +profile to that session, and refreshes expiring leases as needed. Independent +sessions can therefore use different `bt` profiles without exposing tokens to +hook processes or JavaScript plugins. -1. **Embedded in `bt`** (production): `bt` fills `BackendAuth` from its profile - / OAuth / keychain auth. -2. **Standalone binary** (testing): fills it from `BRAINTRUST_API_KEY` etc. +1. **Embedded in `bt`** (production): the provider uses `bt`'s existing + profile, OAuth, refresh, keychain, organization, and backend URL machinery. +2. **Standalone binary** (testing): the provider uses `BRAINTRUST_API_KEY` and + related environment variables. ## Shared plugin settings @@ -41,11 +44,11 @@ or use `config.json` under `BT_DAEMON_DATA_DIR` (by default `~/.braintrust/state/bt-daemon/config.json` on Unix and `%LOCALAPPDATA%\Braintrust\bt-daemon\config.json` on Windows). -See [`config.json.example`](config.json.example). Supported settings are -`traceToBraintrust`, `project`, `flushOnTurnEnd`, and -`additionalMetadata`. File values override plugin environment fallbacks. -Credentials, auth tokens, organization selection, and backend URLs are not -settings here; production resolves them through `bt`. +See [`config.json.example`](config.json.example). `traceToBraintrust` controls +enablement and `route` stores the selected profile, organization, typed +destination, flush mode, and metadata. Omitting `route.auth.profile` selects +the default `bt` profile. Credentials and backend URLs are never stored here; +production resolves and refreshes them through `bt`. ## Build / test diff --git a/bt-daemon/config.json.example b/bt-daemon/config.json.example index 508b630..643c2ac 100644 --- a/bt-daemon/config.json.example +++ b/bt-daemon/config.json.example @@ -1,10 +1,19 @@ { - "_comment": "Shared by every coding-agent plugin connected to bt-daemon. Authentication and backend URLs are resolved by bt and do not belong here.", + "_comment": "Non-secret hook route. Authentication credentials and backend URLs are resolved by bt and do not belong here.", "traceToBraintrust": true, - "project": "my-coding-agents", - "flushOnTurnEnd": false, - "additionalMetadata": { - "team": "platform", - "environment": "development" + "route": { + "auth": { + "profile": "work", + "org_name": "acme" + }, + "destination": { + "type": "project_logs", + "project_name": "my-coding-agents" + }, + "flush_mode": "fire_and_forget", + "additional_metadata": { + "team": "platform", + "environment": "development" + } } } diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index 5247251..1ba9392 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -4,12 +4,15 @@ Status: **frozen for the prototype.** This is the contract between plugin shims (`hook` clients) and the daemon (`serve`), and between the embedded-in-`bt` front-end and the standalone test binary. -All hook clients share one daemon-level, non-credential settings file. Its path +Hook clients may use a non-credential settings file. Its path is `$BT_DAEMON_CONFIG`, falling back to `/config.json` and then the platform default daemon state directory. The hook front-end applies -`traceToBraintrust`, `project`, `flushOnTurnEnd`, and `additionalMetadata` -before constructing `SessionConfig`. Authentication and backend URLs are -resolved by `bt` and never read from this file. +`traceToBraintrust` and the stored `route` before sending an event. Setup can +persist a project-scoped route; a managed run can point `BT_DAEMON_CONFIG` at +an invocation-local route file, or set `BT_TRACE_SESSION_ROUTE` to the selected +route JSON, while keeping generated hook commands stable. +The profile is optional and defaults through `bt`. Credentials and backend +URLs are resolved and refreshed inside the daemon and never enter this file. `PROTOCOL_VERSION = 1`. @@ -89,7 +92,7 @@ Result: { "protocol_version": 1, "daemon_version": "0.1.0", - "capabilities": { "sources": ["codex", "claude-code", "debug"] } + "capabilities": { "sources": ["codex", "claude-code", "opencode", "debug"] } } ``` If `protocol_version` is incompatible the daemon returns an application error; @@ -159,11 +162,9 @@ Used for version handover and by tests. "event": "PostToolUse", "ts_ms": 1753639552123, "payload": { "…raw agent-native hook payload…": true }, - "config": { + "route": { "auth": { - "token": "sk-…", - "api_url": "https://api.braintrust.dev", - "app_url": "https://www.braintrust.dev", + "profile": "work", "org_name": "acme" }, "destination": { @@ -171,9 +172,6 @@ Used for version handover and by tests. "project_id": "project-uuid", "project_name": "codex" }, - "project": "codex", - "parent_span_id": null, - "root_span_id": null, "flush_mode": "fire_and_forget", "additional_metadata": { "…": "…" } } @@ -194,24 +192,23 @@ Field notes: daemon. - **`payload`** is opaque to transport and to everything except the translator for `source`. -- **`config`** carries shim-resolved credentials and trace settings. The shim - attaches it on **every** event (stateless shim); the daemon keeps the latest - per session and only re-inits the Braintrust sink when it changes. `auth` is - filled by `bt`'s `resolve_auth` when embedded, or from env/flags in the - standalone binary. `flush_mode` ∈ `fire_and_forget` | `flush_on_turn_end`. +- **`route`** carries non-secret auth selection and trace settings. `profile` + 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. The older `project`, - `parent_span_id`, `root_span_id`, and `_bt_experiment_id` fields remain - accepted when `destination` is absent. + carries the complete exported `SpanComponents` object. ### Redaction -`config.auth.token` (and any nested secret) is **never** written to the -journal or logs. The journal stores the envelope with `config.auth` reduced to -a non-secret fingerprint (`{ "api_url", "app_url", "org_name", "token_sha256_prefix" }`) -so replay can detect a credential change without persisting the secret; on -replay the live credentials must be re-supplied. +Live credentials returned by the host provider are **never** written to the +journal, logs, status, or RPC response. Envelopes journal only their non-secret +`route`, allowing restart recovery to resolve a fresh lease. ## Daemon lifecycle diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index 602738c..8b4a5ec 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -24,6 +24,7 @@ pub struct Counters { enum SessionMsg { Event(Box), + Configure(Box, oneshot::Sender<()>), Flush(oneshot::Sender), Shutdown(oneshot::Sender<()>), } @@ -102,6 +103,18 @@ impl Session { } } + /// Reconfigure the sink before a refresh-triggered flush. Queue ordering + /// guarantees that all earlier events are processed first. + pub async fn configure(&self, config: crate::wire::SessionConfig) -> anyhow::Result<()> { + let (reply_tx, reply_rx) = oneshot::channel(); + self.tx + .send(SessionMsg::Configure(Box::new(config), reply_tx)) + .map_err(|_| anyhow::anyhow!("session actor is gone"))?; + reply_rx + .await + .map_err(|_| anyhow::anyhow!("session actor dropped configuration reply")) + } + /// Drain, flush, and stop the actor (used on daemon shutdown). pub async fn shutdown(&self) { let (reply_tx, reply_rx) = oneshot::channel(); @@ -171,6 +184,8 @@ impl SessionActor { 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 { @@ -229,6 +244,12 @@ impl SessionActor { } self.counters.queued.fetch_sub(1, Ordering::Relaxed); } + SessionMsg::Configure(config, reply) => { + sink.configure(&config); + ctx.config = Some(*config); + self.refresh_permalink(sink.as_ref()); + let _ = reply.send(()); + } SessionMsg::Flush(reply) => { self.drain_flush(&mut translator, &mut sink, &ctx).await; let _ = reply.send(self.counters.queued.load(Ordering::Relaxed)); diff --git a/bt-daemon/src/journal.rs b/bt-daemon/src/journal.rs index f7b8cc0..c246509 100644 --- a/bt-daemon/src/journal.rs +++ b/bt-daemon/src/journal.rs @@ -5,9 +5,7 @@ //! Format: one [`RedactedEnvelope`] JSON value per line in //! `/journal/.ndjson`. -use crate::wire::{ - AuthFingerprint, BackendAuth, Envelope, RedactedConfig, RedactedEnvelope, SessionConfig, -}; +use crate::wire::{BackendAuth, Envelope, RedactedEnvelope}; use std::path::{Path, PathBuf}; use tokio::io::AsyncWriteExt; @@ -109,6 +107,16 @@ pub async fn gc_old_journals(data_dir: &Path, max_age: std::time::Duration) { /// rebuilding translator state; the sink must be re-supplied live credentials /// if replay needs to actually deliver. pub fn envelope_from_redacted(r: RedactedEnvelope) -> Envelope { + let route = r.route; + let config = route.as_ref().map(|route| { + route.with_auth(BackendAuth { + token: String::new(), + api_url: None, + app_url: None, + org_name: route.auth.org_name.clone(), + org_id: None, + }) + }); Envelope { source: r.source, source_version: r.source_version, @@ -116,30 +124,7 @@ pub fn envelope_from_redacted(r: RedactedEnvelope) -> Envelope { event: r.event, ts_ms: r.ts_ms, payload: r.payload, - config: r.config.map(config_from_redacted), - } -} - -fn config_from_redacted(c: RedactedConfig) -> SessionConfig { - let AuthFingerprint { - api_url, - app_url, - org_name, - .. - } = c.auth; - SessionConfig { - auth: BackendAuth { - token: String::new(), - api_url, - app_url, - org_name, - org_id: None, - }, - destination: c.destination, - project: c.project, - parent_span_id: c.parent_span_id, - root_span_id: c.root_span_id, - flush_mode: c.flush_mode, - additional_metadata: c.additional_metadata, + route, + config, } } diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 5dd2ec8..210fe58 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -1,14 +1,14 @@ //! bt-daemon: the embeddable library behind the Braintrust coding-agent //! tracing daemon. Two front-ends consume it (see `../DESIGN.md`): -//! * `bt` wires the [`clap::Args`] structs into its command tree and fills -//! [`wire::SessionConfig`] from its own auth resolution. +//! * `bt` wires the [`clap::Args`] structs into its command tree and exposes +//! its profile store through [`AuthProvider`]. //! * the feature-gated standalone `bt-daemon` binary does the same with //! env/flag token auth only, for isolated testing. //! -//! The core is credential-passive: it only ever *receives* a resolved -//! [`wire::BackendAuth`] with the session config, so both front-ends share all -//! core behavior. The `cli` feature only gates the standalone binary and its -//! logging subscriber. +//! Hook clients submit a non-secret [`wire::SessionRoute`]. The long-lived +//! daemon resolves and refreshes the selected profile through its host. +//! The `cli` feature only gates the standalone binary and its logging +//! subscriber. pub mod paths; @@ -25,7 +25,7 @@ mod transport; pub mod wire; pub use client::HostInfo; -pub use server::ServeOptions; +pub use server::{AuthLease, AuthProvider, AuthResolveReason, ServeOptions}; pub use sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, Sink, SinkFactory}; pub use translate::{ AgentTranslator, Registry, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory, @@ -36,7 +36,7 @@ use clap::{Args, ValueEnum}; use std::path::PathBuf; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use wire::{method, Envelope, SessionConfig, StatusResult, PROTOCOL_VERSION}; +use wire::{method, Envelope, SessionConfig, SessionRoute, StatusResult, PROTOCOL_VERSION}; /// Arguments for `serve`. #[derive(Debug, Clone, Args)] @@ -84,19 +84,9 @@ pub struct HookArgs { /// Bound an explicit turn/session-end flush. #[arg(long, default_value_t = 10_000)] pub flush_timeout_ms: u64, - /// Attach the agent session below an existing Braintrust span. - #[arg(long)] - pub parent_span_id: Option, - /// Existing trace root when attaching below a non-root parent. - #[arg(long)] - pub root_span_id: Option, /// JSON object merged into root-span metadata. #[arg(long)] pub additional_metadata: Option, - /// Route spans to an existing Braintrust experiment instead of project - /// logs. The Claude shim supplies this from CC_EXPERIMENT_ID. - #[arg(long)] - pub experiment_id: Option, } /// Arguments for `status`. @@ -140,13 +130,13 @@ pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<() /// Capture one hook event from `stdin` and forward it to the daemon. /// -/// `config` is the caller-resolved session config (auth + trace settings). +/// `route` contains only non-secret profile and destination selection. /// Returns `Ok` once the daemon has acked (journaled + enqueued). Callers that /// must never fail the agent's turn should treat any `Err` as non-fatal and /// exit 0. pub async fn run_hook( args: HookArgs, - mut config: SessionConfig, + mut route: SessionRoute, host: HostInfo, ) -> anyhow::Result<()> { let settings = settings::SharedSettings::load(); @@ -163,47 +153,19 @@ pub async fn run_hook( .or_else(|| json_str_field(&payload, &args.event_field)) .unwrap_or_default(); - if let Some(project) = settings.project.filter(|project| !project.is_empty()) { - config.project = Some(project); - } - match settings.flush_on_turn_end { - Some(true) => config.flush_mode = wire::FlushMode::FlushOnTurnEnd, - Some(false) => config.flush_mode = wire::FlushMode::FireAndForget, - None if args.flush_on_turn_end => config.flush_mode = wire::FlushMode::FlushOnTurnEnd, - None => {} + if let Some(configured_route) = settings.route { + route = configured_route; } - if args.parent_span_id.is_some() { - config.parent_span_id = args.parent_span_id.clone(); + if args.flush_on_turn_end { + route.flush_mode = wire::FlushMode::FlushOnTurnEnd; } - if args.root_span_id.is_some() { - config.root_span_id = args.root_span_id.clone(); - } - match (config.parent_span_id.clone(), config.root_span_id.clone()) { - (Some(parent), None) => config.root_span_id = Some(parent), - (None, Some(root)) => config.parent_span_id = Some(root), - _ => {} - } - if let Some(metadata) = settings.additional_metadata { - config.additional_metadata = Some(serde_json::Value::Object(metadata)); - } else if let Some(metadata) = &args.additional_metadata { + if let Some(metadata) = &args.additional_metadata { let value: serde_json::Value = serde_json::from_str(metadata) .map_err(|e| anyhow::anyhow!("invalid --additional-metadata JSON: {e}"))?; if !value.is_object() { anyhow::bail!("--additional-metadata must be a JSON object"); } - config.additional_metadata = Some(value); - } - if let Some(experiment_id) = &args.experiment_id { - let mut metadata = config - .additional_metadata - .take() - .and_then(|value| value.as_object().cloned()) - .unwrap_or_default(); - metadata.insert( - "_bt_experiment_id".to_string(), - serde_json::Value::String(experiment_id.clone()), - ); - config.additional_metadata = Some(serde_json::Value::Object(metadata)); + route.additional_metadata = Some(value); } let env = Envelope { source: args.source.clone(), @@ -212,14 +174,15 @@ pub async fn run_hook( event, ts_ms: now_ms(), payload, - config: Some(config), + route: Some(route), + config: None, }; let socket = paths::socket_path(args.socket.as_deref()); forward_envelope(&env, &socket, &host, args.no_spawn).await?; let should_flush = env.event == "SessionEnd" || (matches!( - env.config.as_ref().map(|c| c.flush_mode), + env.route.as_ref().map(|r| r.flush_mode), Some(wire::FlushMode::FlushOnTurnEnd) ) && matches!(env.event.as_str(), "Stop" | "SubagentStop")); if should_flush { @@ -457,6 +420,7 @@ pub fn debug_serve_options(version: impl Into, data_dir: &std::path::Pat sink_factory: Arc::new(DebugSinkFactory { dir: data_dir.join("spans"), }), + auth_provider: None, } } @@ -473,6 +437,7 @@ pub fn braintrust_serve_options( version: version.into(), translators, sink_factory: Arc::new(BraintrustSinkFactory::new(sink_config)), + auth_provider: None, } } diff --git a/bt-daemon/src/main.rs b/bt-daemon/src/main.rs index f30ff5b..65ab227 100644 --- a/bt-daemon/src/main.rs +++ b/bt-daemon/src/main.rs @@ -4,11 +4,12 @@ //! profiles, OAuth, or keychain here (that lives in `bt`). See //! the crate README's "Dual consumption" section. -use bt_daemon::wire::{BackendAuth, FlushMode, SessionConfig}; +use async_trait::async_trait; +use bt_daemon::wire::{AuthSelection, BackendAuth, SessionRoute, TraceDestination}; use bt_daemon::{ - braintrust_serve_options, paths, run_hook, run_import, run_serve, run_status, - BraintrustSinkConfig, DebugSinkFactory, HookArgs, HostInfo, ImportArgs, Registry, ServeArgs, - ServeOptions, StatusArgs, + braintrust_serve_options, paths, run_hook, run_import, run_serve, run_status, AuthLease, + AuthProvider, AuthResolveReason, BraintrustSinkConfig, DebugSinkFactory, HookArgs, HostInfo, + ImportArgs, Registry, ServeArgs, ServeOptions, StatusArgs, }; use clap::{Args, Parser, Subcommand}; use std::ffi::OsString; @@ -23,6 +24,42 @@ fn debug_serve_options(version: &str, data_dir: &std::path::Path) -> ServeOption sink_factory: Arc::new(DebugSinkFactory { dir: data_dir.join("spans"), }), + auth_provider: None, + } +} + +struct EnvironmentAuthProvider { + require_token: bool, +} + +#[async_trait] +impl AuthProvider for EnvironmentAuthProvider { + async fn resolve( + &self, + selection: &AuthSelection, + _reason: AuthResolveReason, + ) -> anyhow::Result { + let token = std::env::var("BRAINTRUST_API_KEY").unwrap_or_default(); + if self.require_token && token.is_empty() { + anyhow::bail!("BRAINTRUST_API_KEY is not set"); + } + Ok(AuthLease { + profile: selection + .profile + .clone() + .unwrap_or_else(|| "environment".to_string()), + auth: BackendAuth { + token, + api_url: std::env::var("BRAINTRUST_API_URL").ok(), + app_url: std::env::var("BRAINTRUST_APP_URL").ok(), + org_name: selection + .org_name + .clone() + .or_else(|| std::env::var("BRAINTRUST_ORG_NAME").ok()), + org_id: std::env::var("BRAINTRUST_ORG_ID").ok(), + }, + expires_at_ms: None, + }) } } @@ -63,7 +100,7 @@ enum Command { #[command(flatten)] args: HookArgs, #[command(flatten)] - auth: AuthArgs, + route: RouteArgs, }, /// Print daemon/session status. Status(StatusArgs), @@ -71,39 +108,34 @@ enum Command { Import(ImportArgs), } -/// Static-token backend auth from env/flags (no profile resolution). +/// Non-secret session selection. Credentials are resolved by the daemon. #[derive(Args)] -struct AuthArgs { - #[arg(long, env = "BRAINTRUST_API_KEY")] - api_key: Option, - #[arg(long, env = "BRAINTRUST_API_URL")] - api_url: Option, - #[arg(long, env = "BRAINTRUST_APP_URL")] - app_url: Option, +struct RouteArgs { + #[arg(long, env = "BRAINTRUST_PROFILE")] + profile: Option, #[arg(long = "org", env = "BRAINTRUST_ORG_NAME")] org_name: Option, - #[arg(long = "org-id", env = "BRAINTRUST_ORG_ID")] - org_id: Option, - #[arg(long, env = "BRAINTRUST_PROJECT")] + #[arg(long, env = "BRAINTRUST_PROJECT", conflicts_with = "destination")] project: Option, + #[arg(long, env = "BRAINTRUST_DESTINATION")] + destination: Option, } -impl AuthArgs { - fn into_config(self) -> SessionConfig { - SessionConfig { - auth: BackendAuth { - token: self.api_key.unwrap_or_default(), - api_url: self.api_url, - app_url: self.app_url, +impl RouteArgs { + fn into_route(self) -> SessionRoute { + SessionRoute { + auth: AuthSelection { + profile: self.profile, org_name: self.org_name, - org_id: self.org_id, }, - destination: None, - project: self.project, - parent_span_id: None, - root_span_id: None, - flush_mode: FlushMode::FireAndForget, - additional_metadata: None, + destination: self.destination.or_else(|| { + self.project + .map(|project_name| TraceDestination::ProjectLogs { + project_id: None, + project_name: Some(project_name), + }) + }), + ..SessionRoute::default() } } } @@ -137,7 +169,7 @@ async fn main() { app_url, } => { let data_dir = paths::data_dir(args.data_dir.as_deref()); - let opts = if debug_sink { + let mut opts = if debug_sink { debug_serve_options(VERSION, &data_dir) } else { let cfg = BraintrustSinkConfig { @@ -147,15 +179,17 @@ async fn main() { }; braintrust_serve_options(VERSION, cfg, Arc::new(Registry::default_agents())) }; + opts.auth_provider = Some(Arc::new(EnvironmentAuthProvider { + require_token: !debug_sink, + })); if let Err(e) = run_serve(args, opts).await { eprintln!("bt-daemon serve: {e}"); std::process::exit(1); } } - Command::Hook { args, auth } => { + Command::Hook { args, route } => { // A hook must NEVER fail the agent's turn: log and exit 0 on error. - let config = auth.into_config(); - if let Err(e) = run_hook(args, config, host_info()).await { + if let Err(e) = run_hook(args, route.into_route(), host_info()).await { eprintln!("bt-daemon hook (non-fatal): {e}"); } std::process::exit(0); diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index d15ab88..df6c879 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -12,12 +12,14 @@ use crate::wire::{ InitializeParams, InitializeResult, Message, Request, Response, RpcError, SessionStatus, ShutdownResult, StatusParams, StatusResult, PROTOCOL_VERSION, }; +use crate::wire::{AuthSelection, BackendAuth, SessionRoute}; use crate::{paths, ServeArgs}; +use async_trait::async_trait; use std::collections::HashMap; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::sync::Notify; @@ -27,6 +29,45 @@ pub struct ServeOptions { pub version: String, pub translators: Arc, pub sink_factory: Arc, + /// Host-owned access to Braintrust profiles, OAuth, and keychains. The + /// daemon owns lease timing and session routing; the embedding `bt` + /// process owns the credential store implementation. + pub auth_provider: Option>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthResolveReason { + Initial, + Expiring, + Unauthorized, +} + +/// A live credential lease. Only the canonical profile name is observable; +/// the backend credential remains in daemon memory and is never journaled. +#[derive(Debug, Clone)] +pub struct AuthLease { + pub profile: String, + pub auth: BackendAuth, + /// Epoch milliseconds. `None` is appropriate for non-expiring API keys. + pub expires_at_ms: Option, +} + +#[async_trait] +pub trait AuthProvider: Send + Sync { + /// Resolve without prompting. On refresh, `selection.profile` is the + /// canonical profile returned by the initial lease, keeping an active + /// session pinned even if the user's default profile changes. + async fn resolve( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result; +} + +#[derive(Clone)] +struct SessionAuthState { + route: SessionRoute, + lease: AuthLease, } pub struct Daemon { @@ -34,6 +75,9 @@ pub struct Daemon { data_dir: PathBuf, translators: Arc, sink_factory: Arc, + auth_provider: Option>, + session_auth: tokio::sync::Mutex>, + auth_errors: Mutex>, sessions: Mutex>>, started: Instant, last_activity: Mutex, @@ -48,6 +92,9 @@ impl Daemon { data_dir, translators: opts.translators, sink_factory: opts.sink_factory, + auth_provider: opts.auth_provider, + session_auth: tokio::sync::Mutex::new(HashMap::new()), + auth_errors: Mutex::new(HashMap::new()), sessions: Mutex::new(HashMap::new()), started: Instant::now(), last_activity: Mutex::new(Instant::now()), @@ -56,6 +103,119 @@ impl Daemon { }) } + 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"); + }; + let route = env + .route + .clone() + .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; + if route.destination.is_none() { + anyhow::bail!( + "session route is missing its trace destination; select a project or destination during `bt trace setup` or `bt trace run`" + ); + } + let (selection, reason, expected_profile) = { + let states = self.session_auth.lock().await; + match states.get(&env.session_id) { + 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(()); + } + ( + AuthSelection { + profile: Some(state.lease.profile.clone()), + org_name: state.lease.auth.org_name.clone(), + }, + AuthResolveReason::Expiring, + Some(state.lease.profile.clone()), + ) + } + None => (route.auth.clone(), AuthResolveReason::Initial, None), + } + }; + + let lease = provider.resolve(&selection, reason).await.map_err(|error| { + let message = format!( + "could not resolve Braintrust profile for {}: {error}; run `bt auth login` or select a profile explicitly", + env.source + ); + self.auth_errors.lock().unwrap().insert( + env.session_id.clone(), + (env.source.clone(), message.clone()), + ); + anyhow::anyhow!(message) + })?; + if let Some(expected) = expected_profile { + if lease.profile != expected { + anyhow::bail!( + "credential refresh changed profile from {expected:?} to {:?}", + lease.profile + ); + } + } + if let Some(expected_org) = route.auth.org_name.as_deref() { + if lease.auth.org_name.as_deref() != Some(expected_org) { + anyhow::bail!( + "profile {:?} resolved organization {:?}, expected {:?}", + lease.profile, + lease.auth.org_name, + expected_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(()) + } + + async fn refresh_session_before_flush(&self, session_id: &str) -> anyhow::Result<()> { + let Some(provider) = &self.auth_provider else { + return Ok(()); + }; + let Some(state) = self.session_auth.lock().await.get(session_id).cloned() else { + return Ok(()); + }; + if !lease_is_expiring(&state.lease) { + return Ok(()); + } + let selection = AuthSelection { + profile: Some(state.lease.profile.clone()), + org_name: state.lease.auth.org_name.clone(), + }; + let lease = provider + .resolve(&selection, AuthResolveReason::Expiring) + .await?; + if lease.profile != state.lease.profile { + anyhow::bail!("credential refresh changed the session profile"); + } + let config = state.route.with_auth(lease.auth.clone()); + self.session_auth.lock().await.insert( + session_id.to_string(), + SessionAuthState { + route: state.route, + lease, + }, + ); + let session = { self.sessions.lock().unwrap().get(session_id).cloned() }; + if let Some(session) = session { + session.configure(config).await?; + } + Ok(()) + } + fn touch(&self) { *self.last_activity.lock().unwrap() = Instant::now(); } @@ -120,6 +280,21 @@ impl Daemon { } } +fn lease_is_expiring(lease: &AuthLease) -> bool { + const REFRESH_WINDOW_MS: i64 = 60_000; + let Some(expires_at_ms) = lease.expires_at_ms else { + return false; + }; + expires_at_ms <= now_ms().saturating_add(REFRESH_WINDOW_MS) +} + +fn now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis() as i64) + .unwrap_or(0) +} + /// Bind the socket (handling a stale/rival socket), serve until shutdown, then /// drain sessions and remove the socket. pub async fn run(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { @@ -260,7 +435,7 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: Ok(()) } -async fn accept_event(daemon: &Arc, env: Envelope) -> Result<(), String> { +async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), String> { let source = env.source.clone(); let event = env.event.clone(); let session_id = env.session_id.clone(); @@ -268,6 +443,10 @@ async fn accept_event(daemon: &Arc, env: Envelope) -> Result<(), String> daemon.touch(); let result = async { + daemon + .configure_event(&mut env) + .await + .map_err(|error| format!("session auth failed: {error}"))?; let session = daemon .session_for(&env) .await @@ -340,6 +519,15 @@ 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 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, @@ -375,7 +563,7 @@ async fn handle_request(daemon: &Arc, req: Request) -> Response { impl Daemon { fn status(&self, p: StatusParams) -> StatusResult { let map = self.sessions.lock().unwrap(); - let sessions = map + let mut sessions: Vec<_> = map .iter() .filter(|(sid, _)| p.session_id.as_ref().is_none_or(|want| *want == **sid)) .map(|(sid, s)| SessionStatus { @@ -387,6 +575,22 @@ impl Daemon { 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) { + continue; + } + if map.contains_key(session_id) { + continue; + } + sessions.push(SessionStatus { + session_id: session_id.clone(), + source: source.clone(), + queued: 0, + spans_emitted: 0, + permalink: None, + last_error: Some(error.clone()), + }); + } StatusResult { daemon_version: self.version.clone(), uptime_ms: self.started.elapsed().as_millis() as u64, diff --git a/bt-daemon/src/settings.rs b/bt-daemon/src/settings.rs index 711cd87..807e042 100644 --- a/bt-daemon/src/settings.rs +++ b/bt-daemon/src/settings.rs @@ -5,22 +5,29 @@ //! across Codex, Claude Code, and future agent plugins. use crate::paths; +use crate::wire::SessionRoute; use serde::Deserialize; -use serde_json::{Map, Value}; use std::path::Path; +pub(crate) const SESSION_ROUTE_ENV: &str = "BT_TRACE_SESSION_ROUTE"; + #[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct SharedSettings { pub trace_to_braintrust: Option, - pub project: Option, - pub flush_on_turn_end: Option, - pub additional_metadata: Option>, + pub route: Option, } impl SharedSettings { pub(crate) fn load() -> Self { - Self::load_from(&paths::settings_path(None)) + let mut settings = Self::load_from(&paths::settings_path(None)); + if let Ok(raw) = std::env::var(SESSION_ROUTE_ENV) { + match serde_json::from_str(&raw) { + Ok(route) => settings.route = Some(route), + Err(error) => tracing::warn!("managed run route ignored: {error}"), + } + } + settings } fn load_from(path: &Path) -> Self { @@ -75,9 +82,11 @@ mod tests { &path, r#"{ "traceToBraintrust": true, - "project": "agents", - "flushOnTurnEnd": false, - "additionalMetadata": {"team": "platform"}, + "route": { + "destination": {"type": "project_logs", "project_name": "agents"}, + "flush_mode": "fire_and_forget", + "additional_metadata": {"team": "platform"} + }, "apiKey": "ignored", "apiUrl": "https://ignored.example" }"#, @@ -86,23 +95,20 @@ mod tests { let settings = SharedSettings::load_from(&path); assert_eq!(settings.trace_to_braintrust, Some(true)); - assert_eq!(settings.project.as_deref(), Some("agents")); - assert_eq!(settings.flush_on_turn_end, Some(false)); - assert_eq!( - settings.additional_metadata.unwrap()["team"], - Value::String("platform".into()) - ); + let route = settings.route.unwrap(); + assert_eq!(route.destination.unwrap().project_name(), Some("agents")); + assert_eq!(route.additional_metadata.unwrap()["team"], "platform"); } #[test] fn malformed_or_missing_settings_are_fail_open() { let temp = tempfile::tempdir().unwrap(); assert!(SharedSettings::load_from(&temp.path().join("missing.json")) - .project + .route .is_none()); let malformed = temp.path().join("malformed.json"); std::fs::write(&malformed, "{").unwrap(); - assert!(SharedSettings::load_from(&malformed).project.is_none()); + assert!(SharedSettings::load_from(&malformed).route.is_none()); } #[test] diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index e240791..f660e31 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -112,10 +112,16 @@ struct Creds { org_id: String, org_name: Option, destination: Option, - project: Option, - experiment_id: Option, - parent_span_id: Option, - root_span_id: Option, +} + +impl Creds { + fn same_as(&self, other: &Self) -> bool { + self.token == other.token + && self.org_id == other.org_id + && self.org_name == other.org_name + && serde_json::to_value(&self.destination).ok() + == serde_json::to_value(&other.destination).ok() + } } struct BraintrustSink { @@ -143,7 +149,7 @@ impl BraintrustSink { { return project_name.clone(); } - creds.project.clone().unwrap_or_else(|| self.source.clone()) + self.source.clone() } fn parent_info( @@ -156,16 +162,8 @@ impl BraintrustSink { if let Some(destination) = &creds.destination { return root_destination(destination, project); } - // Session root: attach under an external trace if the shim supplied - // one, else land it directly in the project's logs. - Ok(match (&creds.parent_span_id, &creds.root_span_id) { - (Some(p), Some(r)) => full_span(creds, project, p.clone(), r.clone()), - _ if creds.experiment_id.is_some() => ParentSpanInfo::Experiment { - object_id: creds.experiment_id.clone().unwrap(), - }, - _ => ParentSpanInfo::ProjectName { - project_name: project.to_string(), - }, + Ok(ParentSpanInfo::ProjectName { + project_name: project.to_string(), }) } else { Ok(full_span( @@ -251,7 +249,8 @@ impl Sink for BraintrustSink { .or_else(|| self.default_app_url.clone()) .unwrap_or_else(|| DEFAULT_APP_URL.to_string()); let new_urls = (api, app); - if self.urls.as_ref() != Some(&new_urls) { + let urls_changed = self.urls.as_ref() != Some(&new_urls); + if urls_changed { // A session shouldn't change backend URLs mid-flight; if it does, // rebind the client on the next emit. Pre-change open handles stay // bound to the old client (pathological; just noted). @@ -264,21 +263,24 @@ impl Sink for BraintrustSink { self.urls = Some(new_urls); self.client = None; } - self.creds = Some(Creds { + let next_creds = Creds { token: config.auth.token.clone(), org_id: config.auth.org_id.clone().unwrap_or_default(), org_name: config.auth.org_name.clone(), destination: config.destination.clone(), - project: config.project.clone(), - experiment_id: config - .additional_metadata + }; + // Span handles capture their credentials when built. Recreate them + // after a profile token or routing change; deterministic row ids make + // subsequent updates merge into the same Braintrust rows. + if urls_changed + || self + .creds .as_ref() - .and_then(|v| v.get("_bt_experiment_id")) - .and_then(Value::as_str) - .map(ToOwned::to_owned), - parent_span_id: config.parent_span_id.clone(), - root_span_id: config.root_span_id.clone(), - }); + .is_some_and(|old| !old.same_as(&next_creds)) + { + self.open.clear(); + } + self.creds = Some(next_creds); } async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { @@ -323,17 +325,6 @@ fn full_span( propagated_event: components.propagated_event, }; } - if let Some(experiment_id) = &creds.experiment_id { - return ParentSpanInfo::FullSpan { - object_type: SpanObjectType::Experiment, - object_id: Some(experiment_id.clone()), - compute_object_metadata_args: None, - span_id, - root_span_id, - span_parents: None, - propagated_event: None, - }; - } let mut cma = Map::new(); cma.insert( "project_name".to_string(), @@ -382,7 +373,7 @@ fn root_destination( fn destination_root(creds: &Creds) -> Option { match &creds.destination { Some(TraceDestination::ParentSpan { components }) => components.root_span_id.clone(), - _ => creds.root_span_id.clone(), + _ => None, } } diff --git a/bt-daemon/src/transcript_import.rs b/bt-daemon/src/transcript_import.rs index 5d94d00..0a18150 100644 --- a/bt-daemon/src/transcript_import.rs +++ b/bt-daemon/src/transcript_import.rs @@ -396,6 +396,7 @@ fn envelope( event: event.into(), ts_ms, payload, + route: None, config: None, } } diff --git a/bt-daemon/src/translate/codex.rs b/bt-daemon/src/translate/codex.rs index 8d44534..187bb2b 100644 --- a/bt-daemon/src/translate/codex.rs +++ b/bt-daemon/src/translate/codex.rs @@ -142,7 +142,7 @@ impl AgentTranslator for CodexTranslator { if let Some(config) = &ctx.config { self.external_parent_span_id = config.attached_span_ids().0; - self.project = config.project.clone(); + self.project = config.project_name().map(ToOwned::to_owned); self.additional_metadata = config .additional_metadata .as_ref() diff --git a/bt-daemon/src/wire/envelope.rs b/bt-daemon/src/wire/envelope.rs index 89c0482..67f3fd6 100644 --- a/bt-daemon/src/wire/envelope.rs +++ b/bt-daemon/src/wire/envelope.rs @@ -3,7 +3,6 @@ use braintrust_sdk_rust::SpanComponents; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; /// One captured hook event, forwarded from a shim to the daemon. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -22,26 +21,64 @@ pub struct Envelope { pub ts_ms: i64, /// The raw agent-native hook payload; opaque except to the translator. pub payload: serde_json::Value, - /// Shim-resolved credentials + trace settings. Present on every event from - /// a stateless shim; the daemon keeps the latest per session. + /// Non-secret, immutable routing intent for this session. New clients use + /// this instead of resolving credentials themselves. The daemon host maps + /// the selected profile and organization to live credentials. #[serde(default, skip_serializing_if = "Option::is_none")] + pub route: Option, + /// Daemon-internal resolved configuration. This field is never serialized; + /// clients can only submit `route`. + #[serde(skip)] pub config: Option, } +/// Non-secret profile selection. A profile identifies the stored Braintrust +/// user credentials; an optional organization constrains profiles that can +/// address more than one organization. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct AuthSelection { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub profile: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub org_name: Option, +} + +/// Immutable, journal-safe routing and trace settings for one agent session. +/// Credentials are deliberately absent and are resolved inside the daemon. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SessionRoute { + #[serde(default)] + pub auth: AuthSelection, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub destination: Option, + #[serde(default)] + pub flush_mode: FlushMode, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub additional_metadata: Option, +} + +impl SessionRoute { + pub fn with_auth(&self, auth: BackendAuth) -> SessionConfig { + SessionConfig { + auth, + destination: self.destination.clone(), + flush_mode: self.flush_mode, + additional_metadata: self.additional_metadata.clone(), + } + } + + pub fn same_route(&self, other: &Self) -> bool { + serde_json::to_value(self).ok() == serde_json::to_value(other).ok() + } +} + /// Trace settings and backend credentials resolved by the shim. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SessionConfig { pub auth: BackendAuth, - /// Typed destination for new front-ends. When present, this takes - /// precedence over the legacy project and span-attachment fields below. + /// Typed trace destination. #[serde(default, skip_serializing_if = "Option::is_none")] pub destination: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub project: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub parent_span_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub root_span_id: Option, #[serde(default)] pub flush_mode: FlushMode, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -72,7 +109,14 @@ impl SessionConfig { if let Some(TraceDestination::ParentSpan { components }) = &self.destination { return (components.span_id.clone(), components.root_span_id.clone()); } - (self.parent_span_id.clone(), self.root_span_id.clone()) + (None, None) + } + + pub fn project_name(&self) -> Option<&str> { + match &self.destination { + Some(TraceDestination::ProjectLogs { project_name, .. }) => project_name.as_deref(), + _ => None, + } } } @@ -102,8 +146,17 @@ impl std::str::FromStr for TraceDestination { } } +impl TraceDestination { + pub fn project_name(&self) -> Option<&str> { + match self { + Self::ProjectLogs { project_name, .. } => project_name.as_deref(), + _ => None, + } + } +} + /// Backend credentials. `token` is an API key or an OAuth access token; the -/// daemon does not care which. Never persisted (see [`SessionConfig::redacted`]). +/// daemon does not care which. Never serialized or persisted. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BackendAuth { pub token: String, @@ -130,43 +183,9 @@ pub enum FlushMode { FlushOnTurnEnd, } -/// A non-secret fingerprint of [`BackendAuth`], written to the journal in -/// place of the token so replay can detect a credential change without ever -/// persisting the secret. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct AuthFingerprint { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub api_url: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub app_url: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub org_name: Option, - /// First 12 hex chars of SHA-256(token). Enough to detect rotation, far - /// too little to recover the token. - pub token_sha256_prefix: String, -} - -impl BackendAuth { - pub fn fingerprint(&self) -> AuthFingerprint { - let digest = Sha256::digest(self.token.as_bytes()); - let hex = digest.iter().fold(String::with_capacity(64), |mut s, b| { - use std::fmt::Write; - let _ = write!(s, "{b:02x}"); - s - }); - AuthFingerprint { - api_url: self.api_url.clone(), - app_url: self.app_url.clone(), - org_name: self.org_name.clone(), - token_sha256_prefix: hex[..12].to_string(), - } - } -} - impl Envelope { - /// A copy of this envelope safe to write to the journal: the live token is - /// replaced by an [`AuthFingerprint`]. The rest of `config` (project, - /// span-attach ids, flush mode, metadata) is retained — none of it secret. + /// A copy safe to journal. Only the non-secret route is serializable; the + /// daemon-internal resolved config is excluded by construction. pub fn redacted(&self) -> RedactedEnvelope { RedactedEnvelope { source: self.source.clone(), @@ -175,15 +194,7 @@ impl Envelope { event: self.event.clone(), ts_ms: self.ts_ms, payload: self.payload.clone(), - config: self.config.as_ref().map(|c| RedactedConfig { - auth: c.auth.fingerprint(), - destination: c.destination.clone(), - project: c.project.clone(), - parent_span_id: c.parent_span_id.clone(), - root_span_id: c.root_span_id.clone(), - flush_mode: c.flush_mode, - additional_metadata: c.additional_metadata.clone(), - }), + route: self.route.clone(), } } } @@ -200,24 +211,7 @@ pub struct RedactedEnvelope { pub ts_ms: i64, pub payload: serde_json::Value, #[serde(default, skip_serializing_if = "Option::is_none")] - pub config: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct RedactedConfig { - pub auth: AuthFingerprint, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub destination: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub project: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub parent_span_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub root_span_id: Option, - #[serde(default)] - pub flush_mode: FlushMode, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub additional_metadata: Option, + pub route: Option, } #[cfg(test)] @@ -232,6 +226,17 @@ mod tests { event: "PostToolUse".into(), ts_ms: 1_753_639_552_123, payload: serde_json::json!({ "session_id": "sess-1", "tool_name": "shell" }), + route: Some(SessionRoute { + auth: AuthSelection { + profile: Some("work".into()), + org_name: Some("acme".into()), + }, + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("codex".into()), + }), + ..SessionRoute::default() + }), config: Some(SessionConfig { auth: BackendAuth { token: "sk-super-secret".into(), @@ -241,9 +246,6 @@ mod tests { org_id: None, }, destination: None, - project: Some("codex".into()), - parent_span_id: None, - root_span_id: None, flush_mode: FlushMode::FireAndForget, additional_metadata: None, }), @@ -256,11 +258,13 @@ mod tests { let s = serde_json::to_string(&e).unwrap(); let back: Envelope = serde_json::from_str(&s).unwrap(); assert_eq!(back.session_id, "sess-1"); - assert_eq!(back.config.unwrap().auth.token, "sk-super-secret"); + assert_eq!(back.route.unwrap().auth.profile.as_deref(), Some("work")); + assert!(back.config.is_none()); + assert!(!s.contains("sk-super-secret")); } #[test] - fn redaction_drops_the_token_but_keeps_settings() { + fn journal_form_contains_only_the_route() { let e = sample(); let r = e.redacted(); let s = serde_json::to_string(&r).unwrap(); @@ -268,19 +272,42 @@ mod tests { !s.contains("sk-super-secret"), "token leaked into journal form: {s}" ); - let cfg = r.config.unwrap(); - assert_eq!(cfg.project.as_deref(), Some("codex")); - assert_eq!(cfg.auth.org_name.as_deref(), Some("acme")); - assert_eq!(cfg.auth.token_sha256_prefix.len(), 12); + assert_eq!(r.route.unwrap().auth.profile.as_deref(), Some("work")); } #[test] - fn fingerprint_changes_with_token() { - let mut a = sample().config.unwrap().auth; - let f1 = a.fingerprint(); - a.token = "sk-different".into(); - let f2 = a.fingerprint(); - assert_ne!(f1.token_sha256_prefix, f2.token_sha256_prefix); + fn route_is_journal_safe_and_builds_resolved_config() { + let route = SessionRoute { + auth: AuthSelection { + profile: Some("work".into()), + org_name: Some("acme".into()), + }, + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("agent-traces".into()), + }), + ..SessionRoute::default() + }; + let config = route.with_auth(BackendAuth { + token: "secret".into(), + api_url: None, + app_url: None, + org_name: Some("acme".into()), + org_id: None, + }); + assert!(matches!( + config.destination, + Some(TraceDestination::ProjectLogs { project_name: Some(ref name), .. }) + if name == "agent-traces" + )); + assert_eq!(route.auth.profile.as_deref(), Some("work")); + + let mut envelope = sample(); + envelope.route = Some(route); + envelope.config = Some(config); + let journal = serde_json::to_string(&envelope.redacted()).unwrap(); + assert!(journal.contains("work")); + assert!(!journal.contains("secret")); } #[test] diff --git a/bt-daemon/src/wire/mod.rs b/bt-daemon/src/wire/mod.rs index 7e93121..8ede81d 100644 --- a/bt-daemon/src/wire/mod.rs +++ b/bt-daemon/src/wire/mod.rs @@ -10,8 +10,8 @@ mod methods; mod rpc; pub use envelope::{ - AuthFingerprint, BackendAuth, Envelope, FlushMode, RedactedConfig, RedactedEnvelope, - SessionConfig, TraceDestination, + AuthSelection, BackendAuth, Envelope, FlushMode, RedactedEnvelope, SessionConfig, SessionRoute, + TraceDestination, }; pub use methods::{ method, Capabilities, ClientInfo, EventLogResult, FlushParams, FlushResult, InitializeParams, diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 3423b0a..331aa26 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -20,10 +20,10 @@ fn session_config(base: &str) -> SessionConfig { org_name: Some("acme".into()), org_id: None, }, - destination: None, - project: Some("my-project".into()), - parent_span_id: None, - root_span_id: None, + destination: Some(TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("my-project".into()), + }), flush_mode: FlushMode::FireAndForget, additional_metadata: None, } @@ -187,8 +187,11 @@ async fn attached_trace_children_keep_the_external_root() { }); let mut sink = factory.create("sess-1", "codex").unwrap(); let mut config = session_config(&base); - config.parent_span_id = Some("external-parent".into()); - config.root_span_id = Some("external-root".into()); + let mut components = SpanComponents::new(SpanObjectType::ProjectLogs); + components.object_id = Some("proj-parent".into()); + components.span_id = Some("external-parent".into()); + components.root_span_id = Some("external-root".into()); + config.destination = Some(TraceDestination::ParentSpan { components }); sink.configure(&config); sink.emit(&[SpanOp::Insert(row( "child", diff --git a/bt-daemon/tests/claude_translator.rs b/bt-daemon/tests/claude_translator.rs index 666a504..2ed2c02 100644 --- a/bt-daemon/tests/claude_translator.rs +++ b/bt-daemon/tests/claude_translator.rs @@ -50,6 +50,7 @@ fn replay(name: &str) -> Vec { event: record["hook"].as_str().unwrap().into(), ts_ms, payload, + route: None, config: None, }; ops.extend(translator.handle(&env, &ctx).unwrap()); @@ -257,6 +258,7 @@ fn claude_permission_denied_and_failed_tools_are_first_class_spans() { event: name.into(), ts_ms: 1, payload, + route: None, config: None, }; let mut ops = translator @@ -343,6 +345,7 @@ fn claude_pairs_tool_lifecycle_and_marks_explicit_skills_and_stop_failures() { event: name.into(), ts_ms, payload, + route: None, config: None, }; let mut ops = Vec::new(); @@ -459,6 +462,7 @@ fn claude_groups_streamed_rows_and_reads_late_final_output_at_session_end() { event: name.into(), ts_ms, payload, + route: None, config: None, }; let mut ops = translator diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 6797c36..fdb795e 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -61,6 +61,7 @@ fn envelope(session: &str, event: &str, transcript_path: &str, extra: Value) -> event: event.into(), ts_ms: 0, payload, + route: None, config: None, } } @@ -388,10 +389,10 @@ fn configured_ctx(session_id: &str, additional_metadata: Value) -> SessionCtx { org_name: None, org_id: None, }, - destination: None, - project: Some("team-project".into()), - parent_span_id: None, - root_span_id: None, + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("team-project".into()), + }), flush_mode: FlushMode::FireAndForget, additional_metadata: Some(additional_metadata), }), diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index f7d1fa4..f035718 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -2,13 +2,16 @@ //! debug sink. Runs the daemon in-process on a temp socket (no process //! spawning, so it's deterministic). -use bt_daemon::wire::{BackendAuth, Envelope, FlushMode, SessionConfig}; +use async_trait::async_trait; +use bt_daemon::wire::{AuthSelection, BackendAuth, Envelope, SessionRoute}; use bt_daemon::{ debug_serve_options, flush_session, forward_envelope, run_serve, run_status, shutdown_daemon, - HostInfo, ServeArgs, StatusArgs, + AuthLease, AuthProvider, AuthResolveReason, HostInfo, ServeArgs, StatusArgs, }; use std::ffi::OsString; use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; fn dummy_host() -> HostInfo { @@ -20,24 +23,6 @@ fn dummy_host() -> HostInfo { } } -fn config_with_secret() -> SessionConfig { - SessionConfig { - auth: BackendAuth { - token: "sk-TOP-SECRET-abc123".into(), - api_url: Some("https://api.braintrust.dev".into()), - app_url: None, - org_name: Some("acme".into()), - org_id: None, - }, - destination: None, - project: Some("codex".into()), - parent_span_id: None, - root_span_id: None, - flush_mode: FlushMode::FireAndForget, - additional_metadata: None, - } -} - fn envelope(session_id: &str, event: &str, ts_ms: i64) -> Envelope { Envelope { source: "debug".into(), @@ -46,10 +31,96 @@ fn envelope(session_id: &str, event: &str, ts_ms: i64) -> Envelope { event: event.into(), ts_ms, payload: serde_json::json!({ "session_id": session_id, "hook_event_name": event, "n": ts_ms }), - config: Some(config_with_secret()), + route: Some(SessionRoute { + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("codex".into()), + }), + ..SessionRoute::default() + }), + config: None, + } +} + +fn routed_envelope(session_id: &str, profile: &str, org: &str, event: &str) -> Envelope { + let mut env = envelope(session_id, event, 1); + env.route = Some(SessionRoute { + auth: AuthSelection { + profile: Some(profile.into()), + org_name: Some(org.into()), + }, + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some(format!("{profile}-traces")), + }), + ..SessionRoute::default() + }); + env +} + +struct TestAuthProvider { + calls: Mutex>, + fail: bool, + first_lease_expired: bool, +} + +#[async_trait] +impl AuthProvider for TestAuthProvider { + async fn resolve( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result { + let mut calls = self.calls.lock().unwrap(); + calls.push((selection.clone(), reason)); + let call_index = calls.len(); + drop(calls); + if self.fail { + anyhow::bail!("profile credential unavailable") + } + let profile = selection + .profile + .clone() + .unwrap_or_else(|| "default".into()); + Ok(AuthLease { + profile: profile.clone(), + auth: BackendAuth { + token: format!("secret-{profile}-{call_index}"), + api_url: Some(format!("https://{profile}.example.test")), + app_url: None, + org_name: selection.org_name.clone(), + org_id: Some(format!("org-{profile}")), + }, + expires_at_ms: (self.first_lease_expired && call_index == 1).then_some(0), + }) } } +async fn start_routed_daemon( + provider: 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 args = ServeArgs { + socket: Some(socket.clone()), + data_dir: Some(data_dir.clone()), + idle_timeout_secs: 0, + }; + let mut opts = debug_serve_options("test", &data_dir); + opts.auth_provider = Some(provider); + let handle = tokio::spawn(async move { + let _ = run_serve(args, opts).await; + }); + wait_for(&socket).await; + (data_dir, socket, handle, tmp) +} + fn test_endpoint(tmp: &Path) -> PathBuf { #[cfg(unix)] { @@ -115,7 +186,12 @@ async fn start_daemon() -> ( data_dir: Some(data_dir.clone()), idle_timeout_secs: 0, // disable the watchdog for the test }; - let opts = debug_serve_options("test", &data_dir); + let mut opts = debug_serve_options("test", &data_dir); + opts.auth_provider = Some(Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + })); let handle = tokio::spawn(async move { let _ = run_serve(args, opts).await; }); @@ -129,7 +205,12 @@ async fn start_daemon_at(data_dir: PathBuf, socket: PathBuf) -> tokio::task::Joi data_dir: Some(data_dir.clone()), idle_timeout_secs: 0, }; - let opts = debug_serve_options("test", &data_dir); + let mut opts = debug_serve_options("test", &data_dir); + opts.auth_provider = Some(Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + })); let handle = tokio::spawn(async move { let _ = run_serve(args, opts).await; }); @@ -141,6 +222,160 @@ async fn shutdown(socket: &Path) { shutdown_daemon(socket).await.unwrap(); } +#[tokio::test] +async fn routed_sessions_resolve_multiple_profiles_without_journaling_credentials() { + 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.clone()).await; + let host = dummy_host(); + + for (session, profile, org) in [ + ("work-session", "work", "work-org"), + ("personal-session", "personal", "personal-org"), + ] { + forward_envelope( + &routed_envelope(session, profile, org, "SessionStart"), + &socket, + &host, + false, + ) + .await + .unwrap(); + flush_session(session, &socket, 5000).await.unwrap(); + } + + let calls = provider.calls.lock().unwrap().clone(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].0.profile.as_deref(), Some("work")); + assert_eq!(calls[0].0.org_name.as_deref(), Some("work-org")); + assert_eq!(calls[0].1, AuthResolveReason::Initial); + assert_eq!(calls[1].0.profile.as_deref(), Some("personal")); + assert_eq!(calls[1].0.org_name.as_deref(), Some("personal-org")); + assert_eq!(calls[1].1, AuthResolveReason::Initial); + + for session in ["work-session", "personal-session"] { + let journal = + std::fs::read_to_string(data_dir.join("journal").join(format!("{session}.ndjson"))) + .unwrap(); + assert!(journal.contains("\"route\"")); + assert!(!journal.contains("secret-")); + assert!(!journal.contains("token_sha256_prefix")); + assert!(!journal.contains("\"config\"")); + } + + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[tokio::test] +async fn expiring_profile_lease_is_refreshed_for_the_pinned_profile() { + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: true, + }); + let (_data_dir, socket, handle, _tmp) = start_routed_daemon(provider.clone()).await; + let host = dummy_host(); + + forward_envelope( + &routed_envelope("refresh", "work", "work-org", "SessionStart"), + &socket, + &host, + false, + ) + .await + .unwrap(); + forward_envelope( + &routed_envelope("refresh", "work", "work-org", "Stop"), + &socket, + &host, + false, + ) + .await + .unwrap(); + + let calls = provider.calls.lock().unwrap().clone(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].1, AuthResolveReason::Initial); + assert_eq!(calls[1].1, AuthResolveReason::Expiring); + assert_eq!(calls[1].0.profile.as_deref(), Some("work")); + assert_eq!(calls[1].0.org_name.as_deref(), Some("work-org")); + + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[tokio::test] +async fn active_session_rejects_route_changes() { + 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 host = dummy_host(); + + forward_envelope( + &routed_envelope("pinned", "work", "work-org", "SessionStart"), + &socket, + &host, + false, + ) + .await + .unwrap(); + let error = forward_envelope( + &routed_envelope("pinned", "personal", "personal-org", "Stop"), + &socket, + &host, + false, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("session route changed")); + + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[tokio::test] +async fn auth_resolution_failure_is_reported_without_exposing_credentials() { + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: true, + first_lease_expired: false, + }); + let (data_dir, socket, handle, _tmp) = start_routed_daemon(provider.clone()).await; + let host = dummy_host(); + + let error = forward_envelope( + &routed_envelope("login-needed", "missing", "missing-org", "SessionStart"), + &socket, + &host, + false, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("bt auth login")); + let status = run_status(StatusArgs { + socket: Some(socket.clone()), + session_id: Some("login-needed".into()), + }) + .await + .unwrap() + .unwrap(); + let status_error = status.sessions[0].last_error.as_deref().unwrap(); + assert!(status_error.contains("select a profile explicitly")); + assert!(!status_error.contains("secret-")); + assert!(!data_dir.join("journal/login-needed.ndjson").exists()); + assert_eq!(provider.calls.lock().unwrap().len(), 1); + + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn events_are_ordered_journaled_and_emitted() { let (data_dir, socket, handle, _tmp) = start_daemon().await; @@ -156,7 +391,7 @@ async fn events_are_ordered_journaled_and_emitted() { assert!(flushed.flushed, "flush did not complete: {flushed:?}"); assert_eq!(flushed.pending, 0); - // Journal: three events, in order, token redacted. + // Journal: three events, in order, with only the non-secret route. let journal = data_dir.join("journal").join("sess-1.ndjson"); let jtext = std::fs::read_to_string(&journal).unwrap(); let jlines: Vec<&str> = jtext.lines().filter(|l| !l.trim().is_empty()).collect(); @@ -170,10 +405,7 @@ async fn events_are_ordered_journaled_and_emitted() { !jtext.contains("sk-TOP-SECRET-abc123"), "token leaked into journal!" ); - assert!( - jtext.contains("token_sha256_prefix"), - "journal missing auth fingerprint" - ); + assert!(!jtext.contains("token_sha256_prefix")); let events: Vec = jlines .iter() diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 26e7972..c4d0e16 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -15,6 +15,7 @@ fn options(output: &std::path::Path) -> ServeOptions { version: "test".into(), translators: Arc::new(Registry::default_agents()), sink_factory: Arc::new(DebugSinkFactory { dir: output.into() }), + auth_provider: None, } } diff --git a/bt-daemon/tests/support/agent_process.rs b/bt-daemon/tests/support/agent_process.rs index 820c25d..f240494 100644 --- a/bt-daemon/tests/support/agent_process.rs +++ b/bt-daemon/tests/support/agent_process.rs @@ -60,9 +60,11 @@ impl AgentTestWorld { &config_path, serde_json::to_vec_pretty(&json!({ "traceToBraintrust": true, - "project": "agent-e2e", - "flushOnTurnEnd": true, - "additionalMetadata": {"test_harness": true} + "route": { + "destination": {"type": "project_logs", "project_name": "agent-e2e"}, + "flush_mode": "flush_on_turn_end", + "additional_metadata": {"test_harness": true} + } })) .unwrap(), ) @@ -86,6 +88,7 @@ impl AgentTestWorld { .kill_on_drop(true); if ingest_mode == TestBackendMode::Mock { command + .env("BRAINTRUST_API_KEY", "test-key") .env("BRAINTRUST_API_URL", collector_server.uri()) .env("BRAINTRUST_APP_URL", collector_server.uri()); }