Flush managed runs when child processes exit

bt-daemon/docs/protocol.md

@@ -124,6 +124,17 @@ Result:
 `flushed: false` with `pending > 0` means the timeout was hit with work
 outstanding. Used by session-end hooks and flush-on-turn-end mode.
 
+### `managed_run.flush` (request)
+
+Block until every session accepted from one `bt trace run` child process tree
+has flushed. Params:
+```json
+{ "managed_run_id": "…", "timeout_ms": 10000 }
+```
+The result has the same shape as `session.flush`. The managed-run identifier is
+invocation-local and prevents the shared daemon from flushing unrelated agent
+sessions.
+
 ### `status.get` (request)
 
 Params: `{ "session_id": "…" }` (omit `session_id` for daemon-wide status).
@@ -161,6 +172,7 @@ Used for version handover and by tests.
   "session_id": "0f9d…",
   "event": "PostToolUse",
   "ts_ms": 1753639552123,
+  "managed_run_id": "invocation-uuid",
   "payload": { "…raw agent-native hook payload…": true },
   "route": {
     "auth": {
@@ -192,6 +204,9 @@ Field notes:
   daemon.
 - **`payload`** is opaque to transport and to everything except the translator
   for `source`.
+- **`managed_run_id`** is present only for events inherited from a
+  `bt trace run` process tree. It groups native sessions for the final
+  invocation flush and is not trace metadata.
 - **`route`** carries non-secret auth selection and trace settings. `profile`
   is optional and resolves through `bt`'s default profile when absent;
   `org_name` optionally constrains organization selection. The daemon resolves

bt-daemon/src/journal.rs

@@ -124,6 +124,7 @@ pub fn envelope_from_redacted(r: RedactedEnvelope) -> Envelope {
         session_id: r.session_id,
         event: r.event,
         ts_ms: r.ts_ms,
+        managed_run_id: r.managed_run_id,
         payload: r.payload,
         route,
         config,

bt-daemon/src/lib.rs

@@ -41,7 +41,13 @@ use std::ffi::OsString;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
-use wire::{method, Envelope, SessionConfig, SessionRoute, StatusResult, PROTOCOL_VERSION};
+use wire::{
+    method, Envelope, ManagedRunFlushParams, SessionConfig, SessionRoute, StatusResult,
+    PROTOCOL_VERSION,
+};
+
+const MANAGED_RUN_ID_ENV: &str = "BT_TRACE_MANAGED_RUN_ID";
+const MANAGED_RUN_FLUSH_TIMEOUT_MS: u64 = 10_000;
 
 /// Arguments for `serve`.
 #[derive(Debug, Clone, Args)]
@@ -225,6 +231,9 @@ pub async fn run_hook(
         session_id,
         event,
         ts_ms: now_ms(),
+        managed_run_id: std::env::var(MANAGED_RUN_ID_ENV)
+            .ok()
+            .filter(|value| !value.is_empty()),
         payload,
         route: Some(route),
         config: None,
@@ -327,6 +336,39 @@ pub async fn flush_session(
     Ok(serde_json::from_value(value)?)
 }
 
+/// Flush every daemon session accepted from one managed child process tree.
+/// A missing daemon means the child emitted no accepted trace events.
+pub async fn flush_managed_run(
+    managed_run_id: &str,
+    socket: &std::path::Path,
+    timeout_ms: u64,
+) -> anyhow::Result<wire::FlushResult> {
+    let stream = match client::connect(socket).await {
+        Ok(stream) => stream,
+        Err(_) => {
+            return Ok(wire::FlushResult {
+                flushed: true,
+                pending: 0,
+            })
+        }
+    };
+    let mut conn = client::Conn::new(stream);
+    conn.request(
+        method::INITIALIZE,
+        serde_json::json!({
+            "protocol_version": PROTOCOL_VERSION,
+            "client": { "source": "managed-run-flush" }
+        }),
+    )
+    .await?;
+    let params = ManagedRunFlushParams {
+        managed_run_id: managed_run_id.to_string(),
+        timeout_ms,
+    };
+    let value = conn.request(method::MANAGED_RUN_FLUSH, params).await?;
+    Ok(serde_json::from_value(value)?)
+}
+
 /// Query daemon status. `Ok(None)` means no daemon is running.
 pub async fn run_status(args: StatusArgs) -> anyhow::Result<Option<StatusResult>> {
     let socket = paths::socket_path(args.socket.as_deref());
@@ -399,12 +441,14 @@ pub async fn run_traced(
     let executable =
         std::env::var_os(executable_env).unwrap_or_else(|| OsString::from(default_executable));
     let injected_args = managed_run_args(args.source, &hook_command)?;
+    let managed_run_id = uuid::Uuid::new_v4().to_string();
     let invocation_settings = serde_json::to_string(&settings::InvocationSettings::enabled(route))?;
     let mut command = tokio::process::Command::new(&executable);
     command
         .args(injected_args)
         .args(args.agent_args)
         .env("_BT_TRACE_MANAGED_RUN", "1")
+        .env(MANAGED_RUN_ID_ENV, &managed_run_id)
         .env(settings::INVOCATION_SETTINGS_ENV, invocation_settings);
     if args.source == RunSource::OpenCode {
         command.env(
@@ -418,14 +462,32 @@ pub async fn run_traced(
     let interrupt = tokio::signal::ctrl_c();
     tokio::pin!(interrupt);
 
-    tokio::select! {
-        status = child.wait() => Ok(status?),
+    let status = tokio::select! {
+        status = child.wait() => status.map_err(anyhow::Error::from),
         result = &mut interrupt => {
-            result?;
-            child.start_kill()?;
-            Ok(child.wait().await?)
+            match result {
+                Ok(()) => {
+                    let kill_result = child.start_kill();
+                    let wait_result = child.wait().await;
+                    kill_result
+                        .map_err(anyhow::Error::from)
+                        .and_then(|()| wait_result.map_err(anyhow::Error::from))
+                }
+                Err(error) => Err(error.into()),
+            }
         }
-    }
+    };
+    let socket = paths::socket_path(None);
+    match flush_managed_run(&managed_run_id, &socket, MANAGED_RUN_FLUSH_TIMEOUT_MS).await {
+        Ok(result) if result.flushed => {}
+        Ok(result) => tracing::warn!(
+            managed_run_id,
+            pending = result.pending,
+            "managed run trace flush timed out"
+        ),
+        Err(error) => tracing::warn!(managed_run_id, %error, "managed run trace flush failed"),
+    }
+    status
 }
 
 fn managed_run_args(

bt-daemon/src/server.rs

@@ -9,13 +9,13 @@ use crate::translate::Registry;
 use crate::transport::{self, Listener, ServerStream};
 use crate::wire::{
     error_code, method, Capabilities, Envelope, EventLogResult, FlushParams, FlushResult,
-    InitializeParams, InitializeResult, Message, Request, Response, RpcError, SessionStatus,
-    ShutdownResult, StatusParams, StatusResult, PROTOCOL_VERSION,
+    InitializeParams, InitializeResult, ManagedRunFlushParams, Message, Request, Response,
+    RpcError, SessionStatus, ShutdownResult, StatusParams, StatusResult, PROTOCOL_VERSION,
 };
 use crate::wire::{AuthSelection, BackendAuth, SessionRoute};
 use crate::{paths, ServeArgs};
 use async_trait::async_trait;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
@@ -77,6 +77,7 @@ pub struct Daemon {
     sink_factory: Arc<dyn SinkFactory>,
     auth_provider: Option<Arc<dyn AuthProvider>>,
     session_auth: tokio::sync::Mutex<HashMap<String, SessionAuthState>>,
+    managed_run_sessions: Mutex<HashMap<String, HashSet<String>>>,
     auth_errors: Mutex<HashMap<String, (String, String)>>,
     sessions: Mutex<HashMap<String, Arc<Session>>>,
     started: Instant,
@@ -94,6 +95,7 @@ impl Daemon {
             sink_factory: opts.sink_factory,
             auth_provider: opts.auth_provider,
             session_auth: tokio::sync::Mutex::new(HashMap::new()),
+            managed_run_sessions: Mutex::new(HashMap::new()),
             auth_errors: Mutex::new(HashMap::new()),
             sessions: Mutex::new(HashMap::new()),
             started: Instant::now(),
@@ -275,6 +277,50 @@ impl Daemon {
             .sum()
     }
 
+    fn record_managed_run_session(&self, managed_run_id: &str, session_id: &str) {
+        self.managed_run_sessions
+            .lock()
+            .unwrap()
+            .entry(managed_run_id.to_string())
+            .or_default()
+            .insert(session_id.to_string());
+    }
+
+    async fn flush_managed_run(&self, params: ManagedRunFlushParams) -> FlushResult {
+        let session_ids = self
+            .managed_run_sessions
+            .lock()
+            .unwrap()
+            .get(&params.managed_run_id)
+            .cloned()
+            .unwrap_or_default();
+        let mut result = FlushResult {
+            flushed: true,
+            pending: 0,
+        };
+        for session_id in session_ids {
+            if let Err(error) = self.refresh_session_before_flush(&session_id).await {
+                tracing::warn!(
+                    managed_run_id = %params.managed_run_id,
+                    session_id,
+                    %error,
+                    "managed run session auth refresh failed"
+                );
+                result.flushed = false;
+                continue;
+            }
+            let session = { self.sessions.lock().unwrap().get(&session_id).cloned() };
+            if let Some(session) = session {
+                let (flushed, pending) = session
+                    .flush(Duration::from_millis(params.timeout_ms))
+                    .await;
+                result.flushed &= flushed;
+                result.pending = result.pending.saturating_add(pending);
+            }
+        }
+        result
+    }
+
     fn trigger_shutdown(&self) {
         self.shutting_down.store(true, Ordering::SeqCst);
         self.shutdown.notify_waiters();
@@ -448,6 +494,7 @@ async fn accept_event(daemon: &Arc<Daemon>, mut env: Envelope) -> Result<(), Str
     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();
     tracing::info!(source, event, session_id, "event received");
     daemon.touch();
 
@@ -468,7 +515,12 @@ async fn accept_event(daemon: &Arc<Daemon>, mut env: Envelope) -> Result<(), Str
     .await;
 
     match &result {
-        Ok(()) => tracing::info!(source, event, session_id, "event accepted"),
+        Ok(()) => {
+            if let Some(managed_run_id) = managed_run_id {
+                daemon.record_managed_run_session(&managed_run_id, &session_id);
+            }
+            tracing::info!(source, event, session_id, "event accepted")
+        }
         Err(error) => tracing::warn!(source, event, session_id, error, "event rejected"),
     }
     result
@@ -547,6 +599,11 @@ async fn handle_request(daemon: &Arc<Daemon>, req: Request) -> Response {
                 serde_json::to_value(FlushResult { flushed, pending }).unwrap(),
             )
         }
+        method::MANAGED_RUN_FLUSH => {
+            let params = parse!(ManagedRunFlushParams);
+            let result = daemon.flush_managed_run(params).await;
+            Response::ok(id, serde_json::to_value(result).unwrap())
+        }
         method::STATUS_GET => {
             let p = parse!(StatusParams);
             Response::ok(id, serde_json::to_value(daemon.status(p)).unwrap())

bt-daemon/src/transcript_import.rs

@@ -521,6 +521,7 @@ fn envelope(
         session_id: session_id.into(),
         event: event.into(),
         ts_ms,
+        managed_run_id: None,
         payload,
         route: None,
         config: None,

bt-daemon/src/wire/envelope.rs

@@ -23,6 +23,10 @@ pub struct Envelope {
     pub event: String,
     /// Epoch milliseconds, stamped by the shim at capture time.
     pub ts_ms: i64,
+    /// Invocation-local identifier supplied by `bt trace run`. The daemon uses
+    /// it only to flush the sessions created by one managed child process tree.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub managed_run_id: Option<String>,
     /// The raw agent-native hook payload; opaque except to the translator.
     pub payload: serde_json::Value,
     /// Non-secret, immutable routing intent for this session. New clients use
@@ -198,6 +202,7 @@ impl Envelope {
             session_id: self.session_id.clone(),
             event: self.event.clone(),
             ts_ms: self.ts_ms,
+            managed_run_id: self.managed_run_id.clone(),
             payload: self.payload.clone(),
             route: self.route.clone(),
         }
@@ -216,6 +221,8 @@ pub struct RedactedEnvelope {
     pub session_id: String,
     pub event: String,
     pub ts_ms: i64,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub managed_run_id: Option<String>,
     pub payload: serde_json::Value,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub route: Option<SessionRoute>,
@@ -233,6 +240,7 @@ mod tests {
             session_id: "sess-1".into(),
             event: "PostToolUse".into(),
             ts_ms: 1_753_639_552_123,
+            managed_run_id: Some("run-1".into()),
             payload: serde_json::json!({ "session_id": "sess-1", "tool_name": "shell" }),
             route: Some(SessionRoute {
                 auth: AuthSelection {
@@ -266,6 +274,7 @@ mod tests {
         let s = serde_json::to_string(&e).unwrap();
         let back: Envelope = serde_json::from_str(&s).unwrap();
         assert_eq!(back.session_id, "sess-1");
+        assert_eq!(back.managed_run_id.as_deref(), Some("run-1"));
         assert_eq!(back.route.unwrap().auth.profile.as_deref(), Some("work"));
         assert!(back.config.is_none());
         assert!(!s.contains("sk-super-secret"));
@@ -281,6 +290,7 @@ mod tests {
             "token leaked into journal form: {s}"
         );
         assert_eq!(r.route.unwrap().auth.profile.as_deref(), Some("work"));
+        assert_eq!(r.managed_run_id.as_deref(), Some("run-1"));
     }
 
     #[test]

bt-daemon/src/wire/methods.rs

@@ -7,6 +7,7 @@ pub mod method {
     pub const INITIALIZE: &str = "initialize";
     pub const EVENT_LOG: &str = "event.log";
     pub const SESSION_FLUSH: &str = "session.flush";
+    pub const MANAGED_RUN_FLUSH: &str = "managed_run.flush";
     pub const STATUS_GET: &str = "status.get";
     pub const DAEMON_SHUTDOWN: &str = "daemon.shutdown";
 }
@@ -52,6 +53,13 @@ pub struct FlushParams {
     pub timeout_ms: u64,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ManagedRunFlushParams {
+    pub managed_run_id: String,
+    #[serde(default = "default_flush_timeout_ms")]
+    pub timeout_ms: u64,
+}
+
 fn default_flush_timeout_ms() -> u64 {
     10_000
 }

bt-daemon/src/wire/mod.rs

@@ -15,7 +15,8 @@ pub use envelope::{
 };
 pub use methods::{
     method, Capabilities, ClientInfo, EventLogResult, FlushParams, FlushResult, InitializeParams,
-    InitializeResult, SessionStatus, ShutdownResult, StatusParams, StatusResult,
+    InitializeResult, ManagedRunFlushParams, SessionStatus, ShutdownResult, StatusParams,
+    StatusResult,
 };
 pub use rpc::{error_code, Message, Request, RequestId, Response, RpcError};