fix(pi): move delivery flushing into the daemon

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<oneshot::Receiver<u64>> {
+        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::<Envelope>(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();

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 {

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<Daemon>, 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<Daemon>, 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;
     }
 }
 

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;

src/plugins/pi/content/src/daemon-adapter.test.ts

@@ -5,13 +5,15 @@ const mockState = vi.hoisted(() => ({
   flushes: [] as string[],
   closed: 0,
   claim: true,
+  logGate: undefined as Promise<void> | undefined,
 }));
 
 vi.mock("./runtime/daemon-client.ts", () => ({
   claimManagedTracingInstance: () => mockState.claim,
   DaemonClient: class {
     async log(envelope: Record<string, unknown>): Promise<boolean> {
       mockState.logs.push(envelope);
+      await mockState.logGate;
       return true;
     }
     async flush(sessionId: string): Promise<boolean> {
@@ -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<void>((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"],