From 3e929a7ae72c435b751f0443b45616892afb08b9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 6 Oct 2026 20:26:46 +0800 Subject: [PATCH] refactor(server): split server.rs into concern modules server.rs was about 2,700 lines of daemon state plus every subsystem that operates on it. The Daemon struct, options, startup and the serve entry points stay in server/mod.rs. Each subsystem moves to its own module, contributing its own `impl Daemon` block: - auth: route auth resolution, refresh before flush, diagnosis - sessions: delivery keys, session actors, journal writers, locks - ingress: the capture queue, its worker and overflow recovery - correlation_store: automatic parent correlation and its persisted state - recovery: startup replay of unfinished journals - rpc: listener, request handling, status - maintenance: GC, idle retirement, the idle watchdog, shutdown drain - managed_run: managed-run session tracking and flushing - claude_transcripts: Claude Code's transcript poller, now isolated from the generic daemon code Code moves verbatim apart from `pub(super)` visibility. Submodules share the parent's imports through `use super::*`. Co-Authored-By: Claude Opus 5.5 (1M context) --- bt-daemon/src/server.rs | 2698 -------------------- bt-daemon/src/server/auth.rs | 209 ++ bt-daemon/src/server/claude_transcripts.rs | 148 ++ bt-daemon/src/server/correlation_store.rs | 611 +++++ bt-daemon/src/server/ingress.rs | 342 +++ bt-daemon/src/server/maintenance.rs | 150 ++ bt-daemon/src/server/managed_run.rs | 189 ++ bt-daemon/src/server/mod.rs | 344 +++ bt-daemon/src/server/recovery.rs | 133 + bt-daemon/src/server/rpc.rs | 425 +++ bt-daemon/src/server/sessions.rs | 253 ++ 11 files changed, 2804 insertions(+), 2698 deletions(-) delete mode 100644 bt-daemon/src/server.rs create mode 100644 bt-daemon/src/server/auth.rs create mode 100644 bt-daemon/src/server/claude_transcripts.rs create mode 100644 bt-daemon/src/server/correlation_store.rs create mode 100644 bt-daemon/src/server/ingress.rs create mode 100644 bt-daemon/src/server/maintenance.rs create mode 100644 bt-daemon/src/server/managed_run.rs create mode 100644 bt-daemon/src/server/mod.rs create mode 100644 bt-daemon/src/server/recovery.rs create mode 100644 bt-daemon/src/server/rpc.rs create mode 100644 bt-daemon/src/server/sessions.rs 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() + } +}