diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 0eeb285..d7412d9 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -270,6 +270,13 @@ impl BraintrustSink { Ok(client) } + fn span_origin(&self) -> SpanOrigin { + SpanOrigin::new() + .name(format!("braintrust.plugin.{}", self.source)) + .version(self.version.clone()) + .instrumentation("braintrust-plugin") + } + fn ensure_handle(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { if self.open.contains_key(&row.span_id) { return Ok(()); @@ -287,12 +294,7 @@ impl BraintrustSink { .span_id(row.span_id.clone()) .row_id(row.span_id.clone()) .parent_info(parent) - .span_origin( - SpanOrigin::new() - .name(format!("braintrust.plugin.{}", self.source)) - .version(self.version.clone()) - .instrumentation("braintrust-plugin"), - ); + .span_origin(self.span_origin()); if let Some(org_name) = &creds.org_name { builder = builder.org_name(org_name.clone()); } @@ -306,7 +308,7 @@ impl BraintrustSink { fn update_open(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { self.ensure_handle(client, row)?; let handle = self.open.get(&row.span_id).expect("just inserted"); - handle.log(build_log(row, &self.daemon_version)?); + handle.log(build_log(row, &self.daemon_version, self.span_origin())?); if let Some(end) = row.end_ms { handle.end_with_time(ms_to_secs(end)); // SpanHandle retains the complete accumulated input/output. Once a @@ -329,7 +331,7 @@ impl BraintrustSink { creds.token.clone(), creds.org_id.clone(), &components, - build_log(row, &self.daemon_version)?, + build_log(row, &self.daemon_version, self.span_origin())?, ) .map_err(|error| anyhow::anyhow!("braintrust span merge failed: {error}")) } @@ -627,20 +629,22 @@ fn ms_to_secs(ms: i64) -> f64 { ms as f64 / 1000.0 } -fn build_log(row: &SpanRow, daemon_version: &str) -> anyhow::Result { +fn build_log(row: &SpanRow, daemon_version: &str, origin: SpanOrigin) -> anyhow::Result { + // Repeat the plugin origin so stateless merges cannot use SDK defaults. + let mut builder = SpanLog::builder().span_origin(origin); + // The span's display name is carried on the log event, not the builder. // An empty name means "unchanged" (many merge ops use `..Default::default()` // and don't rename the span) — omitting `.name()` avoids overwriting the // already-set name with an empty string on merge. - let mut lb = SpanLog::builder(); if !row.name.is_empty() { - lb = lb.name(row.name.clone()); + builder = builder.name(row.name.clone()); } if let Some(input) = &row.input { - lb = lb.input(input.clone()); + builder = builder.input(input.clone()); } if let Some(output) = &row.output { - lb = lb.output(output.clone()); + builder = builder.output(output.clone()); } let mut metadata = row .metadata @@ -655,7 +659,7 @@ fn build_log(row: &SpanRow, daemon_version: &str) -> anyhow::Result { "bt_daemon_version".into(), Value::String(daemon_version.to_string()), ); - lb = lb.metadata(metadata); + builder = builder.metadata(metadata); let mut metrics = row .metrics .as_ref() @@ -673,17 +677,18 @@ fn build_log(row: &SpanRow, daemon_version: &str) -> anyhow::Result { metrics.insert("end".into(), ms_to_secs(end)); } if !metrics.is_empty() { - lb = lb.metrics(metrics); + builder = builder.metrics(metrics); } if let Some(err) = &row.error { - lb = lb.error(Value::String(err.clone())); + builder = builder.error(Value::String(err.clone())); } if let Some(tags) = &row.tags { if !tags.is_empty() { - lb = lb.tags(tags.clone()); + builder = builder.tags(tags.clone()); } } - lb.build() + builder + .build() .map_err(|e| anyhow::anyhow!("span log build failed: {e}")) } diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 9f7159c..bf17e8a 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -410,49 +410,81 @@ async fn late_merge_updates_a_completed_span_without_an_open_handle() { app_url: Some(base.clone()), version: "test".into(), }); - let mut sink = factory.create("sess-late", "codex", None).unwrap(); - sink.configure(&session_config(&base)); - - sink.emit(&[SpanOp::Insert(row( - "finished", - "trace-root", - &["turn-parent"], - "original name", - SpanType::Tool, - 1, - Some(2), - ))]) - .await - .unwrap(); - sink.flush().await.unwrap(); - let mut late = SpanRow { - span_id: "finished".into(), - root_span_id: "trace-root".into(), - parent_span_ids: vec!["turn-parent".into()], - output: Some(json!({"status":"late"})), - ..Default::default() - }; - late.name.clear(); - sink.emit(&[SpanOp::Merge(late)]).await.unwrap(); - sink.flush().await.unwrap(); + // The same cached SDK client serves every source. Provenance must stay + // session-specific, including when the plugin version falls back to bt's. + for (source, plugin_version) in [ + ("codex", Some("1.2.3")), + ("claude", Some("2.3.4")), + ("grok", Some("3.4.5")), + ("opencode", Some("4.5.6")), + ("pi", Some("5.6.7")), + ("antigravity", None), + ] { + let session_id = format!("sess-late-{source}"); + let span_id = format!("finished-{source}"); + let mut sink = factory.create(&session_id, source, plugin_version).unwrap(); + sink.configure(&session_config(&base)); + + sink.emit(&[SpanOp::Insert(row( + &span_id, + "trace-root", + &["turn-parent"], + "original name", + SpanType::Tool, + 1, + Some(2), + ))]) + .await + .unwrap(); + sink.flush().await.unwrap(); + let late = SpanRow { + span_id: span_id.clone(), + root_span_id: "trace-root".into(), + parent_span_ids: vec!["turn-parent".into()], + output: Some(json!({"status":"late"})), + ..Default::default() + }; + sink.emit(&[SpanOp::Merge(late.clone())]).await.unwrap(); + sink.flush().await.unwrap(); + + // Recovery also starts without an open handle for an existing span. + drop(sink); + let mut recovered = factory.create(&session_id, source, plugin_version).unwrap(); + recovered.configure(&session_config(&base)); + recovered + .emit(&[SpanOp::Merge(SpanRow { + output: Some(json!({"status":"recovered"})), + ..late + })]) + .await + .unwrap(); + recovered.flush().await.unwrap(); - let bodies = logs3_bodies(&server).await; - assert!( - bodies.contains("original name"), - "initial row absent: {bodies}" - ); - let rows = logs3_rows(&server).await; - let late = rows - .iter() - .find(|row| row.pointer("/output/status") == Some(&json!("late"))) - .unwrap_or_else(|| panic!("late merge absent: {bodies}")); - assert_eq!(late["_is_merge"], true); - assert_eq!(late["root_span_id"], "trace-root"); - assert_eq!( - late["span_parents"], - json!(["turn-parent"]), - "stateless merge did not repeat the child parent identity: {bodies}" - ); + let rows: Vec<_> = logs3_rows(&server) + .await + .into_iter() + .filter(|row| row["span_id"] == span_id) + .collect(); + assert!(rows + .iter() + .any(|row| row.pointer("/span_attributes/name") == Some(&json!("original name")))); + for status in ["late", "recovered"] { + let update = rows + .iter() + .find(|row| row.pointer("/output/status") == Some(&json!(status))) + .unwrap_or_else(|| panic!("{source}: {status} merge absent: {rows:?}")); + assert_eq!(update["_is_merge"], true); + assert_eq!(update["root_span_id"], "trace-root"); + assert_eq!(update["span_parents"], json!(["turn-parent"])); + } + for row in rows { + let origin = &row["context"]["span_origin"]; + assert_eq!(origin["name"], format!("braintrust.plugin.{source}")); + assert_eq!(origin["version"], plugin_version.unwrap_or("test")); + assert_eq!(origin["instrumentation"]["name"], "braintrust-plugin"); + assert_eq!(row["metadata"]["bt_daemon_version"], "test"); + } + } } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/src/plugins/pi/content/src/index.ts b/src/plugins/pi/content/src/index.ts index 98b4f6d..75feff5 100644 --- a/src/plugins/pi/content/src/index.ts +++ b/src/plugins/pi/content/src/index.ts @@ -89,7 +89,9 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { requestTimeoutMs: UI_STATUS_TIMEOUT_MS, }); - const remember = async (ctx: ExtensionContext): Promise> => { + const updateSession = async ( + ctx: ExtensionContext, + ): Promise> => { lastContext = ctx; const descriptor = sessionDescriptor(ctx); if (sessionId !== descriptor.sessionId) uiGeneration += 1; @@ -132,7 +134,7 @@ export default function braintrustPiExtension(pi: ExtensionAPI): void { ctx?: ExtensionContext, updateUi = false, ): Promise => { - const descriptor = ctx ? await remember(ctx) : undefined; + const descriptor = ctx ? await updateSession(ctx) : undefined; if (!sessionId) return; await client.log({ source: "pi",