diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index f10f67e..024b42f 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -162,14 +162,23 @@ impl Session { self.last_activity.lock().unwrap().elapsed() } + /// Insert a flush in actor order without waiting for backend delivery. + pub(crate) async fn enqueue_flush(&self) -> anyhow::Result> { + self.touch(); + let (reply_tx, reply_rx) = oneshot::channel(); + self.tx + .send(SessionMsg::Flush(reply_tx)) + .await + .map_err(|_| anyhow::anyhow!("session actor is gone"))?; + Ok(reply_rx) + } + /// Ask the actor to drain and flush its sink, bounded by `timeout`. /// Returns `(flushed, pending)`. pub async fn flush(&self, timeout: std::time::Duration) -> (bool, u64) { - self.touch(); - let (reply_tx, reply_rx) = oneshot::channel(); - if self.tx.send(SessionMsg::Flush(reply_tx)).await.is_err() { + let Ok(reply_rx) = self.enqueue_flush().await else { return (false, self.counters.queued.load(Ordering::Relaxed)); - } + }; match tokio::time::timeout(timeout, reply_rx).await { Ok(Ok(pending)) => (pending == 0, pending), _ => (false, self.counters.queued.load(Ordering::Relaxed)), @@ -847,6 +856,54 @@ pub(crate) fn is_tool_lifecycle_event(event: &str) -> bool { mod tests { use super::*; + #[tokio::test] + async fn enqueue_flush_preserves_boundary_without_waiting_for_delivery() { + let (tx, mut rx) = mpsc::channel(4); + let session = Session { + source: "pi".into(), + tx, + counters: Arc::new(Counters::default()), + last_error: Arc::new(Mutex::new(None)), + permalink: Arc::new(Mutex::new(None)), + last_activity: Mutex::new(Instant::now()), + }; + let envelope = |event| { + serde_json::from_value::(serde_json::json!({ + "source": "pi", "session_id": "ordered-turns", "event": event, + "ts_ms": 1, "payload": {"event": {}} + })) + .unwrap() + }; + session.enqueue(envelope("agent_end"), 1).await.unwrap(); + // No actor is reading yet: enqueue must complete independently of delivery. + let mut completion = + tokio::time::timeout(std::time::Duration::from_secs(1), session.enqueue_flush()) + .await + .expect("enqueue waited for delivery") + .unwrap(); + session + .enqueue(envelope("before_agent_start"), 2) + .await + .unwrap(); + assert!(matches!( + completion.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + assert!( + matches!(rx.recv().await, Some(SessionMsg::Event(env, 1)) if env.event == "agent_end") + ); + let Some(SessionMsg::Flush(reply)) = rx.recv().await else { + panic!("next turn overtook the boundary flush"); + }; + assert!( + matches!(rx.recv().await, Some(SessionMsg::Event(env, 2)) if env.event == "before_agent_start") + ); + reply.send(1).unwrap(); + assert_eq!(completion.await.unwrap(), 1); + drop(rx); + assert!(session.enqueue_flush().await.is_err()); + } + #[tokio::test] async fn grok_hydration_mirrors_transcripts_and_system_prompt_at_one_boundary() { let tmp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 25d5516..8ae2ef0 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -342,6 +342,27 @@ pub(crate) fn should_flush_hook_event(event: &str, flush_on_turn_end: bool) -> b )) } +pub(crate) fn should_flush_ingress_event(env: &wire::Envelope) -> bool { + let native = env.payload.get("event").unwrap_or(&env.payload); + let flush_on_turn_end = matches!( + env.route.as_ref().map(|route| route.flush_mode), + Some(wire::FlushMode::FlushOnTurnEnd) + ); + should_flush_hook_event(&env.event, flush_on_turn_end) + || (env.source == "pi" + && match env.event.as_str() { + // Preserve Pi's explicit lifecycle flushes even in batched mode. + "session_shutdown" | "session_compact" | "session_tree" => true, + // OMP can end an attempt while keeping the same turn open for retry. + "agent_end" => { + flush_on_turn_end + && native.get("willRetry").and_then(serde_json::Value::as_bool) + != Some(true) + } + _ => false, + }) +} + /// Capture one hook event from `stdin` and forward it to the daemon. /// /// `route` contains only non-secret profile and destination selection. @@ -1451,6 +1472,48 @@ mod tests { assert!(!should_flush_hook_event("turn_completed", true)); } + #[test] + fn pi_agent_end_flushes_only_completed_turns_when_enabled() { + let mut env: wire::Envelope = serde_json::from_value(serde_json::json!({ + "source": "pi", "session_id": "pi-turn", "event": "agent_end", + "ts_ms": 1, "payload": {"event": {}}, + "route": {"flush_mode": "flush_on_turn_end"} + })) + .unwrap(); + for native in [ + serde_json::json!({}), + serde_json::json!({"willRetry": false}), + ] { + env.payload = serde_json::json!({"event": native}); + assert!(should_flush_ingress_event(&env)); + } + for payload in [ + serde_json::json!({"event": {"willRetry": true}}), + serde_json::json!({"willRetry": true}), + ] { + env.payload = payload; + assert!(!should_flush_ingress_event(&env)); + } + env.payload = serde_json::json!({"event": {}}); + env.source = "opencode".into(); + assert!(!should_flush_ingress_event(&env)); + env.source = "pi".into(); + env.event = "message_end".into(); + assert!(!should_flush_ingress_event(&env)); + env.event = "agent_end".into(); + env.route.as_mut().unwrap().flush_mode = wire::FlushMode::FireAndForget; + assert!(!should_flush_ingress_event(&env)); + env.route = None; + assert!(!should_flush_ingress_event(&env)); + for event in ["session_shutdown", "session_compact", "session_tree"] { + env.event = event.into(); + assert!(should_flush_ingress_event(&env)); + env.source = "opencode".into(); + assert!(!should_flush_ingress_event(&env)); + env.source = "pi".into(); + } + } + #[test] fn additional_metadata_overrides_a_route_only_with_a_json_object() { let mut route = SessionRoute { diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index 5d9f43c..3148b63 100644 --- a/bt-daemon/src/server.rs +++ b/bt-daemon/src/server.rs @@ -753,10 +753,16 @@ impl Daemon { result } - /// Flush every live delivery route for one source session. This is used + /// 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 flush_source_session(&self, source: &str, session_id: &str, timeout: Duration) { + async fn enqueue_source_session_flush( + &self, + source: &str, + session_id: &str, + timeout: Duration, + ) { let sessions: Vec<_> = self .sessions .lock() @@ -766,15 +772,26 @@ impl Daemon { .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" - ); - } + 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" + ); + } + }); } } @@ -845,13 +862,7 @@ async fn dispatch_ingress_event(daemon: &Arc, event: PendingEvent) { return; } - let schedule_flush = crate::should_flush_hook_event( - &event.env.event, - matches!( - event.env.route.as_ref().map(|route| route.flush_mode), - Some(crate::wire::FlushMode::FlushOnTurnEnd) - ), - ); + let schedule_flush = crate::should_flush_ingress_event(&event.env); let flush_source = event.env.source.clone(); let flush_session_id = event.env.session_id.clone(); let journal_through = event.journal_through; @@ -874,12 +885,9 @@ async fn dispatch_ingress_event(daemon: &Arc, event: PendingEvent) { } } 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; - }); + daemon + .enqueue_source_session_flush(&flush_source, &flush_session_id, Duration::from_secs(10)) + .await; } } diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 203f9bf..5cd08e2 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -1257,6 +1257,53 @@ async fn distinct_sessions_are_isolated() { handle.abort(); } +#[tokio::test] +async fn pi_lifecycle_flushes_without_an_explicit_client_flush() { + let (socket, handle, flushes, _tmp) = start_tracking_daemon("test").await; + let host = dummy_host(); + for event in [ + "agent_end", + "session_compact", + "session_tree", + "session_shutdown", + ] { + let session_id = format!("pi-background-{event}"); + for (index, name) in ["session_start", "before_agent_start", event] + .iter() + .enumerate() + { + let mut env = envelope(&session_id, name, index as i64 + 1); + env.source = "pi".into(); + env.payload = serde_json::json!({"event": {"prompt": "hello"}}); + // Session lifecycle flushing must also work when turn flushing is off. + if event == "agent_end" { + env.route.as_mut().unwrap().flush_mode = bt_daemon::wire::FlushMode::FlushOnTurnEnd; + } + forward_envelope(&env, &socket, &host, false).await.unwrap(); + } + // forward_envelope has already closed its socket: delivery belongs to the daemon. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if flushes + .lock() + .unwrap() + .get(&session_id) + .copied() + .unwrap_or_default() + > 0 + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("Pi {event} was not flushed by the daemon")); + } + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn hook_capture_stops_at_the_durable_journal_boundary() { let (socket, handle, tmp) = start_slow_daemon().await; diff --git a/src/plugins/pi/content/src/daemon-adapter.test.ts b/src/plugins/pi/content/src/daemon-adapter.test.ts index c942e1f..1254992 100644 --- a/src/plugins/pi/content/src/daemon-adapter.test.ts +++ b/src/plugins/pi/content/src/daemon-adapter.test.ts @@ -5,6 +5,7 @@ const mockState = vi.hoisted(() => ({ flushes: [] as string[], closed: 0, claim: true, + logGate: undefined as Promise | undefined, })); vi.mock("./runtime/daemon-client.ts", () => ({ @@ -12,6 +13,7 @@ vi.mock("./runtime/daemon-client.ts", () => ({ DaemonClient: class { async log(envelope: Record): Promise { mockState.logs.push(envelope); + await mockState.logGate; return true; } async flush(sessionId: string): Promise { @@ -63,6 +65,7 @@ describe("Pi daemon adapter", () => { mockState.flushes.length = 0; mockState.closed = 0; mockState.claim = true; + mockState.logGate = undefined; }); it("does not register a duplicate managed adapter instance", async () => { @@ -123,13 +126,35 @@ describe("Pi daemon adapter", () => { await handlers.get("session_start")?.({ reason: "new" }, ctx); await handlers.get("before_agent_start")?.({ prompt: "hello" }, ctx); - await handlers.get("agent_end")?.({ messages: [] }); + let acknowledge!: () => void; + mockState.logGate = new Promise((resolve) => { + acknowledge = resolve; + }); + let turnEnded = false; + const turnEnd = handlers + .get("agent_end")?.({ messages: [] }) + .then(() => { + turnEnded = true; + }); + await Promise.resolve(); + expect(turnEnded).toBe(false); + acknowledge(); + await turnEnd; + mockState.logGate = undefined; + expect(mockState.flushes).toHaveLength(0); + await handlers.get("input")?.({ text: "next turn" }); + expect(mockState.logs.at(-1)?.event).toBe("input"); + await handlers.get("session_compact")?.({}, ctx); + await handlers.get("session_tree")?.({}, ctx); await handlers.get("session_shutdown")?.({ reason: "quit" }, ctx); expect(mockState.logs.map((log) => log.event)).toEqual([ "session_start", "before_agent_start", "agent_end", + "input", + "session_compact", + "session_tree", "session_shutdown", ]); expect( @@ -149,7 +174,7 @@ describe("Pi daemon adapter", () => { cwd: "/tmp/project", model: { provider: "openai", id: "gpt-5" }, }); - expect(mockState.flushes).toHaveLength(2); + expect(mockState.flushes).toHaveLength(0); expect(widgets).toContainEqual([ "braintrust-trace-link", ["Braintrust trace", "https://www.braintrust.dev/trace/1"], diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 71d2c0f..0445fe2 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -88,7 +88,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { name: string, event: unknown, ctx?: ExtensionContext, - flush = false, + updateUi = false, ): Promise => { const descriptor = ctx ? remember(ctx) : undefined; if (!sessionId) return; @@ -108,10 +108,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { }, route: config.route, }); - if (flush) { - await client.flush(sessionId); - if (ctx) await refreshUi(ctx); - } + if (updateUi && ctx) await refreshUi(ctx); }; pi.on("session_start", async (event, ctx) => { @@ -134,9 +131,10 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { pi.on("session_compact", async (event, ctx) => forward("session_compact", event, ctx, true)); pi.on("session_before_tree", async (event, ctx) => forward("session_before_tree", event, ctx)); pi.on("session_tree", async (event, ctx) => forward("session_tree", event, ctx, true)); - pi.on("agent_end", async (event) => forward("agent_end", event, undefined, true)); + // All delivery flushing belongs to the daemon, including session shutdown. + pi.on("agent_end", async (event) => forward("agent_end", event)); pi.on("session_shutdown", async (event, ctx) => { - await forward("session_shutdown", event, ctx, true); + await forward("session_shutdown", event, ctx); if (ctx.hasUI) { ctx.ui.setStatus(STATUS_KEY, undefined); ctx.ui.setWidget(WIDGET_KEY, undefined);