diff --git a/bt-daemon/README.md b/bt-daemon/README.md index faa5411..7dce617 100644 --- a/bt-daemon/README.md +++ b/bt-daemon/README.md @@ -141,12 +141,15 @@ different profiles, organizations, projects, experiments, or parent spans. Phases 0–5 are implemented: protocol, daemon lifecycle, Braintrust sink, Codex and Claude translators, `bt daemon` integration, and thin hook shims for -both shipped plugins. Restart recovery replays the redacted journal with -deterministic span ids, so resubmitted rows merge into the same spans instead -of creating duplicates. Claude lifecycle entries reference a daemon-owned -transcript mirror, so recovery does not depend on mutable external paths -without re-recording the transcript on every turn. Explicit turn/session-end -flushes are bounded, and sessions can target project logs or an experiment. +both shipped plugins. Every coding-agent capture request returns after the raw +event is flushed to its journal; authentication, correlation, translation, and +reporting run on daemon-owned workers. Restart recovery replays the redacted +journal with deterministic span ids, so resubmitted rows merge into the same +spans instead of creating duplicates. Claude and Codex lifecycle entries +reference a daemon-owned transcript mirror, so recovery does not depend on +mutable external paths without re-recording the transcript on every turn. +Explicit turn/session-end flushes are bounded, and sessions can target project +logs or an experiment. Memory is bounded end to end, while on-disk records stay complete: the daemon never holds a transcript or a whole journal in memory, mirroring and replay diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index 57acb25..d5e0dc5 100644 --- a/bt-daemon/docs/protocol.md +++ b/bt-daemon/docs/protocol.md @@ -64,15 +64,13 @@ not found, `-32602` invalid params, `-32603` internal); application errors use ### Ordering & delivery -A subprocess-style shim opens a connection, does one `event.log`, and exits. -Per-session ordering is guaranteed because (a) the agent runs hooks in blocking -mode, so it does not fire the next hook until the current one returns, and (b) -`event.log` is a **request** whose success response means *the event has been -appended to that session's ordered queue* (not that it has been delivered to -Braintrust). The shim must await that response before exiting. Long-lived -in-process clients (opencode/pi, later) hold one connection and may send -`event.log` as a **notification** for the hot path, relying on the single -connection for ordering. +A capture adapter sends `event.log` as a **request** and waits only for the raw +event and any transcript high-water reference to be flushed to the daemon's +journal. The daemon then queues correlation, translation, and reporting on its +own workers. Per-session journal appends are serialized, and the daemon keeps +derived processing ordered without making the hook wait for it. Subprocess +shims may exit as soon as they receive the response; long-lived in-process +clients use the same request boundary. ## Methods @@ -105,18 +103,14 @@ The hot path. Params are the **Envelope** (see below). Request result: ```json { "accepted": true } ``` -`accepted: true` means durably recorded: normally journaled and enqueued to the -session's ordered queue, or held in the daemon's private correlation journal -while multiple parent calls remain indistinguishable. -For an event that opens a tool call, it also means the daemon has made that -active-tool marker visible to local child-session correlation. A child hook -that runs immediately after its parent's blocking pre-tool hook can therefore -attach without an intervening flush. -The daemon never fails the caller's turn for a downstream (Braintrust) error; -those are handled asynchronously and surfaced via `status.get`. The queue is -bounded, so a session whose sink has stalled applies backpressure here instead -of accumulating events without limit; the event is already journaled by then, -so this costs latency, never data. +`accepted: true` means the raw capture is durably recorded in the source +journal. It does not mean authentication, correlation, translation, session +queueing, or Braintrust delivery has completed. Those steps run out-of-band; +errors are surfaced via `status.get`. A child that arrives before its parent's +tool marker is translated is held in durable pending-correlation state and +reconciled by the daemon. On restart, uncheckpointed journal entries are queued +again automatically. Explicit status and flush requests act as daemon-worker +barriers, but hook capture never does. ### `session.flush` (request) diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index 3289630..6b4e1da 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -3,11 +3,9 @@ //! are processed strictly in arrival order. Different sessions run //! concurrently. //! -//! Ack semantics: `event.log` is acked once the event is journaled and its -//! first translation batch has updated local correlation state. Tool-start -//! events drain bounded translator continuations before ack so the spawning -//! marker is guaranteed visible. A downstream error never fails the caller's -//! turn. +//! Hook acknowledgement happens before this layer, immediately after the raw +//! event is durably journaled. Translation, correlation, and delivery here are +//! entirely out-of-band from hook execution. use crate::delivery_ledger::LedgerSink; use crate::journal::JournalWriter; @@ -33,8 +31,9 @@ pub struct Counters { } enum SessionMsg { - Event(Box, u64, oneshot::Sender<()>), + Event(Box, u64), Configure(Box, oneshot::Sender<()>), + Barrier(oneshot::Sender<()>), Flush(oneshot::Sender), Finalize(oneshot::Sender), Shutdown(oneshot::Sender<()>), @@ -65,6 +64,7 @@ pub(crate) struct SessionOptions { pub correlation: Arc, pub data_dir: PathBuf, pub journal: Arc>, + pub correlation_changed: Arc, } /// Handle to one live session: its queue plus observable counters/state. @@ -96,6 +96,7 @@ impl Session { correlation, data_dir, journal, + correlation_changed, } = options; let (tx, rx) = mpsc::channel(QUEUE_CAPACITY); let counters = Arc::new(Counters::default()); @@ -119,6 +120,7 @@ impl Session { correlation, data_dir, journal, + correlation_changed, }; tokio::spawn(actor.run(rx)); @@ -132,24 +134,29 @@ impl Session { }) } - /// Enqueue an event after the daemon has journaled it. + /// Queue a journaled event without waiting for translation or delivery. pub async fn enqueue(&self, env: Envelope, journal_through: u64) -> anyhow::Result<()> { self.touch(); self.counters.queued.fetch_add(1, Ordering::Relaxed); - let (reply_tx, reply_rx) = oneshot::channel(); self.tx - .send(SessionMsg::Event(Box::new(env), journal_through, reply_tx)) + .send(SessionMsg::Event(Box::new(env), journal_through)) .await - .map_err(|_| anyhow::anyhow!("session actor is gone"))?; - reply_rx - .await - .map_err(|_| anyhow::anyhow!("session actor dropped event acknowledgement")) + .map_err(|_| anyhow::anyhow!("session actor is gone")) } fn touch(&self) { *self.last_activity.lock().unwrap() = Instant::now(); } + /// Wait until events already accepted by this daemon worker have updated + /// translator and correlation state. Hook capture never calls this. + pub async fn barrier(&self) { + let (reply_tx, reply_rx) = oneshot::channel(); + if self.tx.send(SessionMsg::Barrier(reply_tx)).await.is_ok() { + let _ = reply_rx.await; + } + } + /// How long since this session last saw traffic. Drives idle retirement. pub fn idle_for(&self) -> std::time::Duration { self.last_activity.lock().unwrap().elapsed() @@ -206,19 +213,22 @@ impl Session { } } -/// Claude transcript files are external mutable state. Mirror them into +/// Agent transcript files are external mutable state. Mirror them into /// daemon-owned storage at lifecycle boundaries and journal only a reference, /// so recovery/replay does not depend on a path that Claude may later rewrite /// or delete — and so the transcript is stored once rather than re-copied into /// every event. Fail open: without a reference the translator reads the live /// path exactly as before. pub(crate) async fn hydrate_transcript_reference(data_dir: &std::path::Path, env: &mut Envelope) { - if env.source != "claude-code" - || !matches!( + let should_capture = match env.source.as_str() { + "codex" => true, + "claude-code" => matches!( env.event.as_str(), "UserPromptSubmit" | "Stop" | "StopFailure" | "SubagentStop" | "SessionEnd" - ) - { + ), + _ => false, + }; + if !should_capture { return; } let field = if env.event == "SubagentStop" { @@ -272,6 +282,7 @@ struct SessionActor { correlation: Arc, data_dir: PathBuf, journal: Arc>, + correlation_changed: Arc, } #[derive(Clone, Copy)] @@ -328,13 +339,15 @@ impl SessionActor { // callers waiting on flush don't hang. while let Some(msg) = rx.recv().await { match msg { - SessionMsg::Event(_, _, reply) => { + SessionMsg::Event(_, _) => { self.counters.queued.fetch_sub(1, Ordering::Relaxed); - let _ = reply.send(()); } SessionMsg::Configure(_, r) => { let _ = r.send(()); } + SessionMsg::Barrier(r) => { + let _ = r.send(()); + } SessionMsg::Flush(r) => { let _ = r.send(0); } @@ -381,18 +394,14 @@ impl SessionActor { while let Some(msg) = rx.recv().await { match msg { - SessionMsg::Event(env, journal_through, reply) => { + SessionMsg::Event(env, journal_through) => { let correlation_barrier = is_tool_lifecycle_event(&env.event); - let mut reply = Some(reply); if let Some(cfg) = &env.config { sink.configure(cfg); ctx.config = Some(cfg.clone()); self.refresh_permalink(sink.as_ref()); } let translated = translator.handle(&env, &ctx); - if !correlation_barrier { - let _ = reply.take().expect("event reply").send(()); - } let (correlation_changed, delivered) = self .emit_translator_batches( &mut translator, @@ -412,9 +421,7 @@ impl SessionActor { { self.set_error(error); } - } - if let Some(reply) = reply { - let _ = reply.send(()); + self.correlation_changed.notify_one(); } self.counters.queued.fetch_sub(1, Ordering::Relaxed); if delivered && checkpointable { @@ -429,6 +436,9 @@ impl SessionActor { self.refresh_permalink(sink.as_ref()); let _ = reply.send(()); } + SessionMsg::Barrier(reply) => { + let _ = reply.send(()); + } SessionMsg::Flush(reply) => { checkpointable &= self .checkpoint_and_flush(&mut translator, &mut sink, &ctx) @@ -510,12 +520,16 @@ impl SessionActor { while let Some(ops) = next { if !ops.is_empty() { if mode.observes_correlation() { - correlation_changed |= self.correlation.observe_ops( + let changed = self.correlation.observe_ops( &self.correlation_key, &self.route, ctx.config.as_ref().expect("session config"), &ops, ); + correlation_changed |= changed; + if changed { + self.persist_correlation_if_changed(true).await; + } } match sink.emit(&ops).await { Ok(n) => { diff --git a/bt-daemon/src/journal.rs b/bt-daemon/src/journal.rs index 252af35..be0448d 100644 --- a/bt-daemon/src/journal.rs +++ b/bt-daemon/src/journal.rs @@ -227,6 +227,10 @@ impl JournalWriter { Ok(Self { file, position }) } + pub(crate) fn position(&self) -> u64 { + self.position + } + /// Append one event in redacted form and flush to the OS. Not fsync'd per /// event (that would dominate hook latency); an OS crash can lose the last /// few lines, which replay tolerates. diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 3b73a34..f08d1f0 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -115,12 +115,14 @@ pub struct HookArgs { /// Fail instead of spawning a daemon if none is running. #[arg(long)] pub no_spawn: bool, - /// Flush the session after a turn-ending event. Intended for short-lived - /// CI hosts; SessionEnd is always flushed. + /// Ask the daemon to flush the session after a turn-ending event. The + /// flush is scheduled out-of-band; hook capture still returns immediately + /// after the durable journal write. #[arg(long)] pub flush_on_turn_end: bool, - /// Bound an explicit turn/session-end flush. - #[arg(long, default_value_t = 10_000)] + /// Deprecated compatibility option. Hook capture never waits for daemon + /// translation, reporting, or flushing. + #[arg(long, default_value_t = 10_000, hide = true)] pub flush_timeout_ms: u64, /// JSON object merged into root-span metadata. Deliberately not read from /// the environment: a hook fires automatically on every event, so its @@ -293,7 +295,7 @@ pub async fn run_serve(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<() /// Capture one hook event from `stdin` and forward it to the daemon. /// /// `route` contains only non-secret profile and destination selection. -/// Returns `Ok` once the daemon has acked (journaled + enqueued). Callers that +/// Returns `Ok` once the daemon has durably journaled the event. Callers that /// must never fail the agent's turn should treat any `Err` as non-fatal and /// exit 0. pub async fn run_hook( @@ -350,14 +352,6 @@ pub async fn run_hook( let socket = paths::socket_path(args.socket.as_deref()); forward_envelope(&env, &socket, &host, args.no_spawn).await?; - let should_flush = env.event == "SessionEnd" - || (matches!( - env.route.as_ref().map(|r| r.flush_mode), - Some(wire::FlushMode::FlushOnTurnEnd) - ) && matches!(env.event.as_str(), "Stop" | "SubagentStop")); - if should_flush { - flush_session(&env.session_id, &socket, args.flush_timeout_ms).await?; - } Ok(()) } diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index c0ee7ac..c957bfc 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -24,7 +24,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt, BufReader}; -use tokio::sync::Notify; +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. @@ -79,13 +79,63 @@ struct PendingSession { #[serde(default, skip_serializing_if = "Option::is_none")] linked_route: Option, #[serde(default)] - events: Vec, + events: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] candidate_span_ids: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] evidence: Vec, } +#[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)] @@ -120,7 +170,11 @@ pub struct Daemon { 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>>>, managed_run_sessions: Mutex>>, auth_errors: Mutex>, @@ -128,7 +182,12 @@ pub struct Daemon { correlation: Arc, automatic_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, shutting_down: AtomicBool, @@ -137,7 +196,8 @@ pub struct Daemon { impl Daemon { fn new(opts: ServeOptions, data_dir: PathBuf) -> Arc { - Arc::new(Daemon { + let (ingress_tx, ingress_rx) = mpsc::channel(INGRESS_QUEUE_CAPACITY); + let daemon = Arc::new(Daemon { version: opts.version, data_dir, translators: opts.translators, @@ -146,6 +206,7 @@ impl Daemon { 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()), managed_run_sessions: Mutex::new(HashMap::new()), auth_errors: Mutex::new(HashMap::new()), @@ -153,12 +214,20 @@ impl Daemon { correlation: Arc::new(crate::correlation::CorrelationRegistry::default()), automatic_links: Mutex::new(HashMap::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()), shutting_down: AtomicBool::new(false), shutdown: Notify::new(), - }) + }); + spawn_ingress_worker(daemon.clone(), ingress_rx); + spawn_pending_reconciler(daemon.clone()); + daemon } async fn configure_event(&self, env: &mut Envelope) -> anyhow::Result { @@ -309,7 +378,12 @@ impl Daemon { .clone() } - async fn session_for(&self, env: &Envelope, key: &DeliveryKey) -> anyhow::Result> { + 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) { @@ -339,15 +413,14 @@ impl Daemon { } else { crate::ids::session_namespace(&env.source, &env.session_id) }; - let through = journal::JournalReader::recorded_len(&journal_path).await; let replay = ReplayPlan { acknowledged_through: journal::JournalReader::acknowledged_through( &journal_path, - through, + replay_through, route, ) .await, - through, + through: replay_through, journal_path, }; let journal = self @@ -370,6 +443,7 @@ impl Daemon { correlation: self.correlation.clone(), data_dir: self.data_dir.clone(), journal, + correlation_changed: self.correlation_changed.clone(), }, self.translators.clone(), self.sink_factory.clone(), @@ -383,7 +457,7 @@ impl Daemon { /// 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) { - let lock = self.session_lock(&key.source, &key.session_id); + let lock = self.dispatch_lock(&key.source, &key.session_id); let _guard = lock.lock().await; let session = { self.sessions.lock().unwrap().remove(key) }; @@ -414,8 +488,19 @@ impl Daemon { .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.journals.lock().unwrap().remove(&storage_key); self.session_locks.lock().unwrap().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"); } @@ -458,14 +543,15 @@ impl Daemon { .clone()) } - async fn append_to_journal(&self, env: &mut Envelope) -> anyhow::Result { + async fn append_to_journal(&self, env: &mut Envelope) -> anyhow::Result<(u64, u64)> { hydrate_transcript_reference(&self.data_dir, env).await; - self.journal_writer_for(&env.source, &env.session_id) - .await? - .lock() - .await - .append(env) - .await + 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> { @@ -478,6 +564,16 @@ impl Daemon { .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 total_queued(&self) -> u64 { self.sessions .lock() @@ -487,6 +583,58 @@ impl Daemon { .sum() } + async fn capture_event(&self, mut env: Envelope) -> Result<(), String> { + 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; + let (replay_through, journal_through) = self + .append_to_journal(&mut env) + .await + .map_err(|error| format!("journal failed: {error}"))?; + 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(()) + } + + 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() @@ -605,6 +753,31 @@ impl Daemon { result } + /// Flush every live delivery route for one source session. 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 flush_source_session(&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 (flushed, pending) = session.flush(timeout).await; + if !flushed { + tracing::warn!( + source, + session_id, + pending, + "out-of-band turn-end flush did not complete" + ); + } + } + } + fn trigger_shutdown(&self) { self.shutting_down.store(true, Ordering::SeqCst); self.shutdown.notify_waiters(); @@ -626,6 +799,189 @@ fn now_ms() -> 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 schedule_flush = event.env.event == "SessionEnd" + || (matches!( + event.env.route.as_ref().map(|route| route.flush_mode), + Some(crate::wire::FlushMode::FlushOnTurnEnd) + ) && matches!(event.env.event.as_str(), "Stop" | "SubagentStop")); + let flush_source = event.env.source.clone(); + let flush_session_id = event.env.session_id.clone(); + 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. + if is_session_start(&event.env.event) { + 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"); + return; + } + } + if schedule_flush { + let daemon = daemon.clone(); + tokio::spawn(async move { + daemon + .flush_source_session(&flush_source, &flush_session_id, Duration::from_secs(10)) + .await; + }); + } +} + +/// 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 { + loop { + tokio::select! { + _ = daemon.shutdown.notified() => return, + _ = daemon.correlation_changed.notified() => { + 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<()> { @@ -652,6 +1008,8 @@ pub async fn run(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { 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; let idle_timeout = Duration::from_secs(args.idle_timeout_secs); spawn_idle_watchdog(daemon.clone(), idle_timeout); spawn_session_reaper( @@ -750,7 +1108,7 @@ async fn serve_connection(daemon: Arc, stream: ServerStream) -> anyhow:: match serde_json::from_value::(params) { Ok(mut env) => { attach_process_capture(&mut env, client.as_ref()); - let _ = accept_event(&daemon, env).await; + let _ = daemon.capture_event(env).await; } Err(error) => tracing::warn!( method = %note.method, @@ -795,14 +1153,12 @@ fn attach_process_capture(env: &mut Envelope, client: Option<&crate::wire::Clien } } -async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), String> { - let canonical_source = daemon - .translators - .canonical_source(&env.source) - .ok_or_else(|| format!("unsupported coding-agent source {:?}", env.source))? - .to_string(); - env.source = canonical_source; - daemon.touch(); +async fn accept_event(daemon: &Arc, event: PendingEvent) -> Result<(), String> { + let PendingEvent { + mut env, + replay_through, + journal_through, + } = event; let requested_link_key = automatic_link_key(&env); let correlation_lock = daemon.correlation_lock(&requested_link_key); let _correlation_guard = correlation_lock.lock().await; @@ -832,7 +1188,7 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str if let Some(mut state) = state { if let Some(route) = state.linked_route.clone() { for mut pending in std::mem::take(&mut state.events) { - pending.route = Some(route.clone()); + pending.env.route = Some(route.clone()); accept_resolved_event(daemon, pending).await?; } write_correlation_state(&daemon.data_dir, &requested_link_key, &state).await?; @@ -843,14 +1199,22 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str .insert(requested_link_key, route.clone()); env.route = Some(route); drop(_correlation_guard); - return accept_resolved_and_retry_pending(daemon, env).await; + return accept_resolved_and_retry_pending( + daemon, + PendingEvent { + env, + replay_through, + journal_through, + }, + ) + .await; } let capture = env.capture.as_ref().or_else(|| { state .events .first() - .and_then(|event| event.capture.as_ref()) + .and_then(|event| event.env.capture.as_ref()) }); state.evidence.push(correlation_evidence(&env).await); let evidence = Value::Array(state.evidence.clone()); @@ -862,10 +1226,14 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str ) { crate::correlation::Resolution::Parent(parent) => { state.linked_route = Some(parent.route.clone()); - state.events.push(env); + state.events.push(PendingEvent { + env, + replay_through, + journal_through, + }); write_correlation_state(&daemon.data_dir, &requested_link_key, &state).await?; for mut event in std::mem::take(&mut state.events) { - event.route = Some(parent.route.clone()); + event.env.route = Some(parent.route.clone()); accept_resolved_event(daemon, event).await?; } write_correlation_state(&daemon.data_dir, &requested_link_key, &state).await?; @@ -878,13 +1246,18 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str } crate::correlation::Resolution::Ambiguous(_) | crate::correlation::Resolution::Standalone => { - state.events.push(env); + state.events.push(PendingEvent { + env, + replay_through, + journal_through, + }); write_correlation_state(&daemon.data_dir, &requested_link_key, &state).await?; daemon .pending_sessions .lock() .unwrap() .insert(requested_link_key, state); + daemon.correlation_changed.notify_one(); return Ok(()); } } @@ -915,7 +1288,11 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str let evidence = correlation_evidence(&env).await; let state = PendingSession { linked_route: None, - events: vec![env], + events: vec![PendingEvent { + env, + replay_through, + journal_through, + }], candidate_span_ids, evidence: vec![evidence], }; @@ -925,6 +1302,7 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str .lock() .unwrap() .insert(requested_link_key, state); + daemon.correlation_changed.notify_one(); return Ok(()); } crate::correlation::Resolution::Standalone => {} @@ -932,18 +1310,27 @@ async fn accept_event(daemon: &Arc, mut env: Envelope) -> Result<(), Str } drop(_correlation_guard); - accept_resolved_and_retry_pending(daemon, env).await + accept_resolved_and_retry_pending( + daemon, + PendingEvent { + env, + replay_through, + journal_through, + }, + ) + .await } async fn accept_resolved_and_retry_pending( daemon: &Arc, - env: Envelope, + event: PendingEvent, ) -> Result<(), String> { - accept_resolved_event(daemon, env).await?; + accept_resolved_event(daemon, event).await?; retry_pending_sessions(daemon).await } async fn retry_pending_sessions(daemon: &Arc) -> Result<(), String> { + let _reconcile_guard = daemon.pending_reconcile_lock.lock().await; let keys: Vec = daemon .pending_sessions .lock() @@ -961,13 +1348,16 @@ async fn retry_pending_sessions(daemon: &Arc) -> Result<(), String> { daemon.pending_sessions.lock().unwrap().insert(key, state); continue; } - let capture = state.events.iter().find_map(|event| event.capture.as_ref()); + let capture = state + .events + .iter() + .find_map(|event| event.env.capture.as_ref()); let evidence = Value::Array(state.evidence.clone()); match daemon.correlation.resolve_pending( state .events .first() - .map(|event| event.source.as_str()) + .map(|event| event.env.source.as_str()) .unwrap_or(""), capture, &evidence, @@ -977,7 +1367,7 @@ async fn retry_pending_sessions(daemon: &Arc) -> Result<(), String> { state.linked_route = Some(parent.route.clone()); write_correlation_state(&daemon.data_dir, &key, &state).await?; for mut event in std::mem::take(&mut state.events) { - event.route = Some(parent.route.clone()); + event.env.route = Some(parent.route.clone()); accept_resolved_event(daemon, event).await?; } write_correlation_state(&daemon.data_dir, &key, &state).await?; @@ -1181,14 +1571,151 @@ async fn restore_active_parent_snapshots( } } -async fn accept_resolved_event(daemon: &Arc, mut env: Envelope) -> Result<(), String> { +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 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(); + 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); + } + } + } + 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; 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.session_lock(&source, &session_id); + let session_lock = daemon.dispatch_lock(&source, &session_id); let _session_guard = session_lock.lock().await; let result = async { @@ -1208,16 +1735,12 @@ async fn accept_resolved_event(daemon: &Arc, mut env: Envelope) -> Resul } } let session = daemon - .session_for(&env, &delivery_key) + .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.capture.as_ref()); - let journal_through = daemon - .append_to_journal(&mut env) - .await - .map_err(|error| format!("journal failed: {error}"))?; session .enqueue(env, journal_through) .await @@ -1253,6 +1776,29 @@ fn automatic_link_key(env: &Envelope) -> String { 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, @@ -1303,7 +1849,7 @@ async fn handle_request( method::EVENT_LOG => { let mut env = parse!(Envelope); attach_process_capture(&mut env, client.as_ref()); - match accept_event(daemon, env).await { + match daemon.capture_event(env).await { Ok(()) => Response::ok( id, serde_json::to_value(EventLogResult { accepted: true }).unwrap(), @@ -1313,6 +1859,9 @@ async fn handle_request( } 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; let delivery_keys: Vec<_> = daemon .sessions .lock() @@ -1354,11 +1903,13 @@ async fn handle_request( } method::MANAGED_RUN_FLUSH => { let params = parse!(ManagedRunFlushParams); + daemon.settle_ingress().await; 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::DAEMON_SHUTDOWN => { @@ -1522,6 +2073,7 @@ fn spawn_idle_watchdog(daemon: Arc, idle_timeout: Duration) { } async fn drain_all(daemon: &Arc) { + daemon.settle_ingress().await; let sessions: Vec> = daemon.sessions.lock().unwrap().values().cloned().collect(); for s in sessions { s.shutdown().await; diff --git a/bt-daemon/src/translate/codex.rs b/bt-daemon/src/translate/codex.rs index f6439fa..844491a 100644 --- a/bt-daemon/src/translate/codex.rs +++ b/bt-daemon/src/translate/codex.rs @@ -144,6 +144,7 @@ enum PendingWork { path: String, hook_ts: i64, through_ms: Option, + through_bytes: Option, after: DeferredHook, }, CatchUp { @@ -200,19 +201,14 @@ impl AgentTranslator for CodexTranslator { self.session_source = str_field(payload, "source"); self.permission_mode = str_field(payload, "permission_mode"); } - "SubagentStart" => self.handle_subagent_start(payload), + "SubagentStart" => self.handle_subagent_start(event), "PreCompact" | "PostCompact" => self.record_compaction_trigger(payload, &mut ops), _ => {} } // --- pick the scope and catch up its transcript --- let agent_id = str_field(payload, "agent_id"); - let path = if event.event == "SubagentStop" { - str_field(payload, "agent_transcript_path") - } else { - str_field(payload, "transcript_path") - .or_else(|| str_field(payload, "agent_transcript_path")) - }; + let path = effective_transcript_path(event); if let Some(path) = path { if agent_id.is_none() { @@ -220,14 +216,24 @@ impl AgentTranslator for CodexTranslator { self.ensure_main_scope(&path); } let import_through_ms = payload.get("_bt_import_through_ms").and_then(Value::as_i64); + let through_bytes = payload + .pointer("/_bt_transcript_mirror/through") + .and_then(Value::as_u64); let after = self.deferred_hook(event, agent_id.is_none()); - if self.catch_up_chunk(&path, event.ts_ms, import_through_ms, &mut ops) { + if self.catch_up_chunk( + &path, + event.ts_ms, + import_through_ms, + through_bytes, + &mut ops, + ) { self.finish_deferred_hook(after, &mut ops); } else { self.pending = Some(PendingWork::Hook { path, hook_ts: event.ts_ms, through_ms: import_through_ms, + through_bytes, after, }); } @@ -248,15 +254,17 @@ impl AgentTranslator for CodexTranslator { path, hook_ts, through_ms, + through_bytes, after, } => { - if self.catch_up_chunk(&path, hook_ts, through_ms, &mut ops) { + if self.catch_up_chunk(&path, hook_ts, through_ms, through_bytes, &mut ops) { self.finish_deferred_hook(after, &mut ops); } else { self.pending = Some(PendingWork::Hook { path, hook_ts, through_ms, + through_bytes, after, }); } @@ -268,7 +276,7 @@ impl AgentTranslator for CodexTranslator { } => { while next_path < paths.len() { let path = &paths[next_path]; - if self.catch_up_chunk(path, 0, None, &mut ops) { + if self.catch_up_chunk(path, 0, None, None, &mut ops) { if finalize { if let Some(mut scope) = self.scopes.remove(path) { self.close_dangling(&mut scope, None, &mut ops); @@ -373,10 +381,11 @@ impl CodexTranslator { } } - fn handle_subagent_start(&mut self, payload: &Value) { + fn handle_subagent_start(&mut self, event: &Envelope) { + let payload = &event.payload; let (Some(agent_id), Some(path)) = ( str_field(payload, "agent_id"), - str_field(payload, "transcript_path"), + effective_transcript_path(event), ) else { return; }; @@ -401,7 +410,7 @@ impl CodexTranslator { // the spawn_agent transcript record that establishes call -> turn. "PostToolUse" => DeferredHook::PostToolUse(event.payload.clone()), "SubagentStop" => DeferredHook::SubagentStop { - path: str_field(&event.payload, "agent_transcript_path"), + path: effective_transcript_path(event), ts: event.ts_ms, }, // Codex writes task_complete slightly after the Stop hook in real @@ -450,6 +459,7 @@ impl CodexTranslator { path: &str, hook_ts: i64, through_ms: Option, + through_bytes: Option, ops: &mut Vec, ) -> bool { let Some(mut scope) = self.scopes.remove(path) else { @@ -459,6 +469,7 @@ impl CodexTranslator { &scope.path, &mut scope.offset, through_ms, + through_bytes, CATCH_UP_BYTE_BUDGET, ); for line in read.lines { @@ -1271,6 +1282,22 @@ impl CodexTranslator { } } +fn effective_transcript_path(event: &Envelope) -> Option { + event + .payload + .pointer("/_bt_transcript_mirror/mirror") + .and_then(Value::as_str) + .map(str::to_owned) + .or_else(|| { + if event.event == "SubagentStop" { + str_field(&event.payload, "agent_transcript_path") + } else { + str_field(&event.payload, "transcript_path") + .or_else(|| str_field(&event.payload, "agent_transcript_path")) + } + }) +} + impl Scope { fn new(path: &str, kind: ScopeKind, turn_parent_span_id: String) -> Self { Scope { @@ -1628,6 +1655,7 @@ fn read_new_lines( path: &str, offset: &mut u64, through_ms: Option, + through_bytes: Option, byte_budget: usize, ) -> ReadLines { use std::io::{BufRead, BufReader, Seek, SeekFrom}; @@ -1652,6 +1680,9 @@ fn read_new_lines( let mut lines = Vec::new(); let mut consumed = 0usize; loop { + if through_bytes.is_some_and(|limit| *offset >= limit) { + break; + } if consumed >= byte_budget && !lines.is_empty() { return ReadLines { lines, @@ -1665,6 +1696,9 @@ fn read_new_lines( if bytes == 0 || !line.ends_with('\n') { break; } + if through_bytes.is_some_and(|limit| offset.saturating_add(bytes as u64) > limit) { + break; + } let trimmed = line.trim_end_matches(['\r', '\n']); if let Some(limit) = through_ms { if serde_json::from_str::(trimmed) diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 896cd88..b972c7b 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -958,6 +958,14 @@ fn codex_subagent_nests_under_spawning_turn() { let sub_t = tmp.path().join("sub.jsonl"); let main_p = main_t.to_str().unwrap(); let sub_p = sub_t.to_str().unwrap(); + let original_sub_p = "/rollouts/subagent-a1.jsonl"; + let mirrored_subagent = || { + json!({ + "path": original_sub_p, + "mirror": sub_p, + "through": std::fs::metadata(&sub_t).map(|meta| meta.len()).unwrap_or(0), + }) + }; // Main session opens a turn and runs a spawn_agent tool. for v in [ @@ -993,7 +1001,12 @@ fn codex_subagent_nests_under_spawning_turn() { "s", "SubagentStart", main_p, - json!({ "agent_id": "a1", "transcript_path": sub_p, "agent_type": "reviewer" }), + json!({ + "agent_id": "a1", + "transcript_path": original_sub_p, + "agent_type": "reviewer", + "_bt_transcript_mirror": mirrored_subagent(), + }), ), &ctx, ) @@ -1020,7 +1033,11 @@ fn codex_subagent_nests_under_spawning_turn() { "s", "PostToolUse", main_p, - json!({ "agent_id": "a1", "transcript_path": sub_p }), + json!({ + "agent_id": "a1", + "transcript_path": original_sub_p, + "_bt_transcript_mirror": mirrored_subagent(), + }), ), &ctx, ) @@ -1032,7 +1049,11 @@ fn codex_subagent_nests_under_spawning_turn() { "s", "SubagentStop", main_p, - json!({ "agent_id": "a1", "agent_transcript_path": sub_p }), + json!({ + "agent_id": "a1", + "agent_transcript_path": original_sub_p, + "_bt_transcript_mirror": mirrored_subagent(), + }), ), &ctx, ) diff --git a/bt-daemon/tests/distributed_tracing.rs b/bt-daemon/tests/distributed_tracing.rs index db0aa36..965a1df 100644 --- a/bt-daemon/tests/distributed_tracing.rs +++ b/bt-daemon/tests/distributed_tracing.rs @@ -156,9 +156,9 @@ async fn every_instrumented_parent_child_agent_pair_shares_one_trace() { ) .await; - // The child starts immediately after the blocking start-tool hook. - // No parent flush is inserted here: forwarding the start event must - // not return until its correlation marker is visible. + // The child starts immediately after the journal-only start-tool + // hook. Daemon ingress orders the parent's correlation update + // before resolving the child; hook execution does not wait for it. forward_all( &mut fixtures.start_turn( child_kind, @@ -273,6 +273,11 @@ async fn every_agent_pair_links_when_the_daemon_restarts_between_spawn_and_child ) .await; + // Hook capture guarantees only durable journaling. This explicit + // daemon barrier makes the derived snapshot observable to the + // restart assertion below. + flush(&parent_session, &socket).await; + let parent_snapshot_dir = data_dir.join("correlation").join("parents"); let mut snapshots = tokio::fs::read_dir(&parent_snapshot_dir).await.unwrap(); let snapshot = snapshots @@ -989,6 +994,7 @@ async fn ambiguous_child_evidence_survives_daemon_restart() { 1_700_500_000_020, ); forward(child_start.remove(0), &socket, &host).await; + flush("pending-child", &socket).await; assert!(tokio::fs::read_dir(data_dir.join("correlation")) .await .unwrap() @@ -1622,6 +1628,7 @@ async fn concurrent_agent_sessions_in_one_process_choose_their_own_tools() { } for (_, _, parent_session, _, _, _, _, _, _) in &cases { + flush(parent_session, &socket).await; let parent = recording.session(parent_session); assert!( parent diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 1812dda..3b03eeb 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -43,6 +43,69 @@ struct TrackingSink { flushes: Arc>>, } +struct SlowSink; + +#[async_trait] +impl Sink for SlowSink { + fn configure(&mut self, _config: &SessionConfig) {} + async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { + tokio::time::sleep(Duration::from_millis(500)).await; + Ok(ops.len() as u64) + } + async fn flush(&mut self) -> anyhow::Result<()> { + Ok(()) + } +} + +struct SlowSinkFactory; + +impl SinkFactory for SlowSinkFactory { + fn create(&self, _: &str, _: &str, _: Option<&str>) -> anyhow::Result> { + Ok(Box::new(SlowSink)) + } +} + +struct GateSinkFactory { + blocked: Arc, + gate: Arc, + emitted: Arc, +} + +struct GateSink { + blocked: Arc, + gate: Arc, + emitted: Arc, +} + +#[async_trait] +impl Sink for GateSink { + async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { + if self + .blocked + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + self.gate.notified().await; + } + self.emitted + .fetch_add(ops.len() as u64, std::sync::atomic::Ordering::Relaxed); + Ok(ops.len() as u64) + } + + async fn flush(&mut self) -> anyhow::Result<()> { + Ok(()) + } +} + +impl SinkFactory for GateSinkFactory { + fn create(&self, _: &str, _: &str, _: Option<&str>) -> anyhow::Result> { + Ok(Box::new(GateSink { + blocked: self.blocked.clone(), + gate: self.gate.clone(), + emitted: self.emitted.clone(), + })) + } +} + #[async_trait] impl Sink for TrackingSink { fn configure(&mut self, _config: &SessionConfig) {} @@ -413,6 +476,72 @@ async fn start_tracking_daemon( (socket, handle, flushes, tmp) } +async fn start_slow_daemon() -> (PathBuf, tokio::task::JoinHandle<()>, tempfile::TempDir) { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let opts = ServeOptions { + version: "test".into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(SlowSinkFactory), + auth_provider: Some(Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + })), + }; + let args = ServeArgs { + socket: Some(socket.clone()), + data_dir: Some(data_dir), + idle_timeout_secs: 0, + session_idle_timeout_secs: 0, + }; + let handle = tokio::spawn(async move { + let _ = run_serve(args, opts).await; + }); + wait_for(&socket).await; + (socket, handle, tmp) +} + +async fn start_gated_daemon( + gate: Arc, + emitted: Arc, +) -> ( + PathBuf, + PathBuf, + tokio::task::JoinHandle<()>, + tempfile::TempDir, +) { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let opts = ServeOptions { + version: "test".into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(GateSinkFactory { + blocked: Arc::new(std::sync::atomic::AtomicBool::new(true)), + gate, + emitted, + }), + auth_provider: Some(Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + })), + }; + let args = ServeArgs { + socket: Some(socket.clone()), + data_dir: Some(data_dir.clone()), + idle_timeout_secs: 0, + session_idle_timeout_secs: 0, + }; + let handle = tokio::spawn(async move { + let _ = run_serve(args, opts).await; + }); + wait_for(&socket).await; + (data_dir, socket, handle, tmp) +} + async fn start_daemon_at(data_dir: PathBuf, socket: PathBuf) -> tokio::task::JoinHandle<()> { let args = ServeArgs { socket: Some(socket.clone()), @@ -551,6 +680,7 @@ async fn expiring_profile_lease_is_refreshed_for_the_pinned_profile() { ) .await .unwrap(); + flush_session("refresh", &socket, 5000).await.unwrap(); let calls = provider.calls.lock().unwrap().clone(); assert_eq!(calls.len(), 2); @@ -651,6 +781,60 @@ async fn one_session_reports_to_multiple_routes_and_orgs_concurrently() { handle.await.unwrap(); } +#[tokio::test] +async fn recovery_checkpoints_are_scoped_to_their_delivery_route() { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let journal = source_journal_path(&data_dir, "debug", "route-recovery"); + std::fs::create_dir_all(journal.parent().unwrap()).unwrap(); + + let personal = routed_envelope("route-recovery", "personal", "personal-org", "SessionStart"); + let work = routed_envelope("route-recovery", "work", "work-org", "Stop"); + let personal_line = serde_json::to_vec(&personal.redacted()).unwrap(); + let work_line = serde_json::to_vec(&work.redacted()).unwrap(); + let work_through = (personal_line.len() + 1 + work_line.len() + 1) as u64; + let checkpoint = serde_json::to_vec(&serde_json::json!({ + "_bt_record_type": "delivery_checkpoint", + "route": work.route.as_ref().unwrap(), + "through": work_through, + })) + .unwrap(); + let mut contents = Vec::new(); + for line in [&personal_line, &work_line, &checkpoint] { + contents.extend_from_slice(line); + contents.push(b'\n'); + } + std::fs::write(&journal, contents).unwrap(); + + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + }); + let recording = Arc::new(RouteRecordingSinkFactory::default()); + let handle = start_daemon_at_with(data_dir, socket.clone(), provider, recording.clone()).await; + let flushed = flush_session("route-recovery", &socket, 5000) + .await + .unwrap(); + assert!(flushed.flushed, "flush did not complete: {flushed:?}"); + + let sinks = recording.sinks.lock().unwrap().clone(); + assert_eq!(sinks.len(), 1, "only the unacknowledged route is recovered"); + assert_eq!( + sinks[0].org.lock().unwrap().as_deref(), + Some("personal-org"), + "a later checkpoint for work must not suppress personal recovery" + ); + assert!( + sinks[0].emitted.load(std::sync::atomic::Ordering::Relaxed) > 0, + "the unacknowledged route must reach its sink" + ); + + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn a_route_backfills_observations_captured_by_another_route_once() { let provider = Arc::new(TestAuthProvider { @@ -847,15 +1031,14 @@ async fn auth_resolution_failure_is_reported_without_exposing_credentials() { let (data_dir, socket, handle, _tmp) = start_routed_daemon(provider.clone()).await; let host = dummy_host(); - let error = forward_envelope( + forward_envelope( &routed_envelope("login-needed", "missing", "missing-org", "SessionStart"), &socket, &host, false, ) .await - .unwrap_err(); - assert!(error.to_string().contains("bt login")); + .unwrap(); let status = run_status(StatusArgs { socket: Some(socket.clone()), session_id: Some("login-needed".into()), @@ -866,7 +1049,10 @@ async fn auth_resolution_failure_is_reported_without_exposing_credentials() { let status_error = status.sessions[0].last_error.as_deref().unwrap(); assert!(status_error.contains("select a profile explicitly")); assert!(!status_error.contains("secret-")); - assert!(!source_journal_path(&data_dir, "debug", "login-needed").exists()); + assert!( + source_journal_path(&data_dir, "debug", "login-needed").exists(), + "capture must remain durable even when background auth fails" + ); assert_eq!(provider.calls.lock().unwrap().len(), 1); shutdown(&socket).await; @@ -1068,6 +1254,97 @@ async fn distinct_sessions_are_isolated() { handle.abort(); } +#[tokio::test] +async fn hook_capture_stops_at_the_durable_journal_boundary() { + let (socket, handle, tmp) = start_slow_daemon().await; + let host = dummy_host(); + forward_envelope( + &envelope("fast-capture", "SessionStart", 1), + &socket, + &host, + false, + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(25)).await; + + let mut turn_end = envelope("fast-capture", "Stop", 2); + turn_end.route.as_mut().unwrap().flush_mode = bt_daemon::wire::FlushMode::FlushOnTurnEnd; + let accepted = tokio::time::timeout( + Duration::from_millis(100), + forward_envelope(&turn_end, &socket, &host, false), + ) + .await; + assert!( + matches!(accepted, Ok(Ok(()))), + "turn-end capture waited for translation or flushing: {accepted:?}" + ); + let journal = std::fs::read_to_string(source_journal_path( + &tmp.path().join("data"), + "debug", + "fast-capture", + )) + .unwrap(); + assert_eq!( + journal.lines().count(), + 2, + "hook returned before journaling" + ); + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[tokio::test] +async fn saturated_ingress_uses_the_journal_as_its_bounded_overflow_queue() { + const EVENT_COUNT: i64 = 1_100; + let gate = Arc::new(tokio::sync::Notify::new()); + let emitted = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let (data_dir, socket, handle, _tmp) = start_gated_daemon(gate.clone(), emitted.clone()).await; + let host = dummy_host(); + + forward_envelope( + &envelope("bounded-overflow", "SessionStart", 0), + &socket, + &host, + false, + ) + .await + .unwrap(); + for ts in 1..EVENT_COUNT { + forward_envelope( + &envelope("bounded-overflow", "PostToolUse", ts), + &socket, + &host, + false, + ) + .await + .unwrap(); + } + + let journal = + std::fs::read_to_string(source_journal_path(&data_dir, "debug", "bounded-overflow")) + .unwrap(); + assert_eq!( + journal.lines().count(), + EVENT_COUNT as usize, + "every hook returns only after its event reaches the journal" + ); + + gate.notify_one(); + let flushed = flush_session("bounded-overflow", &socket, 15_000) + .await + .unwrap(); + assert!(flushed.flushed, "flush did not complete: {flushed:?}"); + assert_eq!( + emitted.load(std::sync::atomic::Ordering::Relaxed), + EVENT_COUNT as u64 + 1, + "the root plus every journaled event must reach the sink exactly once" + ); + + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn status_reports_sessions_and_filters_by_session_id() { let (_data_dir, socket, handle, _tmp) = start_daemon().await; @@ -1259,6 +1536,35 @@ async fn cold_worker_rebuilds_acknowledged_journal_without_redelivery() { second.await.unwrap(); } +#[tokio::test] +async fn daemon_startup_recovers_an_event_journaled_before_worker_dispatch() { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let journal = source_journal_path(&data_dir, "debug", "crash-window"); + std::fs::create_dir_all(journal.parent().unwrap()).unwrap(); + let mut raw = + serde_json::to_vec(&envelope("crash-window", "SessionStart", 1).redacted()).unwrap(); + raw.push(b'\n'); + std::fs::write(&journal, raw).unwrap(); + + let daemon = start_daemon_at(data_dir.clone(), socket.clone()).await; + let flushed = flush_session("crash-window", &socket, 5000).await.unwrap(); + assert!( + flushed.flushed, + "recovered event did not drain: {flushed:?}" + ); + let spans = std::fs::read_to_string(data_dir.join("spans/crash-window.ndjson")).unwrap(); + assert_eq!( + spans.lines().count(), + 2, + "startup must translate the uncheckpointed journal event" + ); + + shutdown(&socket).await; + daemon.await.unwrap(); +} + #[tokio::test] async fn claude_boundary_journal_references_a_self_contained_transcript_mirror() { let (data_dir, socket, handle, tmp) = start_daemon().await; diff --git a/src/plugins/pi/content/src/config.ts b/src/plugins/pi/content/src/config.ts index 72deceb..4a6ab27 100644 --- a/src/plugins/pi/content/src/config.ts +++ b/src/plugins/pi/content/src/config.ts @@ -1,8 +1,8 @@ import { existsSync, readFileSync } from "node:fs"; import { homedir } from "node:os"; import { join } from "node:path"; -import * as piCodingAgent from "@earendil-works/pi-coding-agent"; import { type DaemonSessionRoute, resolveDaemonTraceSettings } from "./runtime/daemon-client.ts"; +import { loadPiPackageMetadata } from "./pi-package.ts"; export interface PiConfig { enabled: boolean; @@ -17,10 +17,7 @@ export interface PiConfig { type ConfigRecord = Record; -const PROJECT_CONFIG_DIR_NAME = - typeof (piCodingAgent as { CONFIG_DIR_NAME?: unknown }).CONFIG_DIR_NAME === "string" - ? (piCodingAgent as { CONFIG_DIR_NAME: string }).CONFIG_DIR_NAME - : ".pi"; +const PROJECT_CONFIG_DIR_NAME = loadPiPackageMetadata().configDir; function record(value: unknown): ConfigRecord | undefined { return value !== null && typeof value === "object" && !Array.isArray(value) diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index ffab7a2..71d2c0f 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -1,16 +1,14 @@ import { createHash, randomUUID } from "node:crypto"; import { resolve } from "node:path"; -import { - VERSION as PI_VERSION, - type ExtensionAPI, - type ExtensionContext, -} from "@earendil-works/pi-coding-agent"; +import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; import { loadConfig } from "./config.ts"; +import { loadPiPackageMetadata } from "./pi-package.ts"; import { claimManagedTracingInstance, DaemonClient } from "./runtime/daemon-client.ts"; import { EXTENSION_VERSION } from "./version.ts"; const STATUS_KEY = "braintrust-tracing"; const WIDGET_KEY = "braintrust-trace-link"; +const PI_VERSION = loadPiPackageMetadata().version; function sessionKeyFor( sessionFile: string | undefined, @@ -96,7 +94,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { if (!sessionId) return; await client.log({ source: "pi", - source_version: PI_VERSION, + ...(PI_VERSION ? { source_version: PI_VERSION } : {}), session_id: sessionId, event: name, ts_ms: Date.now(), diff --git a/src/plugins/pi/content/src/pi-package.ts b/src/plugins/pi/content/src/pi-package.ts new file mode 100644 index 0000000..3f03d69 --- /dev/null +++ b/src/plugins/pi/content/src/pi-package.ts @@ -0,0 +1,48 @@ +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { dirname, join } from "node:path"; + +interface PiPackageManifest { + version?: unknown; + piConfig?: { + configDir?: unknown; + }; +} + +export interface PiPackageMetadata { + version?: string; + configDir: string; +} + +/** + * Read Pi's package metadata without importing its root runtime barrel. New Pi + * releases may add optional entrypoints to that barrel whose dependencies are + * irrelevant to extensions, so loading it just for VERSION/CONFIG_DIR_NAME can + * make an otherwise compatible extension fail during module initialization. + */ +export function loadPiPackageMetadata(): PiPackageMetadata { + try { + const entry = createRequire(import.meta.url).resolve("@earendil-works/pi-coding-agent"); + let directory = dirname(entry); + while (true) { + try { + const manifest = JSON.parse( + readFileSync(join(directory, "package.json"), "utf8"), + ) as PiPackageManifest; + return { + version: typeof manifest.version === "string" ? manifest.version : undefined, + configDir: + typeof manifest.piConfig?.configDir === "string" ? manifest.piConfig.configDir : ".pi", + }; + } catch { + const parent = dirname(directory); + if (parent === directory) break; + directory = parent; + } + } + } catch { + // Pi supplies the extension API at runtime. Metadata discovery is useful + // for diagnostics but must not prevent the extension from loading. + } + return { configDir: ".pi" }; +}