diff --git a/bt-daemon/docs/protocol.md b/bt-daemon/docs/protocol.md index f367409..937a091 100644 --- a/bt-daemon/docs/protocol.md +++ b/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 diff --git a/bt-daemon/src/journal.rs b/bt-daemon/src/journal.rs index 083d628..a220f2d 100644 --- a/bt-daemon/src/journal.rs +++ b/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, diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index ac53f71..310f1ca 100644 --- a/bt-daemon/src/lib.rs +++ b/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 { + 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> { 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( diff --git a/bt-daemon/src/server.rs b/bt-daemon/src/server.rs index ad47be2..52d0bb3 100644 --- a/bt-daemon/src/server.rs +++ b/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, auth_provider: Option>, session_auth: tokio::sync::Mutex>, + managed_run_sessions: Mutex>>, auth_errors: Mutex>, sessions: Mutex>>, 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(¶ms.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, 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, 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, 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()) diff --git a/bt-daemon/src/transcript_import.rs b/bt-daemon/src/transcript_import.rs index 432f3e1..648eda0 100644 --- a/bt-daemon/src/transcript_import.rs +++ b/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, diff --git a/bt-daemon/src/wire/envelope.rs b/bt-daemon/src/wire/envelope.rs index 0455b7f..f61cc4f 100644 --- a/bt-daemon/src/wire/envelope.rs +++ b/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, /// 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, pub payload: serde_json::Value, #[serde(default, skip_serializing_if = "Option::is_none")] pub route: Option, @@ -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] diff --git a/bt-daemon/src/wire/methods.rs b/bt-daemon/src/wire/methods.rs index c4c0fbb..9ca17cc 100644 --- a/bt-daemon/src/wire/methods.rs +++ b/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 } diff --git a/bt-daemon/src/wire/mod.rs b/bt-daemon/src/wire/mod.rs index 8ede81d..289222b 100644 --- a/bt-daemon/src/wire/mod.rs +++ b/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}; diff --git a/bt-daemon/tests/claude_translator.rs b/bt-daemon/tests/claude_translator.rs index a2a32a9..c94cac1 100644 --- a/bt-daemon/tests/claude_translator.rs +++ b/bt-daemon/tests/claude_translator.rs @@ -50,6 +50,7 @@ fn replay(name: &str) -> Vec { session_id: session_id.into(), event: record["hook"].as_str().unwrap().into(), ts_ms, + managed_run_id: None, payload, route: None, config: None, @@ -259,6 +260,7 @@ fn claude_permission_denied_and_failed_tools_are_first_class_spans() { session_id: "s".into(), event: name.into(), ts_ms: 1, + managed_run_id: None, payload, route: None, config: None, @@ -347,6 +349,7 @@ fn claude_pairs_tool_lifecycle_and_marks_explicit_skills_and_stop_failures() { session_id: "lifecycle".into(), event: name.into(), ts_ms, + managed_run_id: None, payload, route: None, config: None, @@ -465,6 +468,7 @@ fn claude_groups_streamed_rows_and_reads_late_final_output_at_session_end() { session_id: "streamed".into(), event: name.into(), ts_ms, + managed_run_id: None, payload, route: None, config: None, diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index dc0eae3..8f94d1a 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -61,6 +61,7 @@ fn envelope(session: &str, event: &str, transcript_path: &str, extra: Value) -> session_id: session.into(), event: event.into(), ts_ms: 0, + managed_run_id: None, payload, route: None, config: None, diff --git a/bt-daemon/tests/opencode_translator.rs b/bt-daemon/tests/opencode_translator.rs index ad6d1bf..a63ff98 100644 --- a/bt-daemon/tests/opencode_translator.rs +++ b/bt-daemon/tests/opencode_translator.rs @@ -11,6 +11,7 @@ fn event(name: &str, ts_ms: i64, payload: serde_json::Value) -> Envelope { session_id: "root-session".into(), event: name.into(), ts_ms, + managed_run_id: None, payload, route: None, config: None, diff --git a/bt-daemon/tests/pi_translator.rs b/bt-daemon/tests/pi_translator.rs index 7143f2e..0d53bc7 100644 --- a/bt-daemon/tests/pi_translator.rs +++ b/bt-daemon/tests/pi_translator.rs @@ -11,6 +11,7 @@ fn event(name: &str, ts_ms: i64, native: serde_json::Value) -> Envelope { session_id: "pi-session".into(), event: name.into(), ts_ms, + managed_run_id: None, payload: json!({"event":native,"extension_version":"1.0.0","cwd":"."}), route: None, config: None, diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 520ed27..dfa9d46 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -3,17 +3,63 @@ //! spawning, so it's deterministic). use async_trait::async_trait; -use bt_daemon::wire::{AuthSelection, BackendAuth, Envelope, SessionRoute}; +use bt_daemon::wire::{AuthSelection, BackendAuth, Envelope, SessionConfig, SessionRoute}; use bt_daemon::{ - debug_serve_options, flush_session, forward_envelope, run_serve, run_status, shutdown_daemon, - AuthLease, AuthProvider, AuthResolveReason, HostInfo, ServeArgs, StatusArgs, + debug_serve_options, flush_managed_run, flush_session, forward_envelope, run_serve, run_status, + shutdown_daemon, AuthLease, AuthProvider, AuthResolveReason, HostInfo, Registry, ServeArgs, + ServeOptions, Sink, SinkFactory, SpanOp, StatusArgs, }; +#[cfg(all(feature = "cli", unix))] +use bt_daemon::{run_traced, RunArgs, RunHookCommand, RunSource}; +use std::collections::HashMap; use std::ffi::OsString; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::Mutex; use std::time::Duration; +struct TrackingSinkFactory { + flushes: Arc>>, +} + +impl SinkFactory for TrackingSinkFactory { + fn create( + &self, + session_id: &str, + _source: &str, + _plugin_version: Option<&str>, + ) -> anyhow::Result> { + Ok(Box::new(TrackingSink { + session_id: session_id.to_string(), + flushes: self.flushes.clone(), + })) + } +} + +struct TrackingSink { + session_id: String, + flushes: Arc>>, +} + +#[async_trait] +impl Sink for TrackingSink { + fn configure(&mut self, _config: &SessionConfig) {} + + async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { + Ok(ops.len() as u64) + } + + async fn flush(&mut self) -> anyhow::Result<()> { + *self + .flushes + .lock() + .unwrap() + .entry(self.session_id.clone()) + .or_default() += 1; + Ok(()) + } +} + fn dummy_host() -> HostInfo { // The daemon is started in-process, so the client never spawns; serve_argv // is unused but must be non-empty. @@ -31,6 +77,7 @@ fn envelope(session_id: &str, event: &str, ts_ms: i64) -> Envelope { session_id: session_id.into(), event: event.into(), ts_ms, + managed_run_id: None, payload: serde_json::json!({ "session_id": session_id, "hook_event_name": event, "n": ts_ms }), route: Some(SessionRoute { destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { @@ -200,6 +247,42 @@ async fn start_daemon() -> ( (data_dir, socket, handle, tmp) } +async fn start_tracking_daemon( + version: &str, +) -> ( + PathBuf, + tokio::task::JoinHandle<()>, + Arc>>, + tempfile::TempDir, +) { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let flushes = Arc::new(Mutex::new(HashMap::new())); + let opts = ServeOptions { + version: version.to_string(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(TrackingSinkFactory { + flushes: flushes.clone(), + }), + 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, + }; + let handle = tokio::spawn(async move { + let _ = run_serve(args, opts).await; + }); + wait_for(&socket).await; + (socket, handle, flushes, tmp) +} + async fn start_daemon_at(data_dir: PathBuf, socket: PathBuf) -> tokio::task::JoinHandle<()> { let args = ServeArgs { socket: Some(socket.clone()), @@ -681,3 +764,139 @@ async fn claude_boundary_journal_contains_a_self_contained_transcript_snapshot() assert!(!journal.contains("sk-TOP-SECRET-abc123")); handle.abort(); } + +#[tokio::test] +async fn managed_run_flush_is_scoped_to_its_accepted_sessions() { + let (socket, handle, flushes, _tmp) = start_tracking_daemon("test").await; + let host = dummy_host(); + for (session_id, managed_run_id) in [ + ("run-a-main", "run-a"), + ("run-a-child", "run-a"), + ("run-b-main", "run-b"), + ] { + let mut env = envelope(session_id, "SessionStart", 1); + env.managed_run_id = Some(managed_run_id.to_string()); + forward_envelope(&env, &socket, &host, false).await.unwrap(); + } + + let result = flush_managed_run("run-a", &socket, 5_000).await.unwrap(); + assert!(result.flushed, "managed run did not flush: {result:?}"); + assert_eq!(result.pending, 0); + assert_eq!( + *flushes.lock().unwrap(), + HashMap::from([ + ("run-a-main".to_string(), 1), + ("run-a-child".to_string(), 1) + ]) + ); + + let result = flush_managed_run("run-b", &socket, 5_000).await.unwrap(); + assert!(result.flushed, "managed run did not flush: {result:?}"); + assert_eq!(flushes.lock().unwrap().get("run-b-main"), Some(&1)); + + shutdown(&socket).await; + handle.await.unwrap(); +} + +#[cfg(all(feature = "cli", unix))] +struct EnvVarGuard { + key: &'static str, + previous: Option, +} + +#[cfg(all(feature = "cli", unix))] +impl EnvVarGuard { + fn set(key: &'static str, value: impl AsRef) -> Self { + let previous = std::env::var_os(key); + std::env::set_var(key, value); + Self { key, previous } + } +} + +#[cfg(all(feature = "cli", unix))] +impl Drop for EnvVarGuard { + fn drop(&mut self) { + match &self.previous { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } +} + +#[cfg(all(feature = "cli", unix))] +#[tokio::test] +async fn managed_run_flushes_after_success_failure_and_signal_exit() { + use std::os::unix::fs::PermissionsExt; + + let version = env!("CARGO_PKG_VERSION"); + let (socket, handle, flushes, tmp) = start_tracking_daemon(version).await; + let agent = tmp.path().join("fake-codex.sh"); + std::fs::write( + &agent, + r#"#!/bin/sh +previous= +last= +for argument in "$@"; do + previous=$last + last=$argument +done +session_id=$previous +mode=$last +printf '{"session_id":"%s","hook_event_name":"SessionStart"}\n' "$session_id" | + "$BT_DAEMON_TEST_BIN" hook --source debug --managed-run-hook --no-spawn +case "$mode" in + success) exit 0 ;; + failure) exit 7 ;; + signal) kill -TERM "$$" ;; + *) exit 99 ;; +esac +"#, + ) + .unwrap(); + std::fs::set_permissions(&agent, std::fs::Permissions::from_mode(0o755)).unwrap(); + + let _socket = EnvVarGuard::set("BT_DAEMON_SOCKET", &socket); + let _agent = EnvVarGuard::set("CODEX_BIN", &agent); + let _daemon = EnvVarGuard::set("BT_DAEMON_TEST_BIN", env!("CARGO_BIN_EXE_bt-daemon")); + let route = || SessionRoute { + destination: Some(bt_daemon::wire::TraceDestination::ProjectLogs { + project_id: None, + project_name: Some("managed-run-test".into()), + }), + ..SessionRoute::default() + }; + let run = |session_id: &str, mode: &str| { + run_traced( + RunArgs { + source: RunSource::Codex, + agent_args: vec![session_id.into(), mode.into()], + }, + RunHookCommand { + program: "unused-hook-command".into(), + args: Vec::new(), + }, + route(), + ) + }; + + assert!(run("managed-success", "success").await.unwrap().success()); + assert_eq!( + run("managed-failure", "failure").await.unwrap().code(), + Some(7) + ); + let signal_status = run("managed-signal", "signal").await.unwrap(); + assert!(!signal_status.success()); + assert_eq!(signal_status.code(), None); + + assert_eq!( + *flushes.lock().unwrap(), + HashMap::from([ + ("managed-success".to_string(), 1), + ("managed-failure".to_string(), 1), + ("managed-signal".to_string(), 1), + ]) + ); + + shutdown(&socket).await; + handle.await.unwrap(); +} diff --git a/src/runtime/js-daemon-client/src/index.ts b/src/runtime/js-daemon-client/src/index.ts index 965eaef..1e6c6b9 100644 --- a/src/runtime/js-daemon-client/src/index.ts +++ b/src/runtime/js-daemon-client/src/index.ts @@ -53,6 +53,7 @@ export interface DaemonEnvelope { session_id: string event: string ts_ms: number + managed_run_id?: string payload: unknown route?: DaemonSessionRoute } @@ -81,6 +82,7 @@ export interface DaemonClientOptions { connectAttempts?: number connectDelayMs?: number requestTimeoutMs?: number + managedRunId?: string warn?: (message: string) => void } @@ -134,6 +136,7 @@ export class DaemonClient { connectAttempts: 50, connectDelayMs: 20, requestTimeoutMs: 10_000, + managedRunId: process.env.BT_TRACE_MANAGED_RUN_ID, ...options, } } @@ -143,6 +146,7 @@ export class DaemonClient { const event = { ...envelope, ...(this.options.pluginVersion ? { plugin_version: this.options.pluginVersion } : {}), + ...(this.options.managedRunId ? { managed_run_id: this.options.managedRunId } : {}), } try { const result = (await this.request("event.log", event)) as { accepted?: boolean } diff --git a/src/runtime/js-daemon-client/tests/client.test.ts b/src/runtime/js-daemon-client/tests/client.test.ts index 6bd0a6a..53eff79 100644 --- a/src/runtime/js-daemon-client/tests/client.test.ts +++ b/src/runtime/js-daemon-client/tests/client.test.ts @@ -101,6 +101,7 @@ test("serializes initialize, events, flush, and status over one connection", asy const client = new DaemonClient({ source: "opencode", pluginVersion: "1.0.0", + managedRunId: "run-123", socketPath: endpoint, }) const envelope = (name: string) => ({ @@ -124,6 +125,7 @@ test("serializes initialize, events, flush, and status over one connection", asy "status.get", ]) assert.deepEqual(eventParams.map((params) => params.plugin_version), ["1.0.0", "1.0.0"]) + assert.deepEqual(eventParams.map((params) => params.managed_run_id), ["run-123", "run-123"]) await client.close() await new Promise((resolve) => server.close(() => resolve())) rmSync(temp, { recursive: true, force: true })