diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs deleted file mode 100644 index 786ad20..0000000 --- a/bt-daemon/src/server.rs +++ /dev/null @@ -1,2698 +0,0 @@ -//! The daemon: owns the session map + shared deps, binds the UDS listener, -//! serves JSON-RPC connections, and shuts down gracefully (idle timeout, -//! `daemon.shutdown`, or SIGINT/SIGTERM). - -use crate::dispatch::{hydrate_transcript_reference, ReplayPlan, Session, SessionOptions}; -use crate::journal::{self, JournalWriter}; -use crate::sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, SinkFactory}; -use crate::translate::Registry; -use crate::transport::{self, Listener, ServerStream}; -use crate::wire::{ - error_code, method, AuthDiagnoseParams, AuthDiagnoseResult, Capabilities, ClientInfo, Envelope, - EventLogResult, FlushParams, FlushResult, InitializeParams, InitializeResult, - ManagedRunFlushParams, 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 serde::{Deserialize, Serialize}; -use serde_json::Value; -use sha2::{Digest, Sha256}; -use std::collections::{HashMap, HashSet}; -use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; -use std::sync::Mutex; -use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; -use tokio::sync::{mpsc, oneshot, Notify}; - -/// Injected dependencies for `serve`, so `bt` / tests can supply a sink -/// factory (Braintrust in production, debug in tests) and a version string. -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 non-secret selection is -/// observable; the backend credential remains in daemon memory and is never -/// journaled. -#[derive(Debug, Clone)] -pub struct AuthLease { - pub selection: AuthSelection, - pub auth: BackendAuth, - /// Epoch milliseconds. `None` is appropriate for non-expiring API keys. - pub expires_at_ms: Option, -} - -// `async_trait` marks its boxed futures as `must_use`; Clippy 1.99 flags that -// generated annotation as redundant on async trait methods. -#[allow(clippy::double_must_use)] -#[async_trait] -pub trait AuthProvider: Send + Sync { - /// Resolve without prompting. On refresh, `selection` is the canonical - /// source returned by the initial lease, keeping an active session pinned - /// even if the user's default profile or environment changes. - async fn resolve( - &self, - selection: &AuthSelection, - reason: AuthResolveReason, - ) -> anyhow::Result; -} - -#[derive(Clone)] -struct SessionAuthState { - route: SessionRoute, - lease: AuthLease, -} - -// Correlation is a short, optional discovery step, never a delivery queue. -const CORRELATION_VERSION: u32 = 1; -// Long enough for a slow cold start of the second hook (e.g. an antivirus scan -// of `bt.exe`); only sessions with a frozen parent candidate ever wait. -const CORRELATION_WAIT: Duration = Duration::from_secs(5); -const CORRELATION_EVENT_LIMIT: usize = 3; - -#[derive(Clone, Serialize, Deserialize)] -struct PendingSession { - // Missing versions predate safe ancestry matching; see `restore`. - #[serde(default)] - version: u32, - #[serde(default, skip_serializing_if = "Option::is_none")] - linked_route: Option, - #[serde(default)] - standalone: bool, - #[serde(default, skip_serializing_if = "Option::is_none")] - parent: Option, - #[serde(default)] - first_seen_ms: i64, - #[serde(skip)] - deadline: Option, - // Discovery expiry is not an authentication/delivery retry loop. A failed - // handoff stays visible until a new event, explicit flush, or restart. - #[serde(skip)] - delivery_failed: bool, - #[serde(default)] - events: Vec, -} - -impl Default for PendingSession { - fn default() -> Self { - Self { - version: CORRELATION_VERSION, - linked_route: None, - standalone: false, - parent: None, - first_seen_ms: 0, - deadline: None, - delivery_failed: false, - events: Vec::new(), - } - } -} - -impl PendingSession { - fn restore(mut self) -> Self { - if self.version != CORRELATION_VERSION { - // A legacy link was decided before safe ancestry matching, but - // its earlier rows already live in that parent's trace. Keep later - // rows there too: splitting one session across two roots is worse. - // Undecided legacy buffers were never delivered; send them alone. - self.version = CORRELATION_VERSION; - self.standalone = self.linked_route.is_none(); - self.parent = None; - } - let remaining = self - .first_seen_ms - .saturating_add(CORRELATION_WAIT.as_millis() as i64) - .saturating_sub(now_ms()); - // A missing/future receipt timestamp is not permission to wait longer. - let remaining = if (0..=CORRELATION_WAIT.as_millis() as i64).contains(&remaining) { - remaining as u64 - } else { - 0 - }; - self.deadline = Some(Instant::now() + Duration::from_millis(remaining)); - self - } - - fn decided(&self) -> bool { - self.standalone || self.linked_route.is_some() - } - - fn expired(&self) -> bool { - self.deadline - .is_none_or(|deadline| Instant::now() >= deadline) - } -} - -#[derive(Clone, Serialize)] -struct PendingEvent { - env: Envelope, - replay_through: u64, - journal_through: u64, -} - -impl<'de> Deserialize<'de> for PendingEvent { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - #[derive(Deserialize)] - #[serde(untagged)] - enum StoredPendingEvent { - Current { - env: Envelope, - replay_through: u64, - journal_through: u64, - }, - Legacy(Envelope), - } - Ok(match StoredPendingEvent::deserialize(deserializer)? { - StoredPendingEvent::Current { - env, - replay_through, - journal_through, - } => Self { - env, - replay_through, - journal_through, - }, - StoredPendingEvent::Legacy(env) => Self { - env, - replay_through: 0, - journal_through: 0, - }, - }) - } -} - -enum IngressMsg { - Event(Box), - Barrier(oneshot::Sender<()>), -} - -/// The hook path never waits for this queue. Once it fills, the append-only -/// journals become the overflow queue and the worker catches them up in place. -const INGRESS_QUEUE_CAPACITY: usize = 64; - -/// 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 { - source: String, - session_id: String, - route: String, -} - -impl DeliveryKey { - fn new(source: &str, session_id: &str, route: &SessionRoute) -> anyhow::Result { - Ok(Self { - source: source.to_string(), - session_id: session_id.to_string(), - route: serde_json::to_string(route)?, - }) - } - - fn correlation_key(&self) -> String { - format!( - "{}\u{1f}{}\u{1f}{}", - self.source, self.session_id, self.route - ) - } -} - -#[derive(Clone)] -struct ClaudeTranscriptObservation { - /// Synthetic-event provenance only: never retain native prompts or responses. - template: Envelope, - through: Option, -} - -pub struct Daemon { - version: String, - data_dir: PathBuf, - translators: Arc, - sink_factory: Arc, - auth_provider: Option>, - session_auth: tokio::sync::Mutex>, - route_aliases: Mutex>, - /// Serializes only capture-side journal appends for a source session. - session_locks: Mutex>>>, - /// Serializes daemon-side routing and actor creation independently of - /// capture, which must never inherit actor or sink backpressure. - dispatch_locks: Mutex>>>, - journals: Mutex>>>, - claude_transcripts: Mutex>, - managed_run_sessions: Mutex>>, - auth_errors: Mutex>, - sessions: Mutex>>, - correlation: Arc, - automatic_links: Mutex>, - standalone_links: Mutex>, - pending_sessions: Mutex>, - pending_reconcile_lock: tokio::sync::Mutex<()>, - correlation_locks: Mutex>>>, - correlation_changed: Arc, - ingress_tx: mpsc::Sender, - ingress_overflow: AtomicBool, - ingress_dispatched: Mutex>, - started: Instant, - last_activity: Mutex, - /// Prevents a shutdown drain from racing capture that has passed the - /// quiescing check but has not yet appended and queued its event. - capture_gate: tokio::sync::RwLock<()>, - quiescing: AtomicBool, - drained: tokio::sync::Mutex, - shutting_down: AtomicBool, - shutdown: Notify, -} - -impl Daemon { - fn new(opts: ServeOptions, data_dir: PathBuf) -> Arc { - let (ingress_tx, ingress_rx) = mpsc::channel(INGRESS_QUEUE_CAPACITY); - let daemon = Arc::new(Daemon { - version: opts.version, - data_dir, - translators: opts.translators, - sink_factory: opts.sink_factory, - auth_provider: opts.auth_provider, - session_auth: tokio::sync::Mutex::new(HashMap::new()), - route_aliases: Mutex::new(HashMap::new()), - session_locks: Mutex::new(HashMap::new()), - dispatch_locks: Mutex::new(HashMap::new()), - journals: Mutex::new(HashMap::new()), - claude_transcripts: Mutex::new(HashMap::new()), - managed_run_sessions: Mutex::new(HashMap::new()), - auth_errors: Mutex::new(HashMap::new()), - sessions: Mutex::new(HashMap::new()), - correlation: Arc::new(crate::correlation::CorrelationRegistry::default()), - automatic_links: Mutex::new(HashMap::new()), - standalone_links: Mutex::new(HashSet::new()), - pending_sessions: Mutex::new(HashMap::new()), - pending_reconcile_lock: tokio::sync::Mutex::new(()), - correlation_locks: Mutex::new(HashMap::new()), - correlation_changed: Arc::new(Notify::new()), - ingress_tx, - ingress_overflow: AtomicBool::new(false), - ingress_dispatched: Mutex::new(HashMap::new()), - started: Instant::now(), - last_activity: Mutex::new(Instant::now()), - capture_gate: tokio::sync::RwLock::new(()), - quiescing: AtomicBool::new(false), - drained: tokio::sync::Mutex::new(false), - shutting_down: AtomicBool::new(false), - shutdown: Notify::new(), - }); - spawn_ingress_worker(daemon.clone(), ingress_rx); - 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 requested_route = env - .route - .clone() - .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; - if requested_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 requested_key = DeliveryKey::new(&env.source, &env.session_id, &requested_route)?; - let key = self - .route_aliases - .lock() - .unwrap() - .get(&requested_key) - .cloned() - .unwrap_or_else(|| requested_key.clone()); - let (selection, reason, expected_selection) = { - let states = self.session_auth.lock().await; - match states.get(&key) { - Some(state) => { - if !lease_is_expiring(&state.lease) { - env.config = Some(state.route.with_auth(state.lease.auth.clone())); - return Ok(key); - } - ( - state.lease.selection.clone(), - AuthResolveReason::Expiring, - Some(state.lease.selection.clone()), - ) - } - None => ( - requested_route.auth.clone(), - AuthResolveReason::Initial, - None, - ), - } - }; - - let lease = resolve_route_auth( - provider.as_ref(), - &selection, - reason, - requested_route.auth.org_name.as_deref(), - ) - .await - .map_err(|error| { - let message = format!( - "could not resolve Braintrust auth for {}: {error}; run `bt login` or select a profile explicitly", - env.source - ); - self.record_auth_error(&key, &env.source, message.clone()); - anyhow::anyhow!(message) - })?; - if let Some(expected) = expected_selection { - if lease.selection != expected { - anyhow::bail!( - "credential refresh changed auth selection from {expected:?} to {:?}", - lease.selection - ); - } - } - - let mut route = requested_route; - if route.auth.effective_source() == crate::wire::AuthSource::SavedProfile - && route.auth.profile_id.is_none() - && lease.selection.profile_id.is_some() - { - route.auth = lease.selection.clone(); - if let Err(error) = crate::settings::migrate_persisted_route( - &env.source, - env.route.as_ref().expect("route checked above"), - &route, - ) { - tracing::warn!(source = %env.source, "could not migrate legacy profile route: {error}"); - } - } - let canonical_key = DeliveryKey::new(&env.source, &env.session_id, &route)?; - if canonical_key != requested_key { - self.route_aliases - .lock() - .unwrap() - .insert(requested_key, canonical_key.clone()); - } - env.route = Some(route.clone()); - env.config = Some(route.with_auth(lease.auth.clone())); - self.session_auth - .lock() - .await - .insert(canonical_key.clone(), SessionAuthState { route, lease }); - self.auth_errors.lock().unwrap().remove(&canonical_key); - Ok(canonical_key) - } - - /// Keep a route's auth failure visible to `status.get` until it succeeds. - fn record_auth_error(&self, key: &DeliveryKey, source: &str, message: String) { - self.auth_errors - .lock() - .unwrap() - .insert(key.clone(), (source.to_string(), message)); - } - - /// Resolve a route's credentials exactly as event delivery would, so - /// diagnostics see this process's credential store and environment rather - /// than the caller's. - async fn diagnose_auth(&self, selection: AuthSelection) -> AuthDiagnoseResult { - let result = match &self.auth_provider { - Some(provider) => { - resolve_route_auth( - provider.as_ref(), - &selection, - AuthResolveReason::Initial, - selection.org_name.as_deref(), - ) - .await - } - None => Err(anyhow::anyhow!( - "daemon host has no Braintrust auth provider" - )), - }; - match result { - Ok(lease) => AuthDiagnoseResult { - selection: Some(lease.selection), - org_name: lease.auth.org_name, - expires_at_ms: lease.expires_at_ms, - error: None, - }, - Err(error) => AuthDiagnoseResult { - error: Some(error.to_string()), - ..AuthDiagnoseResult::default() - }, - } - } - - 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(key).cloned() else { - return Ok(()); - }; - if !lease_is_expiring(&state.lease) { - return Ok(()); - } - let selection = state.lease.selection.clone(); - let lease = provider - .resolve(&selection, AuthResolveReason::Expiring) - .await?; - if lease.selection != state.lease.selection { - anyhow::bail!("credential refresh changed the session auth selection"); - } - let config = state.route.with_auth(lease.auth.clone()); - self.session_auth.lock().await.insert( - key.clone(), - SessionAuthState { - route: state.route, - lease, - }, - ); - let session = { self.sessions.lock().unwrap().get(key).cloned() }; - if let Some(session) = session { - session.configure(config).await?; - } - Ok(()) - } - - fn touch(&self) { - *self.last_activity.lock().unwrap() = Instant::now(); - } - - fn correlation_lock(&self, key: &str) -> Arc> { - let mut locks = self.correlation_locks.lock().unwrap(); - locks - .entry(key.to_string()) - .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) - .clone() - } - - async fn session_for( - &self, - env: &Envelope, - key: &DeliveryKey, - replay_through: u64, - ) -> anyhow::Result> { - { - let map = self.sessions.lock().unwrap(); - if let Some(session) = map.get(key) { - return Ok(session.clone()); - } - } - let route = env - .route - .as_ref() - .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; - let config = env.config.clone().ok_or_else(|| { - anyhow::anyhow!("resolved route is missing its session configuration") - })?; - // The actor streams the journal itself, so creating a session stays - // cheap and allocation-free here no matter how long the recorded - // session is. Bound it to what is recorded now, before this event is - // appended, so replay covers prior source observations only. Every - // destination consumes that shared observation stream, with its own - // delivery checkpoint controlling which rows reach its sink. - let journal_path = - journal::ensure_source_journal(&self.data_dir, &env.source, &env.session_id).await?; - let translator_session_id = - if journal::legacy_journal_has_session(&self.data_dir, &env.source, &env.session_id) - .await - { - env.session_id.clone() - } else { - crate::ids::session_namespace(&env.source, &env.session_id) - }; - let replay = ReplayPlan { - acknowledged_through: journal::JournalReader::acknowledged_through( - &journal_path, - replay_through, - route, - ) - .await, - through: replay_through, - journal_path, - }; - let journal = self - .journal_writer_for(&env.source, &env.session_id) - .await?; - let mut map = self.sessions.lock().unwrap(); - if let Some(s) = map.get(key) { - return Ok(s.clone()); - } - let session = Session::spawn( - SessionOptions { - session_id: env.session_id.clone(), - translator_session_id, - source: env.source.clone(), - plugin_version: env.plugin_version.clone(), - replay: Some(replay), - config, - correlation_key: key.correlation_key(), - route: route.clone(), - correlation: self.correlation.clone(), - data_dir: self.data_dir.clone(), - journal, - correlation_changed: self.correlation_changed.clone(), - }, - self.translators.clone(), - self.sink_factory.clone(), - ); - map.insert(key.clone(), session.clone()); - Ok(session) - } - - /// Drop every trace of one delivery pipeline. A session that goes quiet - /// must not pin its translator state, sink handles, credential lease, or - /// journal file for the rest of the daemon's life; deterministic span ids - /// mean a late event simply rebuilds it from the journal. - async fn retire_session(&self, key: &DeliveryKey) { - self.observe_claude_transcripts(Some(key), false).await; - self.ingress_barrier().await; - let lock = self.dispatch_lock(&key.source, &key.session_id); - let _guard = lock.lock().await; - - let session = { self.sessions.lock().unwrap().remove(key) }; - let Some(session) = session else { - return; - }; - session.shutdown().await; - self.correlation.retire_session(&key.correlation_key()); - - self.session_auth.lock().await.remove(key); - self.auth_errors.lock().unwrap().remove(key); - self.route_aliases - .lock() - .unwrap() - .retain(|alias, target| alias != key && target != key); - self.managed_run_sessions.lock().unwrap().retain(|_, keys| { - keys.remove(key); - !keys.is_empty() - }); - - // Several delivery routes can share one source session, so release - // its journal writer and lock only after its final route retires. - let last = !self - .sessions - .lock() - .unwrap() - .keys() - .any(|other| other.source == key.source && other.session_id == key.session_id); - if last { - let storage_key = crate::ids::session_namespace(&key.source, &key.session_id); - // Capture can proceed while the actor flushes above. Only take its - // lock for the brief writer-map cleanup after daemon work ends. - let capture_lock = self.session_lock(&key.source, &key.session_id); - let _capture_guard = capture_lock.lock().await; - self.claude_transcripts - .lock() - .unwrap() - .retain(|candidate, _| { - candidate.source != key.source || candidate.session_id != key.session_id - }); - self.journals.lock().unwrap().remove(&storage_key); - { - let mut locks = self.session_locks.lock().unwrap(); - // A waiting capture or observer must keep using this same - // lock, rather than race a newly allocated mirror writer. - if Arc::strong_count(&capture_lock) == 2 { - locks.remove(&storage_key); - } - } - self.dispatch_locks.lock().unwrap().remove(&storage_key); - self.ingress_dispatched - .lock() - .unwrap() - .retain(|candidate, _| { - candidate.source != key.source || candidate.session_id != key.session_id - }); - } - tracing::info!(session_id = %key.session_id, "session retired"); - } - - /// Delivery pipelines with no traffic for `idle_timeout` and nothing left - /// queued. - fn idle_sessions(&self, idle_timeout: Duration) -> Vec { - self.sessions - .lock() - .unwrap() - .iter() - .filter(|(_, session)| { - session.idle_for() >= idle_timeout - && session.counters.queued.load(Ordering::Relaxed) == 0 - }) - .filter(|(key, _)| !self.correlation.has_active_tools(&key.correlation_key())) - .map(|(key, _)| key.clone()) - .collect() - } - - async fn journal_writer_for( - &self, - source: &str, - session_id: &str, - ) -> anyhow::Result>> { - let storage_key = crate::ids::session_namespace(source, session_id); - if let Some(writer) = self.journals.lock().unwrap().get(&storage_key).cloned() { - return Ok(writer); - } - let path = journal::ensure_source_journal(&self.data_dir, source, session_id).await?; - let writer = Arc::new(tokio::sync::Mutex::new( - JournalWriter::open_path(&path).await?, - )); - Ok(self - .journals - .lock() - .unwrap() - .entry(storage_key) - .or_insert_with(|| writer.clone()) - .clone()) - } - - async fn append_to_journal(&self, env: &mut Envelope) -> anyhow::Result<(u64, u64)> { - let writer = self - .journal_writer_for(&env.source, &env.session_id) - .await?; - let mut writer = writer.lock().await; - let before = writer.position(); - let through = writer.append(env).await?; - Ok((before, through)) - } - - fn session_lock(&self, source: &str, session_id: &str) -> Arc> { - let storage_key = crate::ids::session_namespace(source, session_id); - self.session_locks - .lock() - .unwrap() - .entry(storage_key) - .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) - .clone() - } - - fn dispatch_lock(&self, source: &str, session_id: &str) -> Arc> { - let storage_key = crate::ids::session_namespace(source, session_id); - self.dispatch_locks - .lock() - .unwrap() - .entry(storage_key) - .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) - .clone() - } - - fn correlation_decided(&self, key: &str) -> bool { - self.standalone_links.lock().unwrap().contains(key) - || self.automatic_links.lock().unwrap().contains_key(key) - } - - /// Queued work, excluding pending sessions whose handoff failed. Those are - /// durable on disk and resume after a restart, so they must not keep an - /// otherwise idle daemon alive forever. - fn total_queued(&self) -> u64 { - let dispatched: u64 = self - .sessions - .lock() - .unwrap() - .values() - .map(|s| s.counters.queued.load(Ordering::Relaxed)) - .sum(); - dispatched.saturating_add( - self.pending_sessions - .lock() - .unwrap() - .values() - .filter(|state| !state.delivery_failed) - .map(|state| state.events.len() as u64) - .sum(), - ) - } - - async fn capture_event(&self, mut env: Envelope) -> Result { - let _capture_guard = self.capture_gate.read().await; - if self.quiescing.load(Ordering::SeqCst) { - return Ok(false); - } - env.source = self - .translators - .canonical_source(&env.source) - .ok_or_else(|| format!("unsupported coding-agent source {:?}", env.source))? - .to_string(); - self.touch(); - let lock = self.session_lock(&env.source, &env.session_id); - let _guard = lock.lock().await; - hydrate_transcript_reference(&self.data_dir, &mut env).await; - self.journal_and_enqueue(env).await - } - - /// Caller holds the source-session capture lock; frozen observations enter - /// the same durable queue as hooks without reacquiring that lock. - async fn journal_and_enqueue(&self, mut env: Envelope) -> Result { - let (replay_through, journal_through) = self - .append_to_journal(&mut env) - .await - .map_err(|error| format!("journal failed: {error}"))?; - self.remember_claude_transcripts(&env); - match self - .ingress_tx - .try_send(IngressMsg::Event(Box::new(PendingEvent { - env, - replay_through, - journal_through, - }))) { - Ok(()) => {} - Err(mpsc::error::TrySendError::Full(_)) => { - self.ingress_overflow.store(true, Ordering::Release); - tracing::debug!("ingress queue full; journal will be drained by daemon worker"); - } - Err(mpsc::error::TrySendError::Closed(_)) => { - tracing::warn!("journaled event will be recovered after daemon restart"); - } - } - Ok(true) - } - - fn remember_claude_transcripts(&self, env: &Envelope) { - if env.source != "claude-code" || env.payload.get("_bt_import_through_offset").is_some() { - return; - } - if env.event == "SessionEnd" { - self.claude_transcripts - .lock() - .unwrap() - .retain(|key, _| key.source != env.source || key.session_id != env.session_id); - return; - } - let Some(route) = env.route.as_ref() else { - return; - }; - let Some(path) = env.payload.get("transcript_path").and_then(Value::as_str) else { - return; - }; - let Ok(key) = DeliveryKey::new(&env.source, &env.session_id, route) else { - return; - }; - let through = env - .payload - .get("_bt_claude_transcript_mirrors") - .and_then(|mirrors| mirrors.get(path)) - .or_else(|| { - env.payload - .get("_bt_transcript_mirror") - .filter(|mirror| mirror.get("path").and_then(Value::as_str) == Some(path)) - }) - .and_then(|mirror| mirror.get("through")) - .and_then(Value::as_u64); - let mut payload = serde_json::json!({"transcript_path": path}); - if let Some(cwd) = env.payload.get("cwd").and_then(Value::as_str) { - payload["cwd"] = Value::String(cwd.to_string()); - } - let template = Envelope { - source: env.source.clone(), - source_version: env.source_version.clone(), - plugin_version: env.plugin_version.clone(), - session_id: env.session_id.clone(), - event: "TranscriptUpdate".to_string(), - ts_ms: env.ts_ms, - managed_run_id: env.managed_run_id.clone(), - capture: None, - payload, - route: env.route.clone(), - config: None, - }; - let mut observations = self.claude_transcripts.lock().unwrap(); - let previous = observations - .get(&key) - .filter(|observation| { - observation - .template - .payload - .get("transcript_path") - .and_then(Value::as_str) - == Some(path) - }) - .and_then(|observation| observation.through); - observations.insert( - key, - ClaudeTranscriptObservation { - template, - through: through.or(previous), - }, - ); - } - - /// Poll native paths only on the capture side. Each observed boundary is - /// journaled before translation, including observations made during shutdown. - async fn observe_claude_transcripts(&self, only: Option<&DeliveryKey>, quiesced: bool) { - // Shutdown already owns the exclusive gate. Ordinary polling must not - // append after that final drain begins. - let _gate = if quiesced { - None - } else { - Some(self.capture_gate.read().await) - }; - if !quiesced && self.quiescing.load(Ordering::SeqCst) { - return; - } - let keys: Vec<_> = self - .claude_transcripts - .lock() - .unwrap() - .keys() - .filter(|key| { - only.is_none_or(|only| { - key.source == only.source && key.session_id == only.session_id - }) - }) - .cloned() - .collect(); - for key in keys { - // Routes share mirror files, so serialize by source session, not - // destination. Never hold this lock across an ingress barrier. - let lock = self.session_lock(&key.source, &key.session_id); - let _guard = lock.lock().await; - let observation = self.claude_transcripts.lock().unwrap().get(&key).cloned(); - let Some(observation) = observation else { - continue; - }; - let path = observation.template.payload["transcript_path"] - .as_str() - .unwrap(); - let Ok(metadata) = tokio::fs::metadata(path).await else { - continue; - }; - if observation.through == Some(metadata.len()) { - continue; - } - let mut env = observation.template; - env.ts_ms = now_ms(); - hydrate_transcript_reference(&self.data_dir, &mut env).await; - let through = env - .payload - .get("_bt_claude_transcript_mirrors") - .and_then(|mirrors| mirrors.get(env.payload["transcript_path"].as_str().unwrap())) - .and_then(|mirror| mirror.get("through")) - .and_then(Value::as_u64); - if through.is_none() || through == observation.through { - continue; - } - if let Err(error) = self.journal_and_enqueue(env).await { - tracing::warn!(session_id = %key.session_id, %error, "Claude transcript observation failed"); - } - } - } - - async fn ingress_barrier(&self) { - let (tx, rx) = oneshot::channel(); - if self.ingress_tx.send(IngressMsg::Barrier(tx)).await.is_ok() { - let _ = rx.await; - } - } - - async fn settle_ingress(self: &Arc) { - self.ingress_barrier().await; - self.settle_session_actors().await; - let _ = retry_pending_sessions(self).await; - } - - async fn settle_session_actors(&self) { - let sessions: Vec<_> = self.sessions.lock().unwrap().values().cloned().collect(); - for session in sessions { - session.barrier().await; - } - } - - fn record_managed_run_session(&self, managed_run_id: &str, key: &DeliveryKey) -> bool { - self.managed_run_sessions - .lock() - .unwrap() - .entry(managed_run_id.to_string()) - .or_default() - .insert(key.clone()) - } - - 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 { - source: Some(key.source.clone()), - 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 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 Some(source) = record.source { - if let Ok(key) = DeliveryKey::new(&source, &record.session_id, &record.route) { - delivery_keys.insert(key); - } - } else { - // Legacy records predate source-qualified delivery keys. A - // matching live pipeline is authoritative; otherwise recover - // the source from the legacy journal. - let live: Vec<_> = self - .sessions - .lock() - .unwrap() - .keys() - .filter(|key| { - key.session_id == record.session_id - && serde_json::from_str::(&key.route) - .is_ok_and(|route| route.same_route(&record.route)) - }) - .cloned() - .collect(); - if live.is_empty() { - if let Some(source) = journal::legacy_journal_source( - &self.data_dir, - &record.session_id, - &record.route, - ) - .await - .and_then(|source| { - self.translators - .canonical_source(&source) - .map(str::to_owned) - }) { - if let Ok(key) = - DeliveryKey::new(&source, &record.session_id, &record.route) - { - delivery_keys.insert(key); - } - } - } else { - delivery_keys.extend(live); - } - } - } - let mut result = FlushResult { - flushed: true, - pending: 0, - accepted_sessions: delivery_keys.len() as u64, - }; - for key in &delivery_keys { - self.observe_claude_transcripts(Some(key), false).await; - } - // Managed completion is terminal. Stop polling before the barrier so - // an in-flight observer cannot enqueue after translator finalization. - for key in &delivery_keys { - let lock = self.session_lock(&key.source, &key.session_id); - let _guard = lock.lock().await; - self.claude_transcripts - .lock() - .unwrap() - .retain(|candidate, observation| { - candidate.source != key.source - || candidate.session_id != key.session_id - || observation.template.managed_run_id.as_deref() - != Some(params.managed_run_id.as_str()) - }); - } - self.ingress_barrier().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 = %key.session_id, - %error, - "managed run session auth refresh failed" - ); - result.flushed = false; - continue; - } - let session = { self.sessions.lock().unwrap().get(&key).cloned() }; - if let Some(session) = session { - let (flushed, pending) = session - .finalize(Duration::from_millis(params.timeout_ms)) - .await; - result.flushed &= flushed; - result.pending = result.pending.saturating_add(pending); - } - } - result - } - - /// Queue a flush on every live delivery route before dispatching more ingress. - /// Only the completion wait runs in a separate task. This is used - /// only by the daemon worker after a turn-ending event has already been - /// durably captured and acknowledged to the hook client. - async fn enqueue_source_session_flush( - &self, - source: &str, - session_id: &str, - timeout: Duration, - ) { - let sessions: Vec<_> = self - .sessions - .lock() - .unwrap() - .iter() - .filter(|(key, _)| key.source == source && key.session_id == session_id) - .map(|(_, session)| session.clone()) - .collect(); - for session in sessions { - let reply = match session.enqueue_flush().await { - Ok(reply) => reply, - Err(error) => { - tracing::warn!(source, session_id, %error, "turn-end flush enqueue failed"); - continue; - } - }; - let source = source.to_owned(); - let session_id = session_id.to_owned(); - tokio::spawn(async move { - // Later events may remain queued after this boundary has flushed. - // Their presence does not make this boundary's flush incomplete. - if !matches!(tokio::time::timeout(timeout, reply).await, Ok(Ok(_))) { - tracing::warn!( - source, - session_id, - "out-of-band turn-end flush did not complete" - ); - } - }); - } - } - - fn trigger_shutdown(&self) { - self.quiescing.store(true, Ordering::SeqCst); - self.shutting_down.store(true, Ordering::SeqCst); - self.shutdown.notify_waiters(); - } - - fn begin_quiesce(&self) { - self.quiescing.store(true, Ordering::SeqCst); - } -} - -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) -} - -fn spawn_ingress_worker(daemon: Arc, mut rx: mpsc::Receiver) { - tokio::spawn(async move { - loop { - if daemon.ingress_overflow.swap(false, Ordering::AcqRel) { - recover_ingress_overflow(&daemon).await; - continue; - } - let Some(msg) = rx.recv().await else { - break; - }; - match msg { - IngressMsg::Event(event) => { - dispatch_ingress_event(&daemon, *event).await; - } - IngressMsg::Barrier(reply) => { - loop { - daemon.settle_session_actors().await; - let _ = retry_pending_sessions(&daemon).await; - if !daemon.ingress_overflow.swap(false, Ordering::AcqRel) { - break; - } - recover_ingress_overflow(&daemon).await; - } - let _ = reply.send(()); - } - } - } - }); -} - -async fn dispatch_ingress_event(daemon: &Arc, event: PendingEvent) { - let ingress_key = - event.env.route.as_ref().and_then(|route| { - DeliveryKey::new(&event.env.source, &event.env.session_id, route).ok() - }); - if ingress_key.as_ref().is_some_and(|key| { - daemon - .ingress_dispatched - .lock() - .unwrap() - .get(key) - .is_some_and(|through| *through >= event.journal_through) - }) { - return; - } - - let journal_through = event.journal_through; - // A child may arrive immediately after a parent's tool hook. Catch prior - // session actors up before resolving a new session, entirely in the daemon. - // Only an undecided session resolves a parent. Antigravity marks every - // invocation as a start, so skip the all-session barrier once decided. - if is_session_start(&event.env.source, &event.env.event) - && !daemon.correlation_decided(&automatic_link_key(&event.env)) - { - daemon.settle_session_actors().await; - } - match accept_event(daemon, event).await { - Ok(()) => { - if let Some(key) = ingress_key { - let mut dispatched = daemon.ingress_dispatched.lock().unwrap(); - let through = dispatched.entry(key).or_default(); - *through = (*through).max(journal_through); - } - } - Err(error) => { - tracing::warn!(%error, "journaled ingress event could not be dispatched"); - } - } -} - -/// Drain events omitted from the bounded in-memory queue. The queue contains -/// only a latency fast path; journals remain the complete source of ingress. -async fn recover_ingress_overflow(daemon: &Arc) { - let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { - return; - }; - let mut recovered = 0usize; - while let Ok(Some(entry)) = entries.next_entry().await { - let path = entry.path(); - if path.extension().and_then(|ext| ext.to_str()) != Some("ndjson") { - continue; - } - let recorded_len = journal::JournalReader::recorded_len(&path).await; - let Ok(Some(mut checkpoints)) = journal::JournalReader::open(&path, recorded_len).await - else { - continue; - }; - let mut acknowledged_by_route: HashMap = HashMap::new(); - while let Ok(Some(entry)) = checkpoints.next_record().await { - if let journal::JournalRecord::DeliveryCheckpoint { route, through } = entry.record { - let key = serde_json::to_string(&route).unwrap_or_default(); - let acknowledged = acknowledged_by_route.entry(key).or_default(); - *acknowledged = (*acknowledged).max(through); - } - } - - let Ok(Some(mut reader)) = journal::JournalReader::open(&path, recorded_len).await else { - continue; - }; - let mut before = 0u64; - while let Ok(Some(entry)) = reader.next_record().await { - let through = entry.through; - let journal::JournalRecord::Event(redacted) = entry.record else { - before = through; - continue; - }; - let mut env = journal::envelope_from_redacted(redacted); - let Some(source) = daemon.translators.canonical_source(&env.source) else { - before = through; - continue; - }; - env.source = source.to_string(); - let Some(route) = env.route.as_ref() else { - before = through; - continue; - }; - let route_json = recovered_delivery_route(daemon, &env) - .await - .as_ref() - .and_then(|route| serde_json::to_string(route).ok()) - .unwrap_or_else(|| serde_json::to_string(route).unwrap_or_default()); - let key = match DeliveryKey::new(&env.source, &env.session_id, route) { - Ok(key) => key, - Err(_) => { - before = through; - continue; - } - }; - let dispatched = daemon - .ingress_dispatched - .lock() - .unwrap() - .get(&key) - .copied() - .unwrap_or(0); - let acknowledged = acknowledged_by_route.get(&route_json).copied().unwrap_or(0); - if through > dispatched.max(acknowledged) { - dispatch_ingress_event( - daemon, - PendingEvent { - env, - replay_through: before, - journal_through: through, - }, - ) - .await; - recovered += 1; - } - before = through; - } - } - if recovered > 0 { - tracing::info!(recovered, "drained journal-backed ingress overflow"); - } -} - -fn spawn_pending_reconciler(daemon: Arc) { - tokio::spawn(async move { - let mut tick = tokio::time::interval(Duration::from_millis(100)); - loop { - tokio::select! { - _ = daemon.shutdown.notified() => return, - _ = daemon.correlation_changed.notified() => {}, - _ = tick.tick() => {}, - } - daemon.observe_claude_transcripts(None, false).await; - if let Err(error) = retry_pending_sessions(&daemon).await { - tracing::warn!(%error, "pending child-session reconciliation failed"); - } - } - }); -} - -/// 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<()> { - let socket = paths::socket_path(args.socket.as_deref()); - let data_dir = paths::data_dir(args.data_dir.as_deref()); - paths::ensure_private_dir(&data_dir)?; - #[cfg(unix)] - if let Some(parent) = socket.parent() { - paths::ensure_private_dir(parent)?; - } - - let listener = match transport::claim(&socket, || probe_alive(&socket)).await? { - Some(l) => l, - None => { - tracing::info!( - "another daemon is already serving {}; exiting", - socket.display() - ); - return Ok(()); - } - }; - tracing::info!(socket = %socket.display(), "bt-daemon listening"); - - let daemon = Daemon::new(opts, data_dir); - collect_garbage(&daemon.data_dir).await; - restore_active_parent_snapshots(&daemon.data_dir, &daemon.correlation).await; - restore_pending_sessions(&daemon).await; - recover_unprocessed_journals(&daemon).await; - spawn_pending_reconciler(daemon.clone()); - let idle_timeout = Duration::from_secs(args.idle_timeout_secs); - spawn_idle_watchdog(daemon.clone(), idle_timeout); - spawn_session_reaper( - daemon.clone(), - Duration::from_secs(args.session_idle_timeout_secs), - ); - spawn_gc(daemon.clone()); - - let accept_result = accept_loop(daemon.clone(), listener).await; - - // Graceful drain regardless of why we stopped. - drain_all(&daemon).await; - transport::cleanup(&socket); - accept_result -} - -async fn accept_loop(daemon: Arc, mut listener: Listener) -> anyhow::Result<()> { - loop { - // `Notify::notify_waiters` does not retain a permit. If accepting a - // connection wins the select at the same time shutdown is requested, - // check the sticky flag before waiting again so the notification - // cannot be lost. - if daemon.shutting_down.load(Ordering::SeqCst) { - tracing::info!("shutdown requested"); - return Ok(()); - } - tokio::select! { - _ = daemon.shutdown.notified() => { - tracing::info!("shutdown requested"); - return Ok(()); - } - _ = tokio::signal::ctrl_c() => { - tracing::info!("interrupt received"); - return Ok(()); - } - accepted = listener.accept() => { - match accepted { - Ok(stream) => { - let d = daemon.clone(); - tokio::spawn(async move { - if let Err(e) = serve_connection(d, stream).await { - tracing::debug!("connection ended: {e}"); - } - }); - } - Err(e) => { - tracing::warn!("accept error: {e}"); - } - } - } - } - } -} - -async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow::Result<()> { - let (read_half, mut write_half) = tokio::io::split(stream); - let mut lines = BufReader::new(read_half).lines(); - let mut client = None; - - while let Some(line) = lines.next_line().await? { - if line.trim().is_empty() { - continue; - } - let mut shutdown_after_response = false; - let response = match Message::from_line(&line) { - Ok(Message::Request(req)) => { - let request_id = req.id.clone(); - let method = req.method.clone(); - tracing::info!( - request_id = ?request_id, - method, - "request received" - ); - let response = handle_request(&daemon, req, &mut client).await; - shutdown_after_response = method == method::DAEMON_SHUTDOWN - && response.error.is_none() - && response - .result - .as_ref() - .and_then(|result| result.get("ok")) - .and_then(Value::as_bool) - == Some(true); - if let Some(error) = &response.error { - tracing::warn!( - request_id = ?request_id, - method, - error_code = error.code, - error = %error.message, - "request failed" - ); - } else { - tracing::info!( - request_id = ?request_id, - method, - "request completed" - ); - } - Some(response) - } - Ok(Message::Notification(note)) => { - tracing::info!(method = %note.method, "notification received"); - // Legacy best-effort notifications cannot participate in the - // acknowledged handover retry used by current clients. - if note.method == method::EVENT_LOG { - if let Some(params) = note.params { - match serde_json::from_value::(params) { - Ok(mut env) => { - attach_process_capture(&daemon, &mut env, client.as_ref()); - match daemon.capture_event(env).await { - Ok(true) => {} - Ok(false) => tracing::warn!( - "event notification arrived while daemon was shutting down" - ), - Err(error) => { - tracing::warn!(%error, "event notification was not captured") - } - } - } - Err(error) => tracing::warn!( - method = %note.method, - error = %error, - "notification parameters rejected" - ), - } - } - } - None - } - Ok(Message::Response(_)) => None, // clients don't send us responses - Err(e) => { - tracing::warn!(error = %e, "request parse failed"); - Some(Response::err( - crate::wire::RequestId::Int(0), - RpcError::new(error_code::PARSE, format!("parse error: {e}")), - )) - } - }; - - if let Some(resp) = response { - let mut buf = Message::Response(resp).to_line()?; - buf.push('\n'); - let write_result = async { - write_half.write_all(buf.as_bytes()).await?; - write_half.flush().await - } - .await; - if shutdown_after_response { - daemon.trigger_shutdown(); - } - write_result?; - } - } - Ok(()) -} - -fn attach_process_capture( - daemon: &Daemon, - env: &mut Envelope, - client: Option<&crate::wire::ClientInfo>, -) { - if env.capture.is_some() { - return; - } - if !is_session_start(&env.source, &env.event) - && !daemon - .correlation - .needs_process_capture(&env.source, &env.session_id) - { - return; - } - let Some(pid) = client.and_then(|client| client.pid) else { - return; - }; - let capture = crate::process::capture_process_context(pid); - if !capture.process_chain.is_empty() || capture.truncated { - env.capture = Some(capture); - } -} - -async fn accept_event(daemon: &Arc, event: PendingEvent) -> Result<(), String> { - let key = automatic_link_key(&event.env); - let lock = daemon.correlation_lock(&key); - let _guard = lock.lock().await; - let mut state = daemon.pending_sessions.lock().unwrap().get(&key).cloned(); - if state.is_none() { - if daemon.standalone_links.lock().unwrap().contains(&key) { - return accept_resolved_event(daemon, event).await; - } - let route = daemon.automatic_links.lock().unwrap().get(&key).cloned(); - if let Some(route) = route { - let mut event = event; - event.env.route = Some(route); - return accept_resolved_event(daemon, event).await; - } - state = read_correlation_state(&daemon.data_dir, &key).await; - } - - if let Some(mut state) = state { - if state.decided() { - // Queue behind any earlier events whose handoff failed so a - // second failure keeps this event visible and retryable too. - state.events.push(event); - return finish_correlation(daemon, &key, state).await; - } - - // A new native start is a different process generation, not a second - // observation confirming the previous hook's ancestry. - let new_generation = is_session_start(&event.env.source, &event.env.event); - state.events.push(event); - if !state.expired() && !new_generation { - let mut captures = state - .events - .iter() - .rev() - .filter_map(|event| event.env.capture.as_ref()); - if let (Some(latest), Some(previous), Some(parent)) = - (captures.next(), captures.next(), state.parent.as_ref()) - { - if let Some(agent) = crate::correlation::session_agent_process( - &state.events[0].env.source, - previous, - Some(latest), - ) { - match crate::correlation::CorrelationRegistry::confirm_parent( - parent, latest, &agent, - ) { - crate::correlation::Resolution::Parent(parent) => { - state.linked_route = Some(parent.route); - } - _ => state.standalone = true, - } - } - } - } - if !state.decided() - && (new_generation || state.expired() || state.events.len() >= CORRELATION_EVENT_LIMIT) - { - state.standalone = true; - } - if state.decided() { - return finish_correlation(daemon, &key, state).await; - } - return store_correlation(daemon, &key, &state).await; - } - - let mut state = PendingSession::default(); - // An explicit parent is authoritative. Missing first-start/process evidence - // also means standalone, permanently; a later start cannot reparent rows - // already delivered. - if is_session_start(&event.env.source, &event.env.event) - && !matches!( - event - .env - .route - .as_ref() - .and_then(|route| route.destination.as_ref()), - Some(crate::wire::TraceDestination::ParentSpan { .. }) - ) - { - let child_agent = event.env.capture.as_ref().and_then(|capture| { - crate::correlation::session_agent_process(&event.env.source, capture, None) - }); - let session_prefix = format!("{}\u{1f}{}\u{1f}", event.env.source, event.env.session_id); - match daemon.correlation.resolve( - Some(&session_prefix), - event.env.capture.as_ref(), - child_agent.as_ref(), - |parent| automatic_parent_route(event.env.route.as_ref(), parent), - ) { - crate::correlation::Resolution::Parent(parent) => { - state.linked_route = Some(parent.route); - } - crate::correlation::Resolution::Pending(parent) => { - state.parent = Some(*parent); - state.first_seen_ms = now_ms(); - state.deadline = Some(Instant::now() + CORRELATION_WAIT); - } - crate::correlation::Resolution::Standalone => {} - } - } - state.standalone = state.linked_route.is_none() && state.parent.is_none(); - state.events.push(event); - if state.decided() { - finish_correlation(daemon, &key, state).await - } else { - store_correlation(daemon, &key, &state).await - } -} - -/// An inferred parent may change hierarchy, never the requested credential or -/// destination. Preserve the child's transforms, metadata, tags and flush mode. -fn automatic_parent_route( - child: Option<&SessionRoute>, - parent: &SessionRoute, -) -> Option { - use crate::wire::TraceDestination; - use braintrust_sdk_rust::SpanObjectType; - - let child = child?; - let same_profile = match (&child.auth.profile_id, &parent.auth.profile_id) { - (Some(child), Some(parent)) => child == parent, - _ => child.auth.profile == parent.auth.profile, - }; - if child.auth.effective_source() != parent.auth.effective_source() - || !same_profile - || child.auth.org_name != parent.auth.org_name - { - return None; - } - let Some(TraceDestination::ParentSpan { components }) = &parent.destination else { - return None; - }; - let same_destination = match child.destination.as_ref()? { - TraceDestination::ProjectLogs { - project_id, - project_name, - } => { - if components.object_type != SpanObjectType::ProjectLogs { - false - } else if let Some(id) = project_id { - components.object_id.as_ref() == Some(id) - } else { - project_name.as_deref().is_some_and(|name| { - components - .compute_object_metadata_args - .as_ref() - .and_then(|args| args.get("project_name")) - .and_then(Value::as_str) - == Some(name) - }) - } - } - TraceDestination::Experiment { experiment_id } => { - components.object_type == SpanObjectType::Experiment - && components.object_id.as_ref() == Some(experiment_id) - } - TraceDestination::ParentSpan { .. } => false, - }; - if !same_destination { - return None; - } - let mut route = child.clone(); - route.destination = parent.destination.clone(); - Some(route) -} - -async fn store_correlation( - daemon: &Daemon, - key: &str, - state: &PendingSession, -) -> Result<(), String> { - // Keep accepted events visible/retryable even if persisting the decision - // fails. The immutable raw journal remains authoritative for recovery. - if !state.events.is_empty() { - daemon - .pending_sessions - .lock() - .unwrap() - .insert(key.to_string(), state.clone()); - } - write_correlation_state(&daemon.data_dir, key, state).await?; - if state.events.is_empty() { - daemon.pending_sessions.lock().unwrap().remove(key); - if let Some(route) = &state.linked_route { - daemon - .automatic_links - .lock() - .unwrap() - .insert(key.to_string(), route.clone()); - } else if state.standalone { - daemon - .standalone_links - .lock() - .unwrap() - .insert(key.to_string()); - } - } - Ok(()) -} - -async fn finish_correlation( - daemon: &Arc, - key: &str, - mut state: PendingSession, -) -> Result<(), String> { - state.parent = None; - state.delivery_failed = false; - let result = async { - // Persist the final route before delivering any events. A crash may - // replay deterministic rows, but must never select a different parent. - store_correlation(daemon, key, &state).await?; - while let Some(event) = state.events.first() { - let mut event = event.clone(); - if let Some(route) = &state.linked_route { - event.env.route = Some(route.clone()); - } - accept_resolved_event(daemon, event).await?; - state.events.remove(0); - store_correlation(daemon, key, &state).await?; - } - Ok(()) - } - .await; - if result.is_err() { - state.delivery_failed = true; - daemon - .pending_sessions - .lock() - .unwrap() - .insert(key.to_string(), state); - } - result -} - -async fn retry_pending_sessions(daemon: &Arc) -> Result<(), String> { - resolve_pending_sessions(daemon, |_| false).await -} - -async fn resolve_pending_sessions( - daemon: &Arc, - force: impl Fn(&Envelope) -> bool + Send + Sync, -) -> Result<(), String> { - let _reconcile_guard = daemon.pending_reconcile_lock.lock().await; - let keys: Vec<_> = daemon - .pending_sessions - .lock() - .unwrap() - .keys() - .cloned() - .collect(); - let mut failure = None; - for key in keys { - let lock = daemon.correlation_lock(&key); - let _guard = lock.lock().await; - let Some(mut state) = daemon.pending_sessions.lock().unwrap().get(&key).cloned() else { - continue; - }; - let forced = state.events.first().is_some_and(|event| force(&event.env)); - if state.delivery_failed && !forced { - continue; - } - if !state.decided() && (state.expired() || forced) { - state.standalone = true; - } - if state.decided() { - if let Err(error) = finish_correlation(daemon, &key, state).await { - // One unavailable destination must not prevent other sessions - // from leaving the discovery window. - failure = Some(error); - } - } - } - failure.map_or(Ok(()), Err) -} - -fn correlation_state_path(data_dir: &std::path::Path, key: &str) -> PathBuf { - let digest = Sha256::digest(key.as_bytes()); - data_dir - .join("correlation") - .join(format!("{digest:x}.json")) -} - -async fn read_correlation_state(data_dir: &std::path::Path, key: &str) -> Option { - let bytes = tokio::fs::read(correlation_state_path(data_dir, key)) - .await - .ok()?; - serde_json::from_slice::(&bytes) - .ok() - .map(PendingSession::restore) -} - -async fn write_correlation_state( - data_dir: &std::path::Path, - key: &str, - state: &PendingSession, -) -> Result<(), String> { - let dir = data_dir.join("correlation"); - tokio::fs::create_dir_all(&dir) - .await - .map_err(|error| format!("correlation journal directory failed: {error}"))?; - let path = correlation_state_path(data_dir, key); - let temp = path.with_extension(format!("{}.tmp", uuid::Uuid::new_v4())); - let bytes = serde_json::to_vec(state) - .map_err(|error| format!("correlation journal encoding failed: {error}"))?; - tokio::fs::write(&temp, bytes) - .await - .map_err(|error| format!("correlation journal write failed: {error}"))?; - if let Err(first_error) = tokio::fs::rename(&temp, &path).await { - // Windows cannot replace an existing destination with rename. - let _ = tokio::fs::remove_file(&path).await; - tokio::fs::rename(&temp, &path).await.map_err(|error| { - format!("correlation journal replace failed: {first_error}; retry failed: {error}") - })?; - } - Ok(()) -} - -fn active_parent_snapshot_path(data_dir: &std::path::Path, key: &str) -> PathBuf { - let digest = Sha256::digest(key.as_bytes()); - data_dir - .join("correlation") - .join("parents") - .join(format!("{digest:x}.json")) -} - -pub(crate) async fn persist_active_parent_snapshot( - data_dir: &std::path::Path, - key: &str, - correlation: &crate::correlation::CorrelationRegistry, -) -> Result<(), String> { - let path = active_parent_snapshot_path(data_dir, key); - let Some(snapshot) = correlation.active_parent_snapshot(key) else { - let _ = tokio::fs::remove_file(path).await; - return Ok(()); - }; - let bytes = serde_json::to_vec(&snapshot) - .map_err(|error| format!("active parent encoding failed: {error}"))?; - write_active_parent_snapshot(data_dir, path, bytes).await -} - -async fn mark_active_parent_snapshot_dirty( - data_dir: &std::path::Path, - key: &str, - correlation: &crate::correlation::CorrelationRegistry, -) -> Result<(), String> { - let Some(snapshot) = correlation.dirty_active_parent_snapshot(key) else { - return Ok(()); - }; - let bytes = serde_json::to_vec(&snapshot) - .map_err(|error| format!("active parent encoding failed: {error}"))?; - write_active_parent_snapshot(data_dir, active_parent_snapshot_path(data_dir, key), bytes).await -} - -async fn write_active_parent_snapshot( - data_dir: &std::path::Path, - path: PathBuf, - bytes: Vec, -) -> Result<(), String> { - let dir = data_dir.join("correlation").join("parents"); - tokio::fs::create_dir_all(&dir) - .await - .map_err(|error| format!("active parent directory failed: {error}"))?; - let temp = path.with_extension(format!("{}.tmp", uuid::Uuid::new_v4())); - tokio::fs::write(&temp, bytes) - .await - .map_err(|error| format!("active parent write failed: {error}"))?; - if let Err(first_error) = tokio::fs::rename(&temp, &path).await { - let _ = tokio::fs::remove_file(&path).await; - tokio::fs::rename(&temp, &path).await.map_err(|error| { - format!("active parent replace failed: {first_error}; retry failed: {error}") - })?; - } - Ok(()) -} - -async fn restore_active_parent_snapshots( - data_dir: &std::path::Path, - correlation: &crate::correlation::CorrelationRegistry, -) { - let Ok(mut entries) = tokio::fs::read_dir(data_dir.join("correlation").join("parents")).await - else { - return; - }; - while let Ok(Some(entry)) = entries.next_entry().await { - let Ok(bytes) = tokio::fs::read(entry.path()).await else { - continue; - }; - let Ok(snapshot) = - serde_json::from_slice::(&bytes) - else { - tracing::debug!(path = %entry.path().display(), "invalid active parent snapshot ignored"); - continue; - }; - correlation.restore_active_parent(snapshot); - } -} - -async fn restore_pending_sessions(daemon: &Arc) { - let Ok(mut entries) = tokio::fs::read_dir(daemon.data_dir.join("correlation")).await else { - return; - }; - let mut restored = 0usize; - while let Ok(Some(entry)) = entries.next_entry().await { - let path = entry.path(); - if path.extension().and_then(|ext| ext.to_str()) != Some("json") { - continue; - } - let Ok(bytes) = tokio::fs::read(&path).await else { - continue; - }; - let Ok(state) = serde_json::from_slice::(&bytes) else { - continue; - }; - let state = state.restore(); - let Some(key) = state - .events - .first() - .map(|event| automatic_link_key(&event.env)) - else { - continue; - }; - daemon.pending_sessions.lock().unwrap().insert(key, state); - restored += 1; - } - if restored > 0 { - tracing::info!(restored, "restored pending child sessions"); - } -} - -async fn recover_unprocessed_journals(daemon: &Arc) { - let pending: HashSet<(String, String, u64)> = daemon - .pending_sessions - .lock() - .unwrap() - .values() - .flat_map(|state| { - state.events.iter().map(|event| { - ( - event.env.source.clone(), - event.env.session_id.clone(), - event.journal_through, - ) - }) - }) - .collect(); - let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { - return; - }; - let mut candidates = Vec::new(); - while let Ok(Some(entry)) = entries.next_entry().await { - let path = entry.path(); - if path.extension().and_then(|ext| ext.to_str()) != Some("ndjson") { - continue; - } - let recorded_len = journal::JournalReader::recorded_len(&path).await; - let Ok(Some(mut reader)) = journal::JournalReader::open(&path, recorded_len).await else { - continue; - }; - let mut before = 0u64; - let mut acknowledged_by_route: HashMap = HashMap::new(); - let mut latest_by_route: HashMap = HashMap::new(); - while let Ok(Some(entry)) = reader.next_record().await { - let through = entry.through; - match entry.record { - journal::JournalRecord::Event(redacted) => { - let mut env = journal::envelope_from_redacted(redacted); - let Some(source) = daemon.translators.canonical_source(&env.source) else { - before = through; - continue; - }; - env.source = source.to_string(); - daemon.remember_claude_transcripts(&env); - if env.source == "claude-code" { - if let Some(payload) = env.payload.as_object_mut() { - payload.insert( - "_bt_transcript_replay".to_string(), - serde_json::json!(true), - ); - } - } - if let Some(route) = env.route.as_ref() { - let key = serde_json::to_string(route).unwrap_or_default(); - latest_by_route.insert( - key, - PendingEvent { - env, - replay_through: before, - journal_through: through, - }, - ); - } - } - journal::JournalRecord::DeliveryCheckpoint { route, through } => { - let key = serde_json::to_string(&route).unwrap_or_default(); - let acknowledged = acknowledged_by_route.entry(key).or_default(); - *acknowledged = (*acknowledged).max(through); - } - } - before = through; - } - for event in latest_by_route.into_values() { - if pending.contains(&( - event.env.source.clone(), - event.env.session_id.clone(), - event.journal_through, - )) { - continue; - } - let route = recovered_delivery_route(daemon, &event.env) - .await - .as_ref() - .and_then(|route| serde_json::to_string(route).ok()) - .unwrap_or_default(); - if acknowledged_by_route.get(&route).copied().unwrap_or(0) < event.journal_through { - candidates.push(event); - } else if event.env.source == "claude-code" { - // A fully delivered historical session has no actor to retire - // its observer. Probe once for bytes appended after the final - // checkpoint, retaining only paths that have new work. - if let Some(key) = event.env.route.as_ref().and_then(|route| { - DeliveryKey::new(&event.env.source, &event.env.session_id, route).ok() - }) { - let observation = daemon.claude_transcripts.lock().unwrap().get(&key).cloned(); - if let Some(observation) = observation { - let path = observation.template.payload["transcript_path"] - .as_str() - .unwrap(); - let changed = tokio::fs::metadata(path) - .await - .is_ok_and(|metadata| observation.through != Some(metadata.len())); - if !changed { - daemon.claude_transcripts.lock().unwrap().remove(&key); - } - } - } - } - } - } - candidates.sort_by_key(|event| event.env.ts_ms); - let recovered = candidates.len(); - for event in candidates { - let _ = daemon - .ingress_tx - .send(IngressMsg::Event(Box::new(event))) - .await; - } - // Reconciliation stays in the daemon worker and follows every recovered - // event in queue order. Dropping the receiver is intentional: startup and - // hook capture do not wait for translation or reporting. - let (reply, _ignored) = oneshot::channel(); - let _ = daemon.ingress_tx.send(IngressMsg::Barrier(reply)).await; - if recovered > 0 { - tracing::info!( - recovered, - "queued unprocessed journal sessions for recovery" - ); - } -} - -async fn accept_resolved_event(daemon: &Arc, event: PendingEvent) -> Result<(), String> { - let PendingEvent { - mut env, - replay_through, - journal_through, - } = event; - // Every fresh Claude capture carries a bounded observation marker. Older - // persisted ingress must never consume live bytes while rebuilding state. - if env.source == "claude-code" && env.payload.get("_bt_claude_transcript_mirrors").is_none() { - if let Some(payload) = env.payload.as_object_mut() { - payload.insert("_bt_transcript_replay".to_string(), serde_json::json!(true)); - } - } - let schedule_flush = crate::hook::should_flush_ingress_event(&env) - || (env.source == "claude-code" && env.event == "TranscriptUpdate"); - let source = env.source.clone(); - 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"); - let session_lock = daemon.dispatch_lock(&source, &session_id); - let _session_guard = session_lock.lock().await; - - let result = async { - let delivery_key = daemon - .configure_event(&mut env) - .await - .map_err(|error| format!("session auth failed: {error}"))?; - if crate::dispatch::is_tool_lifecycle_event(&env.event) { - if let Err(error) = mark_active_parent_snapshot_dirty( - &daemon.data_dir, - &delivery_key.correlation_key(), - &daemon.correlation, - ) - .await - { - tracing::warn!(session_id = %env.session_id, %error, "active parent snapshot could not be marked dirty"); - } - } - let session = daemon - .session_for(&env, &delivery_key, replay_through) - .await - .map_err(|error| format!("session init failed: {error}"))?; - daemon - .correlation - .observe_session(&delivery_key.correlation_key(), &env.source, env.capture.as_ref(), - is_session_start(&env.source, &env.event)); - session - .enqueue(env, journal_through) - .await - .map_err(|error| format!("enqueue failed: {error}"))?; - Ok(delivery_key) - } - .await; - - match &result { - 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"), - } - if result.is_ok() && schedule_flush { - // Deferred terminal events need the same flush as immediate events; - // there may have been no session actor when ingress first accepted them. - daemon - .enqueue_source_session_flush(&source, &session_id, Duration::from_secs(10)) - .await; - } - result.map(|_| ()) -} - -fn is_session_start(source: &str, event: &str) -> bool { - matches!(event, "SessionStart" | "session_start" | "session.created") - || (source == "cursor" && event == "sessionStart") - // Antigravity has no dedicated session-start or resume hook. Its - // PreInvocation is the earliest hook in each invocation and is the - // only opportunity to refresh a resumed conversation's process. - || (source == "antigravity" && event == "PreInvocation") -} - -fn automatic_link_key(env: &Envelope) -> String { - let route = env - .route - .as_ref() - .and_then(|route| serde_json::to_string(route).ok()) - .unwrap_or_default(); - format!("{}\u{1f}{}\u{1f}{route}", env.source, env.session_id) -} - -/// Correlated child events retain their setup route in the immutable journal, -/// while their delivery checkpoint belongs to the resolved parent route. -/// Recover against that effective route without rewriting the captured event. -async fn recovered_delivery_route(daemon: &Arc, env: &Envelope) -> Option { - let key = automatic_link_key(env); - if let Some(route) = daemon.automatic_links.lock().unwrap().get(&key).cloned() { - return Some(route); - } - if let Some(route) = daemon - .pending_sessions - .lock() - .unwrap() - .get(&key) - .and_then(|state| state.linked_route.clone()) - { - return Some(route); - } - read_correlation_state(&daemon.data_dir, &key) - .await - .and_then(|state| state.linked_route) - .or_else(|| env.route.clone()) -} - -async fn handle_request( - daemon: &Arc, - req: Request, - client: &mut Option, -) -> Response { - let id = req.id.clone(); - let params = req.params.unwrap_or(serde_json::Value::Null); - - macro_rules! parse { - ($t:ty) => { - match serde_json::from_value::<$t>(params) { - Ok(v) => v, - Err(e) => { - return Response::err( - id, - RpcError::new(error_code::INVALID_PARAMS, format!("invalid params: {e}")), - ) - } - } - }; - } - - match req.method.as_str() { - method::INITIALIZE => { - let p = parse!(InitializeParams); - if p.protocol_version != PROTOCOL_VERSION { - return Response::err( - id, - RpcError::new( - error_code::APP, - format!( - "protocol version mismatch: client {} daemon {}", - p.protocol_version, PROTOCOL_VERSION - ), - ), - ); - } - *client = Some(p.client); - let result = InitializeResult { - protocol_version: PROTOCOL_VERSION, - daemon_version: daemon.version.clone(), - capabilities: Capabilities { - sources: daemon.translators.sources(), - }, - }; - Response::ok(id, serde_json::to_value(result).unwrap()) - } - method::EVENT_LOG => { - let mut env = parse!(Envelope); - attach_process_capture(daemon, &mut env, client.as_ref()); - match daemon.capture_event(env).await { - Ok(accepted) => Response::ok( - id, - serde_json::to_value(EventLogResult { accepted }).unwrap(), - ), - Err(error) => Response::err(id, RpcError::new(error_code::INTERNAL, error)), - } - } - method::SESSION_FLUSH => { - let p = parse!(FlushParams); - // Explicit flushes wait for the daemon-owned ingress queue. Hook - // capture never waits on this barrier. - daemon.settle_ingress().await; - if let Err(error) = - resolve_pending_sessions(daemon, |env| env.session_id == p.session_id).await - { - return Response::err(id, RpcError::new(error_code::INTERNAL, 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); - } - } - Response::ok( - id, - serde_json::to_value(FlushResult { - flushed, - pending, - accepted_sessions, - }) - .unwrap(), - ) - } - method::MANAGED_RUN_FLUSH => { - let params = parse!(ManagedRunFlushParams); - daemon.settle_ingress().await; - if let Err(error) = resolve_pending_sessions(daemon, |env| { - env.managed_run_id.as_deref() == Some(params.managed_run_id.as_str()) - }) - .await - { - return Response::err(id, RpcError::new(error_code::INTERNAL, error)); - } - let result = daemon.flush_managed_run(params).await; - Response::ok(id, serde_json::to_value(result).unwrap()) - } - method::STATUS_GET => { - let p = parse!(StatusParams); - daemon.settle_ingress().await; - Response::ok(id, serde_json::to_value(daemon.status(p)).unwrap()) - } - method::AUTH_DIAGNOSE => { - let p = parse!(AuthDiagnoseParams); - Response::ok( - id, - serde_json::to_value(daemon.diagnose_auth(p.auth).await).unwrap(), - ) - } - method::DAEMON_SHUTDOWN => { - if !client_may_shutdown(client.as_ref(), &daemon.version) { - return Response::ok( - id, - serde_json::to_value(ShutdownResult { ok: false }).unwrap(), - ); - } - drain_all(daemon).await; - Response::ok( - id, - serde_json::to_value(ShutdownResult { ok: true }).unwrap(), - ) - } - other => Response::err( - id, - RpcError::new( - error_code::METHOD_NOT_FOUND, - format!("unknown method: {other}"), - ), - ), - } -} - -/// Resolve a route's lease the way event delivery does, including the -/// organization the route requires. -pub(crate) async fn resolve_route_auth( - provider: &dyn AuthProvider, - selection: &AuthSelection, - reason: AuthResolveReason, - required_org: Option<&str>, -) -> anyhow::Result { - let lease = provider.resolve(selection, reason).await?; - if let Some(expected_org) = required_org { - if lease.auth.org_name.as_deref() != Some(expected_org) { - anyhow::bail!( - "profile {:?} resolved organization {:?}, expected {:?}", - lease.selection, - lease.auth.org_name, - expected_org - ); - } - } - Ok(lease) -} - -fn client_may_shutdown(client: Option<&ClientInfo>, daemon_version: &str) -> bool { - let Some(client) = client else { - return true; // Explicit stop commands do not initialize first. - }; - let Some(client_version) = client.daemon_version.as_deref() else { - return false; // Legacy initialized hooks must not downgrade the daemon. - }; - crate::client::compare_daemon_versions(client_version, daemon_version) - .is_none_or(|ordering| !ordering.is_lt()) -} - -impl Daemon { - fn status(&self, p: StatusParams) -> StatusResult { - let map = self.sessions.lock().unwrap(); - let mut sessions: Vec<_> = map - .iter() - .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 (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(key) { - continue; - } - sessions.push(SessionStatus { - session_id: key.session_id.clone(), - source: source.clone(), - route: serde_json::from_str(&key.route).ok(), - queued: 0, - spans_emitted: 0, - permalink: None, - last_error: Some(error.clone()), - }); - } - for state in self.pending_sessions.lock().unwrap().values() { - let Some(first) = state.events.first() else { - continue; - }; - if p.session_id - .as_ref() - .is_some_and(|want| want != &first.env.session_id) - { - continue; - } - let route = state.linked_route.as_ref().or(first.env.route.as_ref()); - let existing = sessions.iter_mut().find(|session| { - session.session_id == first.env.session_id - && session.source == first.env.source - && match (session.route.as_ref(), route) { - (Some(left), Some(right)) => left.same_route(right), - (None, None) => true, - _ => false, - } - }); - let message = if state.decided() { - "correlation resolved; journaled events awaiting delivery" - } else { - "awaiting parent process confirmation; standalone fallback within 5 seconds" - }; - if let Some(session) = existing { - session.queued = session.queued.saturating_add(state.events.len() as u64); - if session.last_error.is_none() { - session.last_error = Some(message.to_string()); - } - } else { - sessions.push(SessionStatus { - session_id: first.env.session_id.clone(), - source: first.env.source.clone(), - route: route.cloned(), - queued: state.events.len() as u64, - spans_emitted: 0, - permalink: None, - last_error: Some(message.to_string()), - }); - } - } - StatusResult { - daemon_version: self.version.clone(), - uptime_ms: self.started.elapsed().as_millis() as u64, - sessions, - } - } -} - -/// How long recovery state is kept on disk. -const RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60); -/// How often that retention is enforced while the daemon keeps running. -const GC_INTERVAL: Duration = Duration::from_secs(60 * 60); - -async fn collect_garbage(data_dir: &std::path::Path) { - journal::gc_old_journals(data_dir, RETENTION).await; - journal::gc_old_managed_runs(data_dir, RETENTION).await; - crate::transcript_mirror::gc_old_mirrors(data_dir, RETENTION).await; - gc_old_correlation_states(data_dir, RETENTION).await; -} - -async fn gc_old_correlation_states(data_dir: &std::path::Path, max_age: Duration) { - gc_old_correlation_dir(&data_dir.join("correlation"), max_age).await; - gc_old_correlation_dir(&data_dir.join("correlation").join("parents"), max_age).await; -} - -async fn gc_old_correlation_dir(dir: &std::path::Path, max_age: Duration) { - let Ok(mut entries) = tokio::fs::read_dir(dir).await else { - return; - }; - let now = SystemTime::now(); - while let Ok(Some(entry)) = entries.next_entry().await { - let old = entry - .metadata() - .await - .ok() - .and_then(|metadata| metadata.modified().ok()) - .and_then(|modified| now.duration_since(modified).ok()) - .is_some_and(|age| age > max_age); - if old && entry.file_type().await.is_ok_and(|kind| kind.is_file()) { - let _ = tokio::fs::remove_file(entry.path()).await; - } - } -} - -/// Retire delivery pipelines that have gone quiet. Without this, every session -/// the daemon ever saw keeps its translator state, sink handles, and journal -/// file handle alive until the process exits — which for a continuously busy -/// user is never, since the idle watchdog needs *all* sessions quiet. -fn spawn_session_reaper(daemon: Arc, idle_timeout: Duration) { - if idle_timeout.is_zero() { - return; // 0 disables retirement (useful in tests) - } - tokio::spawn(async move { - let tick = (idle_timeout / 4).max(Duration::from_secs(1)); - loop { - tokio::select! { - _ = daemon.shutdown.notified() => return, - _ = tokio::time::sleep(tick) => {} - } - for key in daemon.idle_sessions(idle_timeout) { - daemon.retire_session(&key).await; - } - } - }); -} - -fn spawn_gc(daemon: Arc) { - tokio::spawn(async move { - loop { - tokio::select! { - _ = daemon.shutdown.notified() => return, - _ = tokio::time::sleep(GC_INTERVAL) => {} - } - collect_garbage(&daemon.data_dir).await; - daemon - .correlation - .prune_retired(crate::process::process_is_alive); - } - }); -} - -fn spawn_idle_watchdog(daemon: Arc, idle_timeout: Duration) { - if idle_timeout.is_zero() { - return; // 0 disables the watchdog (useful in tests) - } - tokio::spawn(async move { - let tick = (idle_timeout / 4).max(Duration::from_secs(1)); - loop { - tokio::select! { - _ = daemon.shutdown.notified() => return, - _ = tokio::time::sleep(tick) => {} - } - let idle_for = daemon.last_activity.lock().unwrap().elapsed(); - if idle_for >= idle_timeout - && daemon.total_queued() == 0 - && !daemon.correlation.has_any_active_tools() - { - tracing::info!("idle for {:?}; shutting down", idle_for); - daemon.trigger_shutdown(); - return; - } - } - }); -} - -async fn drain_all(daemon: &Arc) { - let mut drained = daemon.drained.lock().await; - if *drained { - return; - } - daemon.begin_quiesce(); - let _capture_guard = daemon.capture_gate.write().await; - daemon.observe_claude_transcripts(None, true).await; - daemon.settle_ingress().await; - if let Err(error) = resolve_pending_sessions(daemon, |_| true).await { - tracing::warn!(%error, "pending correlation events retained for recovery during shutdown"); - } - let sessions: Vec> = daemon.sessions.lock().unwrap().values().cloned().collect(); - for s in sessions { - s.shutdown().await; - } - *drained = true; -} - -/// Is a live daemon answering at the endpoint? Connect and expect any line -/// back from a well-formed `initialize`. -async fn probe_alive(endpoint: &std::path::Path) -> bool { - let Ok(stream) = transport::connect(endpoint).await else { - return false; - }; - let (read_half, mut write_half) = tokio::io::split(stream); - let init = Request::new( - crate::wire::RequestId::Int(0), - method::INITIALIZE, - serde_json::json!({ - "protocol_version": PROTOCOL_VERSION, - "client": { "source": "probe" } - }), - ); - let Ok(mut line) = Message::Request(init).to_line() else { - return false; - }; - line.push('\n'); - if write_half.write_all(line.as_bytes()).await.is_err() { - return false; - } - let mut lines = BufReader::new(read_half).lines(); - matches!( - tokio::time::timeout(Duration::from_secs(1), lines.next_line()).await, - Ok(Ok(Some(_))) - ) -} - -/// Run the daemon until shutdown. -pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { - run(args, opts).await -} - -/// Build debug [`ServeOptions`]: the production translator registry plus a -/// debug sink writing NDJSON under `/spans/`. -pub fn debug_serve_options(version: impl Into, data_dir: &std::path::Path) -> ServeOptions { - ServeOptions { - version: version.into(), - translators: Arc::new(Registry::default_agents()), - sink_factory: Arc::new(DebugSinkFactory { - dir: data_dir.join("spans"), - }), - auth_provider: None, - } -} - -/// Build [`ServeOptions`] with the Braintrust sink. `translators` lets the -/// caller choose the translator registry (debug-only until Phase 3 adds the -/// Codex/Claude translators). Clients are built lazily per session URL, so this -/// is cheap and infallible. -pub fn braintrust_serve_options( - version: impl Into, - sink_config: BraintrustSinkConfig, - translators: Arc, -) -> ServeOptions { - ServeOptions { - version: version.into(), - translators, - sink_factory: Arc::new(BraintrustSinkFactory::new(sink_config)), - auth_provider: None, - } -} - -#[cfg(test)] -mod tests { - use super::{client_may_shutdown, is_session_start, ClientInfo}; - - #[test] - fn antigravity_invocation_refreshes_a_resumed_conversations_process() { - assert!(is_session_start("antigravity", "PreInvocation")); - assert!(!is_session_start("claude-code", "PreInvocation")); - assert!(is_session_start("claude-code", "SessionStart")); - assert!(is_session_start("pi", "session_start")); - assert!(is_session_start("opencode", "session.created")); - } - - fn initialized(version: Option<&str>) -> ClientInfo { - ClientInfo { - source: "codex".into(), - daemon_version: version.map(str::to_string), - plugin_version: None, - pid: None, - } - } - - #[test] - fn legacy_initialized_clients_cannot_downgrade_the_daemon() { - assert!(client_may_shutdown(None, "0.20.0")); - assert!(!client_may_shutdown(Some(&initialized(None)), "0.20.0")); - assert!(!client_may_shutdown( - Some(&initialized(Some("0.19.3"))), - "0.20.0" - )); - assert!(client_may_shutdown( - Some(&initialized(Some("0.20.0"))), - "0.20.0" - )); - assert!(client_may_shutdown( - Some(&initialized(Some("0.21.0"))), - "0.20.0" - )); - } -} diff --git a/bt-daemon/src/server/auth.rs b/bt-daemon/src/server/auth.rs new file mode 100644 index 0000000..9d56212 --- /dev/null +++ b/bt-daemon/src/server/auth.rs @@ -0,0 +1,209 @@ +//! Per-session authentication: resolving a route to a backend lease through +//! the host, refreshing it before flushes, and diagnosing failures. + +use super::*; + +pub(super) 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) +} + +/// Resolve a route's lease the way event delivery does, including the +/// organization the route requires. +pub(crate) async fn resolve_route_auth( + provider: &dyn AuthProvider, + selection: &AuthSelection, + reason: AuthResolveReason, + required_org: Option<&str>, +) -> anyhow::Result { + let lease = provider.resolve(selection, reason).await?; + if let Some(expected_org) = required_org { + if lease.auth.org_name.as_deref() != Some(expected_org) { + anyhow::bail!( + "profile {:?} resolved organization {:?}, expected {:?}", + lease.selection, + lease.auth.org_name, + expected_org + ); + } + } + Ok(lease) +} + +impl Daemon { + pub(super) 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 requested_route = env + .route + .clone() + .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; + if requested_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 requested_key = DeliveryKey::new(&env.source, &env.session_id, &requested_route)?; + let key = self + .route_aliases + .lock() + .unwrap() + .get(&requested_key) + .cloned() + .unwrap_or_else(|| requested_key.clone()); + let (selection, reason, expected_selection) = { + let states = self.session_auth.lock().await; + match states.get(&key) { + Some(state) => { + if !lease_is_expiring(&state.lease) { + env.config = Some(state.route.with_auth(state.lease.auth.clone())); + return Ok(key); + } + ( + state.lease.selection.clone(), + AuthResolveReason::Expiring, + Some(state.lease.selection.clone()), + ) + } + None => ( + requested_route.auth.clone(), + AuthResolveReason::Initial, + None, + ), + } + }; + + let lease = resolve_route_auth( + provider.as_ref(), + &selection, + reason, + requested_route.auth.org_name.as_deref(), + ) + .await + .map_err(|error| { + let message = format!( + "could not resolve Braintrust auth for {}: {error}; run `bt login` or select a profile explicitly", + env.source + ); + self.record_auth_error(&key, &env.source, message.clone()); + anyhow::anyhow!(message) + })?; + if let Some(expected) = expected_selection { + if lease.selection != expected { + anyhow::bail!( + "credential refresh changed auth selection from {expected:?} to {:?}", + lease.selection + ); + } + } + + let mut route = requested_route; + if route.auth.effective_source() == crate::wire::AuthSource::SavedProfile + && route.auth.profile_id.is_none() + && lease.selection.profile_id.is_some() + { + route.auth = lease.selection.clone(); + if let Err(error) = crate::settings::migrate_persisted_route( + &env.source, + env.route.as_ref().expect("route checked above"), + &route, + ) { + tracing::warn!(source = %env.source, "could not migrate legacy profile route: {error}"); + } + } + let canonical_key = DeliveryKey::new(&env.source, &env.session_id, &route)?; + if canonical_key != requested_key { + self.route_aliases + .lock() + .unwrap() + .insert(requested_key, canonical_key.clone()); + } + env.route = Some(route.clone()); + env.config = Some(route.with_auth(lease.auth.clone())); + self.session_auth + .lock() + .await + .insert(canonical_key.clone(), SessionAuthState { route, lease }); + self.auth_errors.lock().unwrap().remove(&canonical_key); + Ok(canonical_key) + } + + /// Keep a route's auth failure visible to `status.get` until it succeeds. + pub(super) fn record_auth_error(&self, key: &DeliveryKey, source: &str, message: String) { + self.auth_errors + .lock() + .unwrap() + .insert(key.clone(), (source.to_string(), message)); + } + + /// Resolve a route's credentials exactly as event delivery would, so + /// diagnostics see this process's credential store and environment rather + /// than the caller's. + pub(super) async fn diagnose_auth(&self, selection: AuthSelection) -> AuthDiagnoseResult { + let result = match &self.auth_provider { + Some(provider) => { + resolve_route_auth( + provider.as_ref(), + &selection, + AuthResolveReason::Initial, + selection.org_name.as_deref(), + ) + .await + } + None => Err(anyhow::anyhow!( + "daemon host has no Braintrust auth provider" + )), + }; + match result { + Ok(lease) => AuthDiagnoseResult { + selection: Some(lease.selection), + org_name: lease.auth.org_name, + expires_at_ms: lease.expires_at_ms, + error: None, + }, + Err(error) => AuthDiagnoseResult { + error: Some(error.to_string()), + ..AuthDiagnoseResult::default() + }, + } + } + + pub(super) 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(key).cloned() else { + return Ok(()); + }; + if !lease_is_expiring(&state.lease) { + return Ok(()); + } + let selection = state.lease.selection.clone(); + let lease = provider + .resolve(&selection, AuthResolveReason::Expiring) + .await?; + if lease.selection != state.lease.selection { + anyhow::bail!("credential refresh changed the session auth selection"); + } + let config = state.route.with_auth(lease.auth.clone()); + self.session_auth.lock().await.insert( + key.clone(), + SessionAuthState { + route: state.route, + lease, + }, + ); + let session = { self.sessions.lock().unwrap().get(key).cloned() }; + if let Some(session) = session { + session.configure(config).await?; + } + Ok(()) + } +} diff --git a/bt-daemon/src/server/claude_transcripts.rs b/bt-daemon/src/server/claude_transcripts.rs new file mode 100644 index 0000000..404b235 --- /dev/null +++ b/bt-daemon/src/server/claude_transcripts.rs @@ -0,0 +1,148 @@ +//! Claude Code transcript polling. The daemon remembers each Claude session's +//! transcript and emits synthetic `TranscriptUpdate` events for it, so records +//! written after the session's last hook are still translated. + +use super::*; + +#[derive(Clone)] +pub(super) struct ClaudeTranscriptObservation { + /// Synthetic-event provenance only: never retain native prompts or responses. + pub(super) template: Envelope, + pub(super) through: Option, +} + +impl Daemon { + pub(super) fn remember_claude_transcripts(&self, env: &Envelope) { + if env.source != "claude-code" || env.payload.get("_bt_import_through_offset").is_some() { + return; + } + if env.event == "SessionEnd" { + self.claude_transcripts + .lock() + .unwrap() + .retain(|key, _| key.source != env.source || key.session_id != env.session_id); + return; + } + let Some(route) = env.route.as_ref() else { + return; + }; + let Some(path) = env.payload.get("transcript_path").and_then(Value::as_str) else { + return; + }; + let Ok(key) = DeliveryKey::new(&env.source, &env.session_id, route) else { + return; + }; + let through = env + .payload + .get("_bt_claude_transcript_mirrors") + .and_then(|mirrors| mirrors.get(path)) + .or_else(|| { + env.payload + .get("_bt_transcript_mirror") + .filter(|mirror| mirror.get("path").and_then(Value::as_str) == Some(path)) + }) + .and_then(|mirror| mirror.get("through")) + .and_then(Value::as_u64); + let mut payload = serde_json::json!({"transcript_path": path}); + if let Some(cwd) = env.payload.get("cwd").and_then(Value::as_str) { + payload["cwd"] = Value::String(cwd.to_string()); + } + let template = Envelope { + source: env.source.clone(), + source_version: env.source_version.clone(), + plugin_version: env.plugin_version.clone(), + session_id: env.session_id.clone(), + event: "TranscriptUpdate".to_string(), + ts_ms: env.ts_ms, + managed_run_id: env.managed_run_id.clone(), + capture: None, + payload, + route: env.route.clone(), + config: None, + }; + let mut observations = self.claude_transcripts.lock().unwrap(); + let previous = observations + .get(&key) + .filter(|observation| { + observation + .template + .payload + .get("transcript_path") + .and_then(Value::as_str) + == Some(path) + }) + .and_then(|observation| observation.through); + observations.insert( + key, + ClaudeTranscriptObservation { + template, + through: through.or(previous), + }, + ); + } + + /// Poll native paths only on the capture side. Each observed boundary is + /// journaled before translation, including observations made during shutdown. + pub(super) async fn observe_claude_transcripts( + &self, + only: Option<&DeliveryKey>, + quiesced: bool, + ) { + // Shutdown already owns the exclusive gate. Ordinary polling must not + // append after that final drain begins. + let _gate = if quiesced { + None + } else { + Some(self.capture_gate.read().await) + }; + if !quiesced && self.quiescing.load(Ordering::SeqCst) { + return; + } + let keys: Vec<_> = self + .claude_transcripts + .lock() + .unwrap() + .keys() + .filter(|key| { + only.is_none_or(|only| { + key.source == only.source && key.session_id == only.session_id + }) + }) + .cloned() + .collect(); + for key in keys { + // Routes share mirror files, so serialize by source session, not + // destination. Never hold this lock across an ingress barrier. + let lock = self.session_lock(&key.source, &key.session_id); + let _guard = lock.lock().await; + let observation = self.claude_transcripts.lock().unwrap().get(&key).cloned(); + let Some(observation) = observation else { + continue; + }; + let path = observation.template.payload["transcript_path"] + .as_str() + .unwrap(); + let Ok(metadata) = tokio::fs::metadata(path).await else { + continue; + }; + if observation.through == Some(metadata.len()) { + continue; + } + let mut env = observation.template; + env.ts_ms = now_ms(); + hydrate_transcript_reference(&self.data_dir, &mut env).await; + let through = env + .payload + .get("_bt_claude_transcript_mirrors") + .and_then(|mirrors| mirrors.get(env.payload["transcript_path"].as_str().unwrap())) + .and_then(|mirror| mirror.get("through")) + .and_then(Value::as_u64); + if through.is_none() || through == observation.through { + continue; + } + if let Err(error) = self.journal_and_enqueue(env).await { + tracing::warn!(session_id = %key.session_id, %error, "Claude transcript observation failed"); + } + } + } +} diff --git a/bt-daemon/src/server/correlation_store.rs b/bt-daemon/src/server/correlation_store.rs new file mode 100644 index 0000000..1e25da2 --- /dev/null +++ b/bt-daemon/src/server/correlation_store.rs @@ -0,0 +1,611 @@ +//! Automatic parent correlation: holding a new session briefly while its +//! parent agent session is identified, and persisting that decision and the +//! active parent snapshots across daemon restarts. + +use super::*; + +// Correlation is a short, optional discovery step, never a delivery queue. +pub(super) const CORRELATION_VERSION: u32 = 1; + +// Long enough for a slow cold start of the second hook (e.g. an antivirus scan +// of `bt.exe`); only sessions with a frozen parent candidate ever wait. +pub(super) const CORRELATION_WAIT: Duration = Duration::from_secs(5); + +pub(super) const CORRELATION_EVENT_LIMIT: usize = 3; + +#[derive(Clone, Serialize, Deserialize)] +pub(super) struct PendingSession { + // Missing versions predate safe ancestry matching; see `restore`. + #[serde(default)] + pub(super) version: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) linked_route: Option, + #[serde(default)] + pub(super) standalone: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) parent: Option, + #[serde(default)] + pub(super) first_seen_ms: i64, + #[serde(skip)] + pub(super) deadline: Option, + // Discovery expiry is not an authentication/delivery retry loop. A failed + // handoff stays visible until a new event, explicit flush, or restart. + #[serde(skip)] + pub(super) delivery_failed: bool, + #[serde(default)] + pub(super) events: Vec, +} + +impl Default for PendingSession { + fn default() -> Self { + Self { + version: CORRELATION_VERSION, + linked_route: None, + standalone: false, + parent: None, + first_seen_ms: 0, + deadline: None, + delivery_failed: false, + events: Vec::new(), + } + } +} + +impl PendingSession { + pub(super) fn restore(mut self) -> Self { + if self.version != CORRELATION_VERSION { + // A legacy link was decided before safe ancestry matching, but + // its earlier rows already live in that parent's trace. Keep later + // rows there too: splitting one session across two roots is worse. + // Undecided legacy buffers were never delivered; send them alone. + self.version = CORRELATION_VERSION; + self.standalone = self.linked_route.is_none(); + self.parent = None; + } + let remaining = self + .first_seen_ms + .saturating_add(CORRELATION_WAIT.as_millis() as i64) + .saturating_sub(now_ms()); + // A missing/future receipt timestamp is not permission to wait longer. + let remaining = if (0..=CORRELATION_WAIT.as_millis() as i64).contains(&remaining) { + remaining as u64 + } else { + 0 + }; + self.deadline = Some(Instant::now() + Duration::from_millis(remaining)); + self + } + + pub(super) fn decided(&self) -> bool { + self.standalone || self.linked_route.is_some() + } + + pub(super) fn expired(&self) -> bool { + self.deadline + .is_none_or(|deadline| Instant::now() >= deadline) + } +} + +#[derive(Clone, Serialize)] +pub(super) struct PendingEvent { + pub(super) env: Envelope, + pub(super) replay_through: u64, + pub(super) journal_through: u64, +} + +impl<'de> Deserialize<'de> for PendingEvent { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + #[serde(untagged)] + enum StoredPendingEvent { + Current { + env: Envelope, + replay_through: u64, + journal_through: u64, + }, + Legacy(Envelope), + } + Ok(match StoredPendingEvent::deserialize(deserializer)? { + StoredPendingEvent::Current { + env, + replay_through, + journal_through, + } => Self { + env, + replay_through, + journal_through, + }, + StoredPendingEvent::Legacy(env) => Self { + env, + replay_through: 0, + journal_through: 0, + }, + }) + } +} + +pub(super) async fn accept_event(daemon: &Arc, event: PendingEvent) -> Result<(), String> { + let key = automatic_link_key(&event.env); + let lock = daemon.correlation_lock(&key); + let _guard = lock.lock().await; + let mut state = daemon.pending_sessions.lock().unwrap().get(&key).cloned(); + if state.is_none() { + if daemon.standalone_links.lock().unwrap().contains(&key) { + return accept_resolved_event(daemon, event).await; + } + let route = daemon.automatic_links.lock().unwrap().get(&key).cloned(); + if let Some(route) = route { + let mut event = event; + event.env.route = Some(route); + return accept_resolved_event(daemon, event).await; + } + state = read_correlation_state(&daemon.data_dir, &key).await; + } + + if let Some(mut state) = state { + if state.decided() { + // Queue behind any earlier events whose handoff failed so a + // second failure keeps this event visible and retryable too. + state.events.push(event); + return finish_correlation(daemon, &key, state).await; + } + + // A new native start is a different process generation, not a second + // observation confirming the previous hook's ancestry. + let new_generation = is_session_start(&event.env.source, &event.env.event); + state.events.push(event); + if !state.expired() && !new_generation { + let mut captures = state + .events + .iter() + .rev() + .filter_map(|event| event.env.capture.as_ref()); + if let (Some(latest), Some(previous), Some(parent)) = + (captures.next(), captures.next(), state.parent.as_ref()) + { + if let Some(agent) = crate::correlation::session_agent_process( + &state.events[0].env.source, + previous, + Some(latest), + ) { + match crate::correlation::CorrelationRegistry::confirm_parent( + parent, latest, &agent, + ) { + crate::correlation::Resolution::Parent(parent) => { + state.linked_route = Some(parent.route); + } + _ => state.standalone = true, + } + } + } + } + if !state.decided() + && (new_generation || state.expired() || state.events.len() >= CORRELATION_EVENT_LIMIT) + { + state.standalone = true; + } + if state.decided() { + return finish_correlation(daemon, &key, state).await; + } + return store_correlation(daemon, &key, &state).await; + } + + let mut state = PendingSession::default(); + // An explicit parent is authoritative. Missing first-start/process evidence + // also means standalone, permanently; a later start cannot reparent rows + // already delivered. + if is_session_start(&event.env.source, &event.env.event) + && !matches!( + event + .env + .route + .as_ref() + .and_then(|route| route.destination.as_ref()), + Some(crate::wire::TraceDestination::ParentSpan { .. }) + ) + { + let child_agent = event.env.capture.as_ref().and_then(|capture| { + crate::correlation::session_agent_process(&event.env.source, capture, None) + }); + let session_prefix = format!("{}\u{1f}{}\u{1f}", event.env.source, event.env.session_id); + match daemon.correlation.resolve( + Some(&session_prefix), + event.env.capture.as_ref(), + child_agent.as_ref(), + |parent| automatic_parent_route(event.env.route.as_ref(), parent), + ) { + crate::correlation::Resolution::Parent(parent) => { + state.linked_route = Some(parent.route); + } + crate::correlation::Resolution::Pending(parent) => { + state.parent = Some(*parent); + state.first_seen_ms = now_ms(); + state.deadline = Some(Instant::now() + CORRELATION_WAIT); + } + crate::correlation::Resolution::Standalone => {} + } + } + state.standalone = state.linked_route.is_none() && state.parent.is_none(); + state.events.push(event); + if state.decided() { + finish_correlation(daemon, &key, state).await + } else { + store_correlation(daemon, &key, &state).await + } +} + +/// An inferred parent may change hierarchy, never the requested credential or +/// destination. Preserve the child's transforms, metadata, tags and flush mode. +pub(super) fn automatic_parent_route( + child: Option<&SessionRoute>, + parent: &SessionRoute, +) -> Option { + use crate::wire::TraceDestination; + use braintrust_sdk_rust::SpanObjectType; + + let child = child?; + let same_profile = match (&child.auth.profile_id, &parent.auth.profile_id) { + (Some(child), Some(parent)) => child == parent, + _ => child.auth.profile == parent.auth.profile, + }; + if child.auth.effective_source() != parent.auth.effective_source() + || !same_profile + || child.auth.org_name != parent.auth.org_name + { + return None; + } + let Some(TraceDestination::ParentSpan { components }) = &parent.destination else { + return None; + }; + let same_destination = match child.destination.as_ref()? { + TraceDestination::ProjectLogs { + project_id, + project_name, + } => { + if components.object_type != SpanObjectType::ProjectLogs { + false + } else if let Some(id) = project_id { + components.object_id.as_ref() == Some(id) + } else { + project_name.as_deref().is_some_and(|name| { + components + .compute_object_metadata_args + .as_ref() + .and_then(|args| args.get("project_name")) + .and_then(Value::as_str) + == Some(name) + }) + } + } + TraceDestination::Experiment { experiment_id } => { + components.object_type == SpanObjectType::Experiment + && components.object_id.as_ref() == Some(experiment_id) + } + TraceDestination::ParentSpan { .. } => false, + }; + if !same_destination { + return None; + } + let mut route = child.clone(); + route.destination = parent.destination.clone(); + Some(route) +} + +pub(super) async fn store_correlation( + daemon: &Daemon, + key: &str, + state: &PendingSession, +) -> Result<(), String> { + // Keep accepted events visible/retryable even if persisting the decision + // fails. The immutable raw journal remains authoritative for recovery. + if !state.events.is_empty() { + daemon + .pending_sessions + .lock() + .unwrap() + .insert(key.to_string(), state.clone()); + } + write_correlation_state(&daemon.data_dir, key, state).await?; + if state.events.is_empty() { + daemon.pending_sessions.lock().unwrap().remove(key); + if let Some(route) = &state.linked_route { + daemon + .automatic_links + .lock() + .unwrap() + .insert(key.to_string(), route.clone()); + } else if state.standalone { + daemon + .standalone_links + .lock() + .unwrap() + .insert(key.to_string()); + } + } + Ok(()) +} + +pub(super) async fn finish_correlation( + daemon: &Arc, + key: &str, + mut state: PendingSession, +) -> Result<(), String> { + state.parent = None; + state.delivery_failed = false; + let result = async { + // Persist the final route before delivering any events. A crash may + // replay deterministic rows, but must never select a different parent. + store_correlation(daemon, key, &state).await?; + while let Some(event) = state.events.first() { + let mut event = event.clone(); + if let Some(route) = &state.linked_route { + event.env.route = Some(route.clone()); + } + accept_resolved_event(daemon, event).await?; + state.events.remove(0); + store_correlation(daemon, key, &state).await?; + } + Ok(()) + } + .await; + if result.is_err() { + state.delivery_failed = true; + daemon + .pending_sessions + .lock() + .unwrap() + .insert(key.to_string(), state); + } + result +} + +pub(super) async fn retry_pending_sessions(daemon: &Arc) -> Result<(), String> { + resolve_pending_sessions(daemon, |_| false).await +} + +pub(super) async fn resolve_pending_sessions( + daemon: &Arc, + force: impl Fn(&Envelope) -> bool + Send + Sync, +) -> Result<(), String> { + let _reconcile_guard = daemon.pending_reconcile_lock.lock().await; + let keys: Vec<_> = daemon + .pending_sessions + .lock() + .unwrap() + .keys() + .cloned() + .collect(); + let mut failure = None; + for key in keys { + let lock = daemon.correlation_lock(&key); + let _guard = lock.lock().await; + let Some(mut state) = daemon.pending_sessions.lock().unwrap().get(&key).cloned() else { + continue; + }; + let forced = state.events.first().is_some_and(|event| force(&event.env)); + if state.delivery_failed && !forced { + continue; + } + if !state.decided() && (state.expired() || forced) { + state.standalone = true; + } + if state.decided() { + if let Err(error) = finish_correlation(daemon, &key, state).await { + // One unavailable destination must not prevent other sessions + // from leaving the discovery window. + failure = Some(error); + } + } + } + failure.map_or(Ok(()), Err) +} + +pub(super) fn correlation_state_path(data_dir: &std::path::Path, key: &str) -> PathBuf { + let digest = Sha256::digest(key.as_bytes()); + data_dir + .join("correlation") + .join(format!("{digest:x}.json")) +} + +pub(super) async fn read_correlation_state( + data_dir: &std::path::Path, + key: &str, +) -> Option { + let bytes = tokio::fs::read(correlation_state_path(data_dir, key)) + .await + .ok()?; + serde_json::from_slice::(&bytes) + .ok() + .map(PendingSession::restore) +} + +pub(super) async fn write_correlation_state( + data_dir: &std::path::Path, + key: &str, + state: &PendingSession, +) -> Result<(), String> { + let dir = data_dir.join("correlation"); + tokio::fs::create_dir_all(&dir) + .await + .map_err(|error| format!("correlation journal directory failed: {error}"))?; + let path = correlation_state_path(data_dir, key); + let temp = path.with_extension(format!("{}.tmp", uuid::Uuid::new_v4())); + let bytes = serde_json::to_vec(state) + .map_err(|error| format!("correlation journal encoding failed: {error}"))?; + tokio::fs::write(&temp, bytes) + .await + .map_err(|error| format!("correlation journal write failed: {error}"))?; + if let Err(first_error) = tokio::fs::rename(&temp, &path).await { + // Windows cannot replace an existing destination with rename. + let _ = tokio::fs::remove_file(&path).await; + tokio::fs::rename(&temp, &path).await.map_err(|error| { + format!("correlation journal replace failed: {first_error}; retry failed: {error}") + })?; + } + Ok(()) +} + +pub(super) fn active_parent_snapshot_path(data_dir: &std::path::Path, key: &str) -> PathBuf { + let digest = Sha256::digest(key.as_bytes()); + data_dir + .join("correlation") + .join("parents") + .join(format!("{digest:x}.json")) +} + +pub(crate) async fn persist_active_parent_snapshot( + data_dir: &std::path::Path, + key: &str, + correlation: &crate::correlation::CorrelationRegistry, +) -> Result<(), String> { + let path = active_parent_snapshot_path(data_dir, key); + let Some(snapshot) = correlation.active_parent_snapshot(key) else { + let _ = tokio::fs::remove_file(path).await; + return Ok(()); + }; + let bytes = serde_json::to_vec(&snapshot) + .map_err(|error| format!("active parent encoding failed: {error}"))?; + write_active_parent_snapshot(data_dir, path, bytes).await +} + +pub(super) async fn mark_active_parent_snapshot_dirty( + data_dir: &std::path::Path, + key: &str, + correlation: &crate::correlation::CorrelationRegistry, +) -> Result<(), String> { + let Some(snapshot) = correlation.dirty_active_parent_snapshot(key) else { + return Ok(()); + }; + let bytes = serde_json::to_vec(&snapshot) + .map_err(|error| format!("active parent encoding failed: {error}"))?; + write_active_parent_snapshot(data_dir, active_parent_snapshot_path(data_dir, key), bytes).await +} + +pub(super) async fn write_active_parent_snapshot( + data_dir: &std::path::Path, + path: PathBuf, + bytes: Vec, +) -> Result<(), String> { + let dir = data_dir.join("correlation").join("parents"); + tokio::fs::create_dir_all(&dir) + .await + .map_err(|error| format!("active parent directory failed: {error}"))?; + let temp = path.with_extension(format!("{}.tmp", uuid::Uuid::new_v4())); + tokio::fs::write(&temp, bytes) + .await + .map_err(|error| format!("active parent write failed: {error}"))?; + if let Err(first_error) = tokio::fs::rename(&temp, &path).await { + let _ = tokio::fs::remove_file(&path).await; + tokio::fs::rename(&temp, &path).await.map_err(|error| { + format!("active parent replace failed: {first_error}; retry failed: {error}") + })?; + } + Ok(()) +} + +pub(super) async fn restore_active_parent_snapshots( + data_dir: &std::path::Path, + correlation: &crate::correlation::CorrelationRegistry, +) { + let Ok(mut entries) = tokio::fs::read_dir(data_dir.join("correlation").join("parents")).await + else { + return; + }; + while let Ok(Some(entry)) = entries.next_entry().await { + let Ok(bytes) = tokio::fs::read(entry.path()).await else { + continue; + }; + let Ok(snapshot) = + serde_json::from_slice::(&bytes) + else { + tracing::debug!(path = %entry.path().display(), "invalid active parent snapshot ignored"); + continue; + }; + correlation.restore_active_parent(snapshot); + } +} + +pub(super) async fn restore_pending_sessions(daemon: &Arc) { + let Ok(mut entries) = tokio::fs::read_dir(daemon.data_dir.join("correlation")).await else { + return; + }; + let mut restored = 0usize; + while let Ok(Some(entry)) = entries.next_entry().await { + let path = entry.path(); + if path.extension().and_then(|ext| ext.to_str()) != Some("json") { + continue; + } + let Ok(bytes) = tokio::fs::read(&path).await else { + continue; + }; + let Ok(state) = serde_json::from_slice::(&bytes) else { + continue; + }; + let state = state.restore(); + let Some(key) = state + .events + .first() + .map(|event| automatic_link_key(&event.env)) + else { + continue; + }; + daemon.pending_sessions.lock().unwrap().insert(key, state); + restored += 1; + } + if restored > 0 { + tracing::info!(restored, "restored pending child sessions"); + } +} + +pub(super) fn automatic_link_key(env: &Envelope) -> String { + let route = env + .route + .as_ref() + .and_then(|route| serde_json::to_string(route).ok()) + .unwrap_or_default(); + format!("{}\u{1f}{}\u{1f}{route}", env.source, env.session_id) +} + +/// Correlated child events retain their setup route in the immutable journal, +/// while their delivery checkpoint belongs to the resolved parent route. +/// Recover against that effective route without rewriting the captured event. +pub(super) async fn recovered_delivery_route( + daemon: &Arc, + env: &Envelope, +) -> Option { + let key = automatic_link_key(env); + if let Some(route) = daemon.automatic_links.lock().unwrap().get(&key).cloned() { + return Some(route); + } + if let Some(route) = daemon + .pending_sessions + .lock() + .unwrap() + .get(&key) + .and_then(|state| state.linked_route.clone()) + { + return Some(route); + } + read_correlation_state(&daemon.data_dir, &key) + .await + .and_then(|state| state.linked_route) + .or_else(|| env.route.clone()) +} + +impl Daemon { + pub(super) fn correlation_lock(&self, key: &str) -> Arc> { + let mut locks = self.correlation_locks.lock().unwrap(); + locks + .entry(key.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + + pub(super) fn correlation_decided(&self, key: &str) -> bool { + self.standalone_links.lock().unwrap().contains(key) + || self.automatic_links.lock().unwrap().contains_key(key) + } +} diff --git a/bt-daemon/src/server/ingress.rs b/bt-daemon/src/server/ingress.rs new file mode 100644 index 0000000..7047fe9 --- /dev/null +++ b/bt-daemon/src/server/ingress.rs @@ -0,0 +1,342 @@ +//! The ingress queue between hook capture and session dispatch. Capture only +//! journals and enqueues; the worker routes events to session actors, using +//! the journals as overflow when the queue is full. + +use super::*; + +pub(super) enum IngressMsg { + Event(Box), + Barrier(oneshot::Sender<()>), +} + +/// The hook path never waits for this queue. Once it fills, the append-only +/// journals become the overflow queue and the worker catches them up in place. +pub(super) const INGRESS_QUEUE_CAPACITY: usize = 64; + +pub(super) fn spawn_ingress_worker(daemon: Arc, mut rx: mpsc::Receiver) { + tokio::spawn(async move { + loop { + if daemon.ingress_overflow.swap(false, Ordering::AcqRel) { + recover_ingress_overflow(&daemon).await; + continue; + } + let Some(msg) = rx.recv().await else { + break; + }; + match msg { + IngressMsg::Event(event) => { + dispatch_ingress_event(&daemon, *event).await; + } + IngressMsg::Barrier(reply) => { + loop { + daemon.settle_session_actors().await; + let _ = retry_pending_sessions(&daemon).await; + if !daemon.ingress_overflow.swap(false, Ordering::AcqRel) { + break; + } + recover_ingress_overflow(&daemon).await; + } + let _ = reply.send(()); + } + } + } + }); +} + +pub(super) async fn dispatch_ingress_event(daemon: &Arc, event: PendingEvent) { + let ingress_key = + event.env.route.as_ref().and_then(|route| { + DeliveryKey::new(&event.env.source, &event.env.session_id, route).ok() + }); + if ingress_key.as_ref().is_some_and(|key| { + daemon + .ingress_dispatched + .lock() + .unwrap() + .get(key) + .is_some_and(|through| *through >= event.journal_through) + }) { + return; + } + + let journal_through = event.journal_through; + // A child may arrive immediately after a parent's tool hook. Catch prior + // session actors up before resolving a new session, entirely in the daemon. + // Only an undecided session resolves a parent. Antigravity marks every + // invocation as a start, so skip the all-session barrier once decided. + if is_session_start(&event.env.source, &event.env.event) + && !daemon.correlation_decided(&automatic_link_key(&event.env)) + { + daemon.settle_session_actors().await; + } + match accept_event(daemon, event).await { + Ok(()) => { + if let Some(key) = ingress_key { + let mut dispatched = daemon.ingress_dispatched.lock().unwrap(); + let through = dispatched.entry(key).or_default(); + *through = (*through).max(journal_through); + } + } + Err(error) => { + tracing::warn!(%error, "journaled ingress event could not be dispatched"); + } + } +} + +/// Drain events omitted from the bounded in-memory queue. The queue contains +/// only a latency fast path; journals remain the complete source of ingress. +pub(super) async fn recover_ingress_overflow(daemon: &Arc) { + let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { + return; + }; + let mut recovered = 0usize; + while let Ok(Some(entry)) = entries.next_entry().await { + let path = entry.path(); + if path.extension().and_then(|ext| ext.to_str()) != Some("ndjson") { + continue; + } + let recorded_len = journal::JournalReader::recorded_len(&path).await; + let Ok(Some(mut checkpoints)) = journal::JournalReader::open(&path, recorded_len).await + else { + continue; + }; + let mut acknowledged_by_route: HashMap = HashMap::new(); + while let Ok(Some(entry)) = checkpoints.next_record().await { + if let journal::JournalRecord::DeliveryCheckpoint { route, through } = entry.record { + let key = serde_json::to_string(&route).unwrap_or_default(); + let acknowledged = acknowledged_by_route.entry(key).or_default(); + *acknowledged = (*acknowledged).max(through); + } + } + + let Ok(Some(mut reader)) = journal::JournalReader::open(&path, recorded_len).await else { + continue; + }; + let mut before = 0u64; + while let Ok(Some(entry)) = reader.next_record().await { + let through = entry.through; + let journal::JournalRecord::Event(redacted) = entry.record else { + before = through; + continue; + }; + let mut env = journal::envelope_from_redacted(redacted); + let Some(source) = daemon.translators.canonical_source(&env.source) else { + before = through; + continue; + }; + env.source = source.to_string(); + let Some(route) = env.route.as_ref() else { + before = through; + continue; + }; + let route_json = recovered_delivery_route(daemon, &env) + .await + .as_ref() + .and_then(|route| serde_json::to_string(route).ok()) + .unwrap_or_else(|| serde_json::to_string(route).unwrap_or_default()); + let key = match DeliveryKey::new(&env.source, &env.session_id, route) { + Ok(key) => key, + Err(_) => { + before = through; + continue; + } + }; + let dispatched = daemon + .ingress_dispatched + .lock() + .unwrap() + .get(&key) + .copied() + .unwrap_or(0); + let acknowledged = acknowledged_by_route.get(&route_json).copied().unwrap_or(0); + if through > dispatched.max(acknowledged) { + dispatch_ingress_event( + daemon, + PendingEvent { + env, + replay_through: before, + journal_through: through, + }, + ) + .await; + recovered += 1; + } + before = through; + } + } + if recovered > 0 { + tracing::info!(recovered, "drained journal-backed ingress overflow"); + } +} + +pub(super) fn spawn_pending_reconciler(daemon: Arc) { + tokio::spawn(async move { + let mut tick = tokio::time::interval(Duration::from_millis(100)); + loop { + tokio::select! { + _ = daemon.shutdown.notified() => return, + _ = daemon.correlation_changed.notified() => {}, + _ = tick.tick() => {}, + } + daemon.observe_claude_transcripts(None, false).await; + if let Err(error) = retry_pending_sessions(&daemon).await { + tracing::warn!(%error, "pending child-session reconciliation failed"); + } + } + }); +} + +pub(super) async fn accept_resolved_event( + daemon: &Arc, + event: PendingEvent, +) -> Result<(), String> { + let PendingEvent { + mut env, + replay_through, + journal_through, + } = event; + // Every fresh Claude capture carries a bounded observation marker. Older + // persisted ingress must never consume live bytes while rebuilding state. + if env.source == "claude-code" && env.payload.get("_bt_claude_transcript_mirrors").is_none() { + if let Some(payload) = env.payload.as_object_mut() { + payload.insert("_bt_transcript_replay".to_string(), serde_json::json!(true)); + } + } + let schedule_flush = crate::hook::should_flush_ingress_event(&env) + || (env.source == "claude-code" && env.event == "TranscriptUpdate"); + let source = env.source.clone(); + 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"); + let session_lock = daemon.dispatch_lock(&source, &session_id); + let _session_guard = session_lock.lock().await; + + let result = async { + let delivery_key = daemon + .configure_event(&mut env) + .await + .map_err(|error| format!("session auth failed: {error}"))?; + if crate::dispatch::is_tool_lifecycle_event(&env.event) { + if let Err(error) = mark_active_parent_snapshot_dirty( + &daemon.data_dir, + &delivery_key.correlation_key(), + &daemon.correlation, + ) + .await + { + tracing::warn!(session_id = %env.session_id, %error, "active parent snapshot could not be marked dirty"); + } + } + let session = daemon + .session_for(&env, &delivery_key, replay_through) + .await + .map_err(|error| format!("session init failed: {error}"))?; + daemon + .correlation + .observe_session(&delivery_key.correlation_key(), &env.source, env.capture.as_ref(), + is_session_start(&env.source, &env.event)); + session + .enqueue(env, journal_through) + .await + .map_err(|error| format!("enqueue failed: {error}"))?; + Ok(delivery_key) + } + .await; + + match &result { + 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"), + } + if result.is_ok() && schedule_flush { + // Deferred terminal events need the same flush as immediate events; + // there may have been no session actor when ingress first accepted them. + daemon + .enqueue_source_session_flush(&source, &session_id, Duration::from_secs(10)) + .await; + } + result.map(|_| ()) +} + +pub(super) fn is_session_start(source: &str, event: &str) -> bool { + matches!(event, "SessionStart" | "session_start" | "session.created") + || (source == "cursor" && event == "sessionStart") + // Antigravity has no dedicated session-start or resume hook. Its + // PreInvocation is the earliest hook in each invocation and is the + // only opportunity to refresh a resumed conversation's process. + || (source == "antigravity" && event == "PreInvocation") +} + +impl Daemon { + pub(super) async fn capture_event(&self, mut env: Envelope) -> Result { + let _capture_guard = self.capture_gate.read().await; + if self.quiescing.load(Ordering::SeqCst) { + return Ok(false); + } + env.source = self + .translators + .canonical_source(&env.source) + .ok_or_else(|| format!("unsupported coding-agent source {:?}", env.source))? + .to_string(); + self.touch(); + let lock = self.session_lock(&env.source, &env.session_id); + let _guard = lock.lock().await; + hydrate_transcript_reference(&self.data_dir, &mut env).await; + self.journal_and_enqueue(env).await + } + + /// Caller holds the source-session capture lock; frozen observations enter + /// the same durable queue as hooks without reacquiring that lock. + pub(super) async fn journal_and_enqueue(&self, mut env: Envelope) -> Result { + let (replay_through, journal_through) = self + .append_to_journal(&mut env) + .await + .map_err(|error| format!("journal failed: {error}"))?; + self.remember_claude_transcripts(&env); + match self + .ingress_tx + .try_send(IngressMsg::Event(Box::new(PendingEvent { + env, + replay_through, + journal_through, + }))) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + self.ingress_overflow.store(true, Ordering::Release); + tracing::debug!("ingress queue full; journal will be drained by daemon worker"); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + tracing::warn!("journaled event will be recovered after daemon restart"); + } + } + Ok(true) + } + + pub(super) async fn ingress_barrier(&self) { + let (tx, rx) = oneshot::channel(); + if self.ingress_tx.send(IngressMsg::Barrier(tx)).await.is_ok() { + let _ = rx.await; + } + } + + pub(super) async fn settle_ingress(self: &Arc) { + self.ingress_barrier().await; + self.settle_session_actors().await; + let _ = retry_pending_sessions(self).await; + } + + pub(super) async fn settle_session_actors(&self) { + let sessions: Vec<_> = self.sessions.lock().unwrap().values().cloned().collect(); + for session in sessions { + session.barrier().await; + } + } +} diff --git a/bt-daemon/src/server/maintenance.rs b/bt-daemon/src/server/maintenance.rs new file mode 100644 index 0000000..d460cd0 --- /dev/null +++ b/bt-daemon/src/server/maintenance.rs @@ -0,0 +1,150 @@ +//! Background maintenance: retention GC, idle-session retirement, the idle +//! watchdog, and draining on shutdown. + +use super::*; + +/// How long recovery state is kept on disk. +pub(super) const RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60); + +/// How often that retention is enforced while the daemon keeps running. +pub(super) const GC_INTERVAL: Duration = Duration::from_secs(60 * 60); + +pub(super) async fn collect_garbage(data_dir: &std::path::Path) { + journal::gc_old_journals(data_dir, RETENTION).await; + journal::gc_old_managed_runs(data_dir, RETENTION).await; + crate::transcript_mirror::gc_old_mirrors(data_dir, RETENTION).await; + gc_old_correlation_states(data_dir, RETENTION).await; +} + +pub(super) async fn gc_old_correlation_states(data_dir: &std::path::Path, max_age: Duration) { + gc_old_correlation_dir(&data_dir.join("correlation"), max_age).await; + gc_old_correlation_dir(&data_dir.join("correlation").join("parents"), max_age).await; +} + +pub(super) async fn gc_old_correlation_dir(dir: &std::path::Path, max_age: Duration) { + let Ok(mut entries) = tokio::fs::read_dir(dir).await else { + return; + }; + let now = SystemTime::now(); + while let Ok(Some(entry)) = entries.next_entry().await { + let old = entry + .metadata() + .await + .ok() + .and_then(|metadata| metadata.modified().ok()) + .and_then(|modified| now.duration_since(modified).ok()) + .is_some_and(|age| age > max_age); + if old && entry.file_type().await.is_ok_and(|kind| kind.is_file()) { + let _ = tokio::fs::remove_file(entry.path()).await; + } + } +} + +/// Retire delivery pipelines that have gone quiet. Without this, every session +/// the daemon ever saw keeps its translator state, sink handles, and journal +/// file handle alive until the process exits — which for a continuously busy +/// user is never, since the idle watchdog needs *all* sessions quiet. +pub(super) fn spawn_session_reaper(daemon: Arc, idle_timeout: Duration) { + if idle_timeout.is_zero() { + return; // 0 disables retirement (useful in tests) + } + tokio::spawn(async move { + let tick = (idle_timeout / 4).max(Duration::from_secs(1)); + loop { + tokio::select! { + _ = daemon.shutdown.notified() => return, + _ = tokio::time::sleep(tick) => {} + } + for key in daemon.idle_sessions(idle_timeout) { + daemon.retire_session(&key).await; + } + } + }); +} + +pub(super) fn spawn_gc(daemon: Arc) { + tokio::spawn(async move { + loop { + tokio::select! { + _ = daemon.shutdown.notified() => return, + _ = tokio::time::sleep(GC_INTERVAL) => {} + } + collect_garbage(&daemon.data_dir).await; + daemon + .correlation + .prune_retired(crate::process::process_is_alive); + } + }); +} + +pub(super) fn spawn_idle_watchdog(daemon: Arc, idle_timeout: Duration) { + if idle_timeout.is_zero() { + return; // 0 disables the watchdog (useful in tests) + } + tokio::spawn(async move { + let tick = (idle_timeout / 4).max(Duration::from_secs(1)); + loop { + tokio::select! { + _ = daemon.shutdown.notified() => return, + _ = tokio::time::sleep(tick) => {} + } + let idle_for = daemon.last_activity.lock().unwrap().elapsed(); + if idle_for >= idle_timeout + && daemon.total_queued() == 0 + && !daemon.correlation.has_any_active_tools() + { + tracing::info!("idle for {:?}; shutting down", idle_for); + daemon.trigger_shutdown(); + return; + } + } + }); +} + +pub(super) async fn drain_all(daemon: &Arc) { + let mut drained = daemon.drained.lock().await; + if *drained { + return; + } + daemon.begin_quiesce(); + let _capture_guard = daemon.capture_gate.write().await; + daemon.observe_claude_transcripts(None, true).await; + daemon.settle_ingress().await; + if let Err(error) = resolve_pending_sessions(daemon, |_| true).await { + tracing::warn!(%error, "pending correlation events retained for recovery during shutdown"); + } + let sessions: Vec> = daemon.sessions.lock().unwrap().values().cloned().collect(); + for s in sessions { + s.shutdown().await; + } + *drained = true; +} + +/// Is a live daemon answering at the endpoint? Connect and expect any line +/// back from a well-formed `initialize`. +pub(super) async fn probe_alive(endpoint: &std::path::Path) -> bool { + let Ok(stream) = transport::connect(endpoint).await else { + return false; + }; + let (read_half, mut write_half) = tokio::io::split(stream); + let init = Request::new( + crate::wire::RequestId::Int(0), + method::INITIALIZE, + serde_json::json!({ + "protocol_version": PROTOCOL_VERSION, + "client": { "source": "probe" } + }), + ); + let Ok(mut line) = Message::Request(init).to_line() else { + return false; + }; + line.push('\n'); + if write_half.write_all(line.as_bytes()).await.is_err() { + return false; + } + let mut lines = BufReader::new(read_half).lines(); + matches!( + tokio::time::timeout(Duration::from_secs(1), lines.next_line()).await, + Ok(Ok(Some(_))) + ) +} diff --git a/bt-daemon/src/server/managed_run.rs b/bt-daemon/src/server/managed_run.rs new file mode 100644 index 0000000..3de2997 --- /dev/null +++ b/bt-daemon/src/server/managed_run.rs @@ -0,0 +1,189 @@ +//! Managed-run bookkeeping: which sessions a managed child produced, and +//! flushing all of them when the run exits. + +use super::*; + +impl Daemon { + pub(super) fn record_managed_run_session( + &self, + managed_run_id: &str, + key: &DeliveryKey, + ) -> bool { + self.managed_run_sessions + .lock() + .unwrap() + .entry(managed_run_id.to_string()) + .or_default() + .insert(key.clone()) + } + + pub(super) 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 { + source: Some(key.source.clone()), + 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" + ); + } + } + + pub(super) async fn flush_managed_run(&self, params: ManagedRunFlushParams) -> FlushResult { + 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 Some(source) = record.source { + if let Ok(key) = DeliveryKey::new(&source, &record.session_id, &record.route) { + delivery_keys.insert(key); + } + } else { + // Legacy records predate source-qualified delivery keys. A + // matching live pipeline is authoritative; otherwise recover + // the source from the legacy journal. + let live: Vec<_> = self + .sessions + .lock() + .unwrap() + .keys() + .filter(|key| { + key.session_id == record.session_id + && serde_json::from_str::(&key.route) + .is_ok_and(|route| route.same_route(&record.route)) + }) + .cloned() + .collect(); + if live.is_empty() { + if let Some(source) = journal::legacy_journal_source( + &self.data_dir, + &record.session_id, + &record.route, + ) + .await + .and_then(|source| { + self.translators + .canonical_source(&source) + .map(str::to_owned) + }) { + if let Ok(key) = + DeliveryKey::new(&source, &record.session_id, &record.route) + { + delivery_keys.insert(key); + } + } + } else { + delivery_keys.extend(live); + } + } + } + let mut result = FlushResult { + flushed: true, + pending: 0, + accepted_sessions: delivery_keys.len() as u64, + }; + for key in &delivery_keys { + self.observe_claude_transcripts(Some(key), false).await; + } + // Managed completion is terminal. Stop polling before the barrier so + // an in-flight observer cannot enqueue after translator finalization. + for key in &delivery_keys { + let lock = self.session_lock(&key.source, &key.session_id); + let _guard = lock.lock().await; + self.claude_transcripts + .lock() + .unwrap() + .retain(|candidate, observation| { + candidate.source != key.source + || candidate.session_id != key.session_id + || observation.template.managed_run_id.as_deref() + != Some(params.managed_run_id.as_str()) + }); + } + self.ingress_barrier().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 = %key.session_id, + %error, + "managed run session auth refresh failed" + ); + result.flushed = false; + continue; + } + let session = { self.sessions.lock().unwrap().get(&key).cloned() }; + if let Some(session) = session { + let (flushed, pending) = session + .finalize(Duration::from_millis(params.timeout_ms)) + .await; + result.flushed &= flushed; + result.pending = result.pending.saturating_add(pending); + } + } + result + } + + /// Queue a flush on every live delivery route before dispatching more ingress. + /// Only the completion wait runs in a separate task. This is used + /// only by the daemon worker after a turn-ending event has already been + /// durably captured and acknowledged to the hook client. + pub(super) async fn enqueue_source_session_flush( + &self, + source: &str, + session_id: &str, + timeout: Duration, + ) { + let sessions: Vec<_> = self + .sessions + .lock() + .unwrap() + .iter() + .filter(|(key, _)| key.source == source && key.session_id == session_id) + .map(|(_, session)| session.clone()) + .collect(); + for session in sessions { + let reply = match session.enqueue_flush().await { + Ok(reply) => reply, + Err(error) => { + tracing::warn!(source, session_id, %error, "turn-end flush enqueue failed"); + continue; + } + }; + let source = source.to_owned(); + let session_id = session_id.to_owned(); + tokio::spawn(async move { + // Later events may remain queued after this boundary has flushed. + // Their presence does not make this boundary's flush incomplete. + if !matches!(tokio::time::timeout(timeout, reply).await, Ok(Ok(_))) { + tracing::warn!( + source, + session_id, + "out-of-band turn-end flush did not complete" + ); + } + }); + } + } +} diff --git a/bt-daemon/src/server/mod.rs b/bt-daemon/src/server/mod.rs new file mode 100644 index 0000000..9e076ca --- /dev/null +++ b/bt-daemon/src/server/mod.rs @@ -0,0 +1,344 @@ +//! The daemon: owns the session map + shared deps, binds the UDS listener, +//! serves JSON-RPC connections, and shuts down gracefully (idle timeout, +//! `daemon.shutdown`, or SIGINT/SIGTERM). + +mod auth; +mod claude_transcripts; +mod correlation_store; +mod ingress; +mod maintenance; +mod managed_run; +mod recovery; +mod rpc; +mod sessions; + +pub(crate) use auth::resolve_route_auth; +use claude_transcripts::*; +pub(crate) use correlation_store::persist_active_parent_snapshot; +use correlation_store::*; +use ingress::*; +use maintenance::*; +use recovery::*; +use rpc::*; +use sessions::*; + +use crate::dispatch::{hydrate_transcript_reference, ReplayPlan, Session, SessionOptions}; +use crate::journal::{self, JournalWriter}; +use crate::sink::{BraintrustSinkConfig, BraintrustSinkFactory, DebugSinkFactory, SinkFactory}; +use crate::translate::Registry; +use crate::transport::{self, Listener, ServerStream}; +use crate::wire::{ + error_code, method, AuthDiagnoseParams, AuthDiagnoseResult, Capabilities, ClientInfo, Envelope, + EventLogResult, FlushParams, FlushResult, InitializeParams, InitializeResult, + ManagedRunFlushParams, 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 serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::sync::{mpsc, oneshot, Notify}; + +/// Injected dependencies for `serve`, so `bt` / tests can supply a sink +/// factory (Braintrust in production, debug in tests) and a version string. +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 non-secret selection is +/// observable; the backend credential remains in daemon memory and is never +/// journaled. +#[derive(Debug, Clone)] +pub struct AuthLease { + pub selection: AuthSelection, + pub auth: BackendAuth, + /// Epoch milliseconds. `None` is appropriate for non-expiring API keys. + pub expires_at_ms: Option, +} + +// `async_trait` marks its boxed futures as `must_use`; Clippy 1.99 flags that +// generated annotation as redundant on async trait methods. +#[allow(clippy::double_must_use)] +#[async_trait] +pub trait AuthProvider: Send + Sync { + /// Resolve without prompting. On refresh, `selection` is the canonical + /// source returned by the initial lease, keeping an active session pinned + /// even if the user's default profile or environment changes. + async fn resolve( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result; +} + +#[derive(Clone)] +struct SessionAuthState { + route: SessionRoute, + lease: AuthLease, +} + +pub struct Daemon { + version: String, + data_dir: PathBuf, + translators: Arc, + sink_factory: Arc, + auth_provider: Option>, + session_auth: tokio::sync::Mutex>, + route_aliases: Mutex>, + /// Serializes only capture-side journal appends for a source session. + session_locks: Mutex>>>, + /// Serializes daemon-side routing and actor creation independently of + /// capture, which must never inherit actor or sink backpressure. + dispatch_locks: Mutex>>>, + journals: Mutex>>>, + claude_transcripts: Mutex>, + managed_run_sessions: Mutex>>, + auth_errors: Mutex>, + sessions: Mutex>>, + correlation: Arc, + automatic_links: Mutex>, + standalone_links: Mutex>, + pending_sessions: Mutex>, + pending_reconcile_lock: tokio::sync::Mutex<()>, + correlation_locks: Mutex>>>, + correlation_changed: Arc, + ingress_tx: mpsc::Sender, + ingress_overflow: AtomicBool, + ingress_dispatched: Mutex>, + started: Instant, + last_activity: Mutex, + /// Prevents a shutdown drain from racing capture that has passed the + /// quiescing check but has not yet appended and queued its event. + capture_gate: tokio::sync::RwLock<()>, + quiescing: AtomicBool, + drained: tokio::sync::Mutex, + shutting_down: AtomicBool, + shutdown: Notify, +} + +impl Daemon { + fn new(opts: ServeOptions, data_dir: PathBuf) -> Arc { + let (ingress_tx, ingress_rx) = mpsc::channel(INGRESS_QUEUE_CAPACITY); + let daemon = Arc::new(Daemon { + version: opts.version, + data_dir, + translators: opts.translators, + sink_factory: opts.sink_factory, + auth_provider: opts.auth_provider, + session_auth: tokio::sync::Mutex::new(HashMap::new()), + route_aliases: Mutex::new(HashMap::new()), + session_locks: Mutex::new(HashMap::new()), + dispatch_locks: Mutex::new(HashMap::new()), + journals: Mutex::new(HashMap::new()), + claude_transcripts: Mutex::new(HashMap::new()), + managed_run_sessions: Mutex::new(HashMap::new()), + auth_errors: Mutex::new(HashMap::new()), + sessions: Mutex::new(HashMap::new()), + correlation: Arc::new(crate::correlation::CorrelationRegistry::default()), + automatic_links: Mutex::new(HashMap::new()), + standalone_links: Mutex::new(HashSet::new()), + pending_sessions: Mutex::new(HashMap::new()), + pending_reconcile_lock: tokio::sync::Mutex::new(()), + correlation_locks: Mutex::new(HashMap::new()), + correlation_changed: Arc::new(Notify::new()), + ingress_tx, + ingress_overflow: AtomicBool::new(false), + ingress_dispatched: Mutex::new(HashMap::new()), + started: Instant::now(), + last_activity: Mutex::new(Instant::now()), + capture_gate: tokio::sync::RwLock::new(()), + quiescing: AtomicBool::new(false), + drained: tokio::sync::Mutex::new(false), + shutting_down: AtomicBool::new(false), + shutdown: Notify::new(), + }); + spawn_ingress_worker(daemon.clone(), ingress_rx); + daemon + } + + fn touch(&self) { + *self.last_activity.lock().unwrap() = Instant::now(); + } + + /// Queued work, excluding pending sessions whose handoff failed. Those are + /// durable on disk and resume after a restart, so they must not keep an + /// otherwise idle daemon alive forever. + fn total_queued(&self) -> u64 { + let dispatched: u64 = self + .sessions + .lock() + .unwrap() + .values() + .map(|s| s.counters.queued.load(Ordering::Relaxed)) + .sum(); + dispatched.saturating_add( + self.pending_sessions + .lock() + .unwrap() + .values() + .filter(|state| !state.delivery_failed) + .map(|state| state.events.len() as u64) + .sum(), + ) + } + + fn trigger_shutdown(&self) { + self.quiescing.store(true, Ordering::SeqCst); + self.shutting_down.store(true, Ordering::SeqCst); + self.shutdown.notify_waiters(); + } + + fn begin_quiesce(&self) { + self.quiescing.store(true, Ordering::SeqCst); + } +} + +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<()> { + let socket = paths::socket_path(args.socket.as_deref()); + let data_dir = paths::data_dir(args.data_dir.as_deref()); + paths::ensure_private_dir(&data_dir)?; + #[cfg(unix)] + if let Some(parent) = socket.parent() { + paths::ensure_private_dir(parent)?; + } + + let listener = match transport::claim(&socket, || probe_alive(&socket)).await? { + Some(l) => l, + None => { + tracing::info!( + "another daemon is already serving {}; exiting", + socket.display() + ); + return Ok(()); + } + }; + tracing::info!(socket = %socket.display(), "bt-daemon listening"); + + let daemon = Daemon::new(opts, data_dir); + collect_garbage(&daemon.data_dir).await; + restore_active_parent_snapshots(&daemon.data_dir, &daemon.correlation).await; + restore_pending_sessions(&daemon).await; + recover_unprocessed_journals(&daemon).await; + spawn_pending_reconciler(daemon.clone()); + let idle_timeout = Duration::from_secs(args.idle_timeout_secs); + spawn_idle_watchdog(daemon.clone(), idle_timeout); + spawn_session_reaper( + daemon.clone(), + Duration::from_secs(args.session_idle_timeout_secs), + ); + spawn_gc(daemon.clone()); + + let accept_result = accept_loop(daemon.clone(), listener).await; + + // Graceful drain regardless of why we stopped. + drain_all(&daemon).await; + transport::cleanup(&socket); + accept_result +} + +/// Run the daemon until shutdown. +pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { + run(args, opts).await +} + +/// Build debug [`ServeOptions`]: the production translator registry plus a +/// debug sink writing NDJSON under `/spans/`. +pub fn debug_serve_options(version: impl Into, data_dir: &std::path::Path) -> ServeOptions { + ServeOptions { + version: version.into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(DebugSinkFactory { + dir: data_dir.join("spans"), + }), + auth_provider: None, + } +} + +/// Build [`ServeOptions`] with the Braintrust sink. `translators` lets the +/// caller choose the translator registry (debug-only until Phase 3 adds the +/// Codex/Claude translators). Clients are built lazily per session URL, so this +/// is cheap and infallible. +pub fn braintrust_serve_options( + version: impl Into, + sink_config: BraintrustSinkConfig, + translators: Arc, +) -> ServeOptions { + ServeOptions { + version: version.into(), + translators, + sink_factory: Arc::new(BraintrustSinkFactory::new(sink_config)), + auth_provider: None, + } +} + +#[cfg(test)] +mod tests { + use super::{client_may_shutdown, is_session_start, ClientInfo}; + + #[test] + fn antigravity_invocation_refreshes_a_resumed_conversations_process() { + assert!(is_session_start("antigravity", "PreInvocation")); + assert!(!is_session_start("claude-code", "PreInvocation")); + assert!(is_session_start("claude-code", "SessionStart")); + assert!(is_session_start("pi", "session_start")); + assert!(is_session_start("opencode", "session.created")); + } + + fn initialized(version: Option<&str>) -> ClientInfo { + ClientInfo { + source: "codex".into(), + daemon_version: version.map(str::to_string), + plugin_version: None, + pid: None, + } + } + + #[test] + fn legacy_initialized_clients_cannot_downgrade_the_daemon() { + assert!(client_may_shutdown(None, "0.20.0")); + assert!(!client_may_shutdown(Some(&initialized(None)), "0.20.0")); + assert!(!client_may_shutdown( + Some(&initialized(Some("0.19.3"))), + "0.20.0" + )); + assert!(client_may_shutdown( + Some(&initialized(Some("0.20.0"))), + "0.20.0" + )); + assert!(client_may_shutdown( + Some(&initialized(Some("0.21.0"))), + "0.20.0" + )); + } +} diff --git a/bt-daemon/src/server/recovery.rs b/bt-daemon/src/server/recovery.rs new file mode 100644 index 0000000..d93138b --- /dev/null +++ b/bt-daemon/src/server/recovery.rs @@ -0,0 +1,133 @@ +//! Startup recovery of journaled events the previous daemon did not finish. + +use super::*; + +pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { + let pending: HashSet<(String, String, u64)> = daemon + .pending_sessions + .lock() + .unwrap() + .values() + .flat_map(|state| { + state.events.iter().map(|event| { + ( + event.env.source.clone(), + event.env.session_id.clone(), + event.journal_through, + ) + }) + }) + .collect(); + let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { + return; + }; + let mut candidates = Vec::new(); + while let Ok(Some(entry)) = entries.next_entry().await { + let path = entry.path(); + if path.extension().and_then(|ext| ext.to_str()) != Some("ndjson") { + continue; + } + let recorded_len = journal::JournalReader::recorded_len(&path).await; + let Ok(Some(mut reader)) = journal::JournalReader::open(&path, recorded_len).await else { + continue; + }; + let mut before = 0u64; + let mut acknowledged_by_route: HashMap = HashMap::new(); + let mut latest_by_route: HashMap = HashMap::new(); + while let Ok(Some(entry)) = reader.next_record().await { + let through = entry.through; + match entry.record { + journal::JournalRecord::Event(redacted) => { + let mut env = journal::envelope_from_redacted(redacted); + let Some(source) = daemon.translators.canonical_source(&env.source) else { + before = through; + continue; + }; + env.source = source.to_string(); + daemon.remember_claude_transcripts(&env); + if env.source == "claude-code" { + if let Some(payload) = env.payload.as_object_mut() { + payload.insert( + "_bt_transcript_replay".to_string(), + serde_json::json!(true), + ); + } + } + if let Some(route) = env.route.as_ref() { + let key = serde_json::to_string(route).unwrap_or_default(); + latest_by_route.insert( + key, + PendingEvent { + env, + replay_through: before, + journal_through: through, + }, + ); + } + } + journal::JournalRecord::DeliveryCheckpoint { route, through } => { + let key = serde_json::to_string(&route).unwrap_or_default(); + let acknowledged = acknowledged_by_route.entry(key).or_default(); + *acknowledged = (*acknowledged).max(through); + } + } + before = through; + } + for event in latest_by_route.into_values() { + if pending.contains(&( + event.env.source.clone(), + event.env.session_id.clone(), + event.journal_through, + )) { + continue; + } + let route = recovered_delivery_route(daemon, &event.env) + .await + .as_ref() + .and_then(|route| serde_json::to_string(route).ok()) + .unwrap_or_default(); + if acknowledged_by_route.get(&route).copied().unwrap_or(0) < event.journal_through { + candidates.push(event); + } else if event.env.source == "claude-code" { + // A fully delivered historical session has no actor to retire + // its observer. Probe once for bytes appended after the final + // checkpoint, retaining only paths that have new work. + if let Some(key) = event.env.route.as_ref().and_then(|route| { + DeliveryKey::new(&event.env.source, &event.env.session_id, route).ok() + }) { + let observation = daemon.claude_transcripts.lock().unwrap().get(&key).cloned(); + if let Some(observation) = observation { + let path = observation.template.payload["transcript_path"] + .as_str() + .unwrap(); + let changed = tokio::fs::metadata(path) + .await + .is_ok_and(|metadata| observation.through != Some(metadata.len())); + if !changed { + daemon.claude_transcripts.lock().unwrap().remove(&key); + } + } + } + } + } + } + candidates.sort_by_key(|event| event.env.ts_ms); + let recovered = candidates.len(); + for event in candidates { + let _ = daemon + .ingress_tx + .send(IngressMsg::Event(Box::new(event))) + .await; + } + // Reconciliation stays in the daemon worker and follows every recovered + // event in queue order. Dropping the receiver is intentional: startup and + // hook capture do not wait for translation or reporting. + let (reply, _ignored) = oneshot::channel(); + let _ = daemon.ingress_tx.send(IngressMsg::Barrier(reply)).await; + if recovered > 0 { + tracing::info!( + recovered, + "queued unprocessed journal sessions for recovery" + ); + } +} diff --git a/bt-daemon/src/server/rpc.rs b/bt-daemon/src/server/rpc.rs new file mode 100644 index 0000000..74186b3 --- /dev/null +++ b/bt-daemon/src/server/rpc.rs @@ -0,0 +1,425 @@ +//! The JSON-RPC listener, request handling and status reporting. + +use super::*; + +pub(super) async fn accept_loop(daemon: Arc, mut listener: Listener) -> anyhow::Result<()> { + loop { + // `Notify::notify_waiters` does not retain a permit. If accepting a + // connection wins the select at the same time shutdown is requested, + // check the sticky flag before waiting again so the notification + // cannot be lost. + if daemon.shutting_down.load(Ordering::SeqCst) { + tracing::info!("shutdown requested"); + return Ok(()); + } + tokio::select! { + _ = daemon.shutdown.notified() => { + tracing::info!("shutdown requested"); + return Ok(()); + } + _ = tokio::signal::ctrl_c() => { + tracing::info!("interrupt received"); + return Ok(()); + } + accepted = listener.accept() => { + match accepted { + Ok(stream) => { + let d = daemon.clone(); + tokio::spawn(async move { + if let Err(e) = serve_connection(d, stream).await { + tracing::debug!("connection ended: {e}"); + } + }); + } + Err(e) => { + tracing::warn!("accept error: {e}"); + } + } + } + } + } +} + +pub(super) async fn serve_connection( + daemon: Arc, + stream: ServerStream, +) -> anyhow::Result<()> { + let (read_half, mut write_half) = tokio::io::split(stream); + let mut lines = BufReader::new(read_half).lines(); + let mut client = None; + + while let Some(line) = lines.next_line().await? { + if line.trim().is_empty() { + continue; + } + let mut shutdown_after_response = false; + let response = match Message::from_line(&line) { + Ok(Message::Request(req)) => { + let request_id = req.id.clone(); + let method = req.method.clone(); + tracing::info!( + request_id = ?request_id, + method, + "request received" + ); + let response = handle_request(&daemon, req, &mut client).await; + shutdown_after_response = method == method::DAEMON_SHUTDOWN + && response.error.is_none() + && response + .result + .as_ref() + .and_then(|result| result.get("ok")) + .and_then(Value::as_bool) + == Some(true); + if let Some(error) = &response.error { + tracing::warn!( + request_id = ?request_id, + method, + error_code = error.code, + error = %error.message, + "request failed" + ); + } else { + tracing::info!( + request_id = ?request_id, + method, + "request completed" + ); + } + Some(response) + } + Ok(Message::Notification(note)) => { + tracing::info!(method = %note.method, "notification received"); + // Legacy best-effort notifications cannot participate in the + // acknowledged handover retry used by current clients. + if note.method == method::EVENT_LOG { + if let Some(params) = note.params { + match serde_json::from_value::(params) { + Ok(mut env) => { + attach_process_capture(&daemon, &mut env, client.as_ref()); + match daemon.capture_event(env).await { + Ok(true) => {} + Ok(false) => tracing::warn!( + "event notification arrived while daemon was shutting down" + ), + Err(error) => { + tracing::warn!(%error, "event notification was not captured") + } + } + } + Err(error) => tracing::warn!( + method = %note.method, + error = %error, + "notification parameters rejected" + ), + } + } + } + None + } + Ok(Message::Response(_)) => None, // clients don't send us responses + Err(e) => { + tracing::warn!(error = %e, "request parse failed"); + Some(Response::err( + crate::wire::RequestId::Int(0), + RpcError::new(error_code::PARSE, format!("parse error: {e}")), + )) + } + }; + + if let Some(resp) = response { + let mut buf = Message::Response(resp).to_line()?; + buf.push('\n'); + let write_result = async { + write_half.write_all(buf.as_bytes()).await?; + write_half.flush().await + } + .await; + if shutdown_after_response { + daemon.trigger_shutdown(); + } + write_result?; + } + } + Ok(()) +} + +pub(super) fn attach_process_capture( + daemon: &Daemon, + env: &mut Envelope, + client: Option<&crate::wire::ClientInfo>, +) { + if env.capture.is_some() { + return; + } + if !is_session_start(&env.source, &env.event) + && !daemon + .correlation + .needs_process_capture(&env.source, &env.session_id) + { + return; + } + let Some(pid) = client.and_then(|client| client.pid) else { + return; + }; + let capture = crate::process::capture_process_context(pid); + if !capture.process_chain.is_empty() || capture.truncated { + env.capture = Some(capture); + } +} + +pub(super) async fn handle_request( + daemon: &Arc, + req: Request, + client: &mut Option, +) -> Response { + let id = req.id.clone(); + let params = req.params.unwrap_or(serde_json::Value::Null); + + macro_rules! parse { + ($t:ty) => { + match serde_json::from_value::<$t>(params) { + Ok(v) => v, + Err(e) => { + return Response::err( + id, + RpcError::new(error_code::INVALID_PARAMS, format!("invalid params: {e}")), + ) + } + } + }; + } + + match req.method.as_str() { + method::INITIALIZE => { + let p = parse!(InitializeParams); + if p.protocol_version != PROTOCOL_VERSION { + return Response::err( + id, + RpcError::new( + error_code::APP, + format!( + "protocol version mismatch: client {} daemon {}", + p.protocol_version, PROTOCOL_VERSION + ), + ), + ); + } + *client = Some(p.client); + let result = InitializeResult { + protocol_version: PROTOCOL_VERSION, + daemon_version: daemon.version.clone(), + capabilities: Capabilities { + sources: daemon.translators.sources(), + }, + }; + Response::ok(id, serde_json::to_value(result).unwrap()) + } + method::EVENT_LOG => { + let mut env = parse!(Envelope); + attach_process_capture(daemon, &mut env, client.as_ref()); + match daemon.capture_event(env).await { + Ok(accepted) => Response::ok( + id, + serde_json::to_value(EventLogResult { accepted }).unwrap(), + ), + Err(error) => Response::err(id, RpcError::new(error_code::INTERNAL, error)), + } + } + method::SESSION_FLUSH => { + let p = parse!(FlushParams); + // Explicit flushes wait for the daemon-owned ingress queue. Hook + // capture never waits on this barrier. + daemon.settle_ingress().await; + if let Err(error) = + resolve_pending_sessions(daemon, |env| env.session_id == p.session_id).await + { + return Response::err(id, RpcError::new(error_code::INTERNAL, 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); + } + } + Response::ok( + id, + serde_json::to_value(FlushResult { + flushed, + pending, + accepted_sessions, + }) + .unwrap(), + ) + } + method::MANAGED_RUN_FLUSH => { + let params = parse!(ManagedRunFlushParams); + daemon.settle_ingress().await; + if let Err(error) = resolve_pending_sessions(daemon, |env| { + env.managed_run_id.as_deref() == Some(params.managed_run_id.as_str()) + }) + .await + { + return Response::err(id, RpcError::new(error_code::INTERNAL, error)); + } + let result = daemon.flush_managed_run(params).await; + Response::ok(id, serde_json::to_value(result).unwrap()) + } + method::STATUS_GET => { + let p = parse!(StatusParams); + daemon.settle_ingress().await; + Response::ok(id, serde_json::to_value(daemon.status(p)).unwrap()) + } + method::AUTH_DIAGNOSE => { + let p = parse!(AuthDiagnoseParams); + Response::ok( + id, + serde_json::to_value(daemon.diagnose_auth(p.auth).await).unwrap(), + ) + } + method::DAEMON_SHUTDOWN => { + if !client_may_shutdown(client.as_ref(), &daemon.version) { + return Response::ok( + id, + serde_json::to_value(ShutdownResult { ok: false }).unwrap(), + ); + } + drain_all(daemon).await; + Response::ok( + id, + serde_json::to_value(ShutdownResult { ok: true }).unwrap(), + ) + } + other => Response::err( + id, + RpcError::new( + error_code::METHOD_NOT_FOUND, + format!("unknown method: {other}"), + ), + ), + } +} + +pub(super) fn client_may_shutdown(client: Option<&ClientInfo>, daemon_version: &str) -> bool { + let Some(client) = client else { + return true; // Explicit stop commands do not initialize first. + }; + let Some(client_version) = client.daemon_version.as_deref() else { + return false; // Legacy initialized hooks must not downgrade the daemon. + }; + crate::client::compare_daemon_versions(client_version, daemon_version) + .is_none_or(|ordering| !ordering.is_lt()) +} + +impl Daemon { + pub(super) fn status(&self, p: StatusParams) -> StatusResult { + let map = self.sessions.lock().unwrap(); + let mut sessions: Vec<_> = map + .iter() + .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 (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(key) { + continue; + } + sessions.push(SessionStatus { + session_id: key.session_id.clone(), + source: source.clone(), + route: serde_json::from_str(&key.route).ok(), + queued: 0, + spans_emitted: 0, + permalink: None, + last_error: Some(error.clone()), + }); + } + for state in self.pending_sessions.lock().unwrap().values() { + let Some(first) = state.events.first() else { + continue; + }; + if p.session_id + .as_ref() + .is_some_and(|want| want != &first.env.session_id) + { + continue; + } + let route = state.linked_route.as_ref().or(first.env.route.as_ref()); + let existing = sessions.iter_mut().find(|session| { + session.session_id == first.env.session_id + && session.source == first.env.source + && match (session.route.as_ref(), route) { + (Some(left), Some(right)) => left.same_route(right), + (None, None) => true, + _ => false, + } + }); + let message = if state.decided() { + "correlation resolved; journaled events awaiting delivery" + } else { + "awaiting parent process confirmation; standalone fallback within 5 seconds" + }; + if let Some(session) = existing { + session.queued = session.queued.saturating_add(state.events.len() as u64); + if session.last_error.is_none() { + session.last_error = Some(message.to_string()); + } + } else { + sessions.push(SessionStatus { + session_id: first.env.session_id.clone(), + source: first.env.source.clone(), + route: route.cloned(), + queued: state.events.len() as u64, + spans_emitted: 0, + permalink: None, + last_error: Some(message.to_string()), + }); + } + } + StatusResult { + daemon_version: self.version.clone(), + uptime_ms: self.started.elapsed().as_millis() as u64, + sessions, + } + } +} diff --git a/bt-daemon/src/server/sessions.rs b/bt-daemon/src/server/sessions.rs new file mode 100644 index 0000000..d871f19 --- /dev/null +++ b/bt-daemon/src/server/sessions.rs @@ -0,0 +1,253 @@ +//! Delivery keys and the per-session actors, journal writers and locks. + +use super::*; + +/// 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)] +pub(super) struct DeliveryKey { + pub(super) source: String, + pub(super) session_id: String, + pub(super) route: String, +} + +impl DeliveryKey { + pub(super) fn new( + source: &str, + session_id: &str, + route: &SessionRoute, + ) -> anyhow::Result { + Ok(Self { + source: source.to_string(), + session_id: session_id.to_string(), + route: serde_json::to_string(route)?, + }) + } + + pub(super) fn correlation_key(&self) -> String { + format!( + "{}\u{1f}{}\u{1f}{}", + self.source, self.session_id, self.route + ) + } +} + +impl Daemon { + pub(super) async fn session_for( + &self, + env: &Envelope, + key: &DeliveryKey, + replay_through: u64, + ) -> anyhow::Result> { + { + let map = self.sessions.lock().unwrap(); + if let Some(session) = map.get(key) { + return Ok(session.clone()); + } + } + let route = env + .route + .as_ref() + .ok_or_else(|| anyhow::anyhow!("event is missing its session route"))?; + let config = env.config.clone().ok_or_else(|| { + anyhow::anyhow!("resolved route is missing its session configuration") + })?; + // The actor streams the journal itself, so creating a session stays + // cheap and allocation-free here no matter how long the recorded + // session is. Bound it to what is recorded now, before this event is + // appended, so replay covers prior source observations only. Every + // destination consumes that shared observation stream, with its own + // delivery checkpoint controlling which rows reach its sink. + let journal_path = + journal::ensure_source_journal(&self.data_dir, &env.source, &env.session_id).await?; + let translator_session_id = + if journal::legacy_journal_has_session(&self.data_dir, &env.source, &env.session_id) + .await + { + env.session_id.clone() + } else { + crate::ids::session_namespace(&env.source, &env.session_id) + }; + let replay = ReplayPlan { + acknowledged_through: journal::JournalReader::acknowledged_through( + &journal_path, + replay_through, + route, + ) + .await, + through: replay_through, + journal_path, + }; + let journal = self + .journal_writer_for(&env.source, &env.session_id) + .await?; + let mut map = self.sessions.lock().unwrap(); + if let Some(s) = map.get(key) { + return Ok(s.clone()); + } + let session = Session::spawn( + SessionOptions { + session_id: env.session_id.clone(), + translator_session_id, + source: env.source.clone(), + plugin_version: env.plugin_version.clone(), + replay: Some(replay), + config, + correlation_key: key.correlation_key(), + route: route.clone(), + correlation: self.correlation.clone(), + data_dir: self.data_dir.clone(), + journal, + correlation_changed: self.correlation_changed.clone(), + }, + self.translators.clone(), + self.sink_factory.clone(), + ); + map.insert(key.clone(), session.clone()); + Ok(session) + } + + /// Drop every trace of one delivery pipeline. A session that goes quiet + /// must not pin its translator state, sink handles, credential lease, or + /// journal file for the rest of the daemon's life; deterministic span ids + /// mean a late event simply rebuilds it from the journal. + pub(super) async fn retire_session(&self, key: &DeliveryKey) { + self.observe_claude_transcripts(Some(key), false).await; + self.ingress_barrier().await; + let lock = self.dispatch_lock(&key.source, &key.session_id); + let _guard = lock.lock().await; + + let session = { self.sessions.lock().unwrap().remove(key) }; + let Some(session) = session else { + return; + }; + session.shutdown().await; + self.correlation.retire_session(&key.correlation_key()); + + self.session_auth.lock().await.remove(key); + self.auth_errors.lock().unwrap().remove(key); + self.route_aliases + .lock() + .unwrap() + .retain(|alias, target| alias != key && target != key); + self.managed_run_sessions.lock().unwrap().retain(|_, keys| { + keys.remove(key); + !keys.is_empty() + }); + + // Several delivery routes can share one source session, so release + // its journal writer and lock only after its final route retires. + let last = !self + .sessions + .lock() + .unwrap() + .keys() + .any(|other| other.source == key.source && other.session_id == key.session_id); + if last { + let storage_key = crate::ids::session_namespace(&key.source, &key.session_id); + // Capture can proceed while the actor flushes above. Only take its + // lock for the brief writer-map cleanup after daemon work ends. + let capture_lock = self.session_lock(&key.source, &key.session_id); + let _capture_guard = capture_lock.lock().await; + self.claude_transcripts + .lock() + .unwrap() + .retain(|candidate, _| { + candidate.source != key.source || candidate.session_id != key.session_id + }); + self.journals.lock().unwrap().remove(&storage_key); + { + let mut locks = self.session_locks.lock().unwrap(); + // A waiting capture or observer must keep using this same + // lock, rather than race a newly allocated mirror writer. + if Arc::strong_count(&capture_lock) == 2 { + locks.remove(&storage_key); + } + } + self.dispatch_locks.lock().unwrap().remove(&storage_key); + self.ingress_dispatched + .lock() + .unwrap() + .retain(|candidate, _| { + candidate.source != key.source || candidate.session_id != key.session_id + }); + } + tracing::info!(session_id = %key.session_id, "session retired"); + } + + /// Delivery pipelines with no traffic for `idle_timeout` and nothing left + /// queued. + pub(super) fn idle_sessions(&self, idle_timeout: Duration) -> Vec { + self.sessions + .lock() + .unwrap() + .iter() + .filter(|(_, session)| { + session.idle_for() >= idle_timeout + && session.counters.queued.load(Ordering::Relaxed) == 0 + }) + .filter(|(key, _)| !self.correlation.has_active_tools(&key.correlation_key())) + .map(|(key, _)| key.clone()) + .collect() + } + + pub(super) async fn journal_writer_for( + &self, + source: &str, + session_id: &str, + ) -> anyhow::Result>> { + let storage_key = crate::ids::session_namespace(source, session_id); + if let Some(writer) = self.journals.lock().unwrap().get(&storage_key).cloned() { + return Ok(writer); + } + let path = journal::ensure_source_journal(&self.data_dir, source, session_id).await?; + let writer = Arc::new(tokio::sync::Mutex::new( + JournalWriter::open_path(&path).await?, + )); + Ok(self + .journals + .lock() + .unwrap() + .entry(storage_key) + .or_insert_with(|| writer.clone()) + .clone()) + } + + pub(super) async fn append_to_journal(&self, env: &mut Envelope) -> anyhow::Result<(u64, u64)> { + let writer = self + .journal_writer_for(&env.source, &env.session_id) + .await?; + let mut writer = writer.lock().await; + let before = writer.position(); + let through = writer.append(env).await?; + Ok((before, through)) + } + + pub(super) fn session_lock( + &self, + source: &str, + session_id: &str, + ) -> Arc> { + let storage_key = crate::ids::session_namespace(source, session_id); + self.session_locks + .lock() + .unwrap() + .entry(storage_key) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + + pub(super) fn dispatch_lock( + &self, + source: &str, + session_id: &str, + ) -> Arc> { + let storage_key = crate::ids::session_namespace(source, session_id); + self.dispatch_locks + .lock() + .unwrap() + .entry(storage_key) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } +}