diff --git a/.gitignore b/.gitignore index d77f6bd..659c91a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,5 @@ # Build output dist/ + +# Local credentials and environment overrides +.env diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 627f79e..5c4d856 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -75,6 +75,11 @@ pub struct HookArgs { /// Explicit event name (overrides `--event-field` lookup). #[arg(long)] pub event: Option, + /// JSON field holding a transcript path. When present, capture the file + /// length observed by this hook so deterministic journal replay cannot + /// read transcript records written by later lifecycle events. + #[arg(long)] + pub transcript_path_field: Option, /// Fail instead of spawning a daemon if none is running. #[arg(long)] pub no_spawn: bool, @@ -190,7 +195,11 @@ pub async fn run_hook( if !settings.tracing_enabled() { return Ok(()); } - let payload = read_stdin_json()?; + let mut payload = read_stdin_json()?; + + if let Some(field) = &args.transcript_path_field { + add_transcript_observation(&mut payload, field); + } let session_id = json_str_field(&payload, &args.session_id_field) .ok_or_else(|| anyhow::anyhow!("no `{}` field in hook payload", args.session_id_field))?; @@ -771,6 +780,47 @@ fn json_str_field(payload: &serde_json::Value, field: &str) -> Option { } } +/// Stamp the transcript boundary visible when a blocking hook runs. Agent +/// transcripts are append-only, while daemon journal replay may happen after +/// the session has advanced. Recording byte lengths keeps translation causally +/// aligned with each native hook without copying transcript contents into the +/// journal. +fn add_transcript_observation(payload: &mut serde_json::Value, field: &str) { + let Some(path) = json_str_field(payload, field) else { + return; + }; + let transcript = std::path::Path::new(&path); + let mut observation = serde_json::Map::new(); + observation.insert("path".into(), serde_json::Value::String(path.clone())); + if let Ok(metadata) = std::fs::metadata(transcript) { + observation.insert( + "observed_bytes".into(), + serde_json::Value::Number(metadata.len().into()), + ); + } + + if transcript.file_name().and_then(|name| name.to_str()) == Some("transcript.jsonl") { + let full = transcript.with_file_name("transcript_full.jsonl"); + if let Ok(metadata) = std::fs::metadata(&full) { + observation.insert( + "full_path".into(), + serde_json::Value::String(full.to_string_lossy().into_owned()), + ); + observation.insert( + "full_observed_bytes".into(), + serde_json::Value::Number(metadata.len().into()), + ); + } + } + + if let Some(object) = payload.as_object_mut() { + object.insert( + "_bt_transcript_observation".into(), + serde_json::Value::Object(observation), + ); + } +} + #[cfg(test)] mod tests { use super::*; @@ -798,6 +848,26 @@ mod tests { assert!(now_ms() > 0); } + #[test] + fn transcript_observation_captures_compact_and_full_boundaries() { + let dir = tempfile::tempdir().unwrap(); + let compact = dir.path().join("transcript.jsonl"); + let full = dir.path().join("transcript_full.jsonl"); + std::fs::write(&compact, b"compact\n").unwrap(); + std::fs::write(&full, b"complete record\n").unwrap(); + let mut payload = serde_json::json!({ + "transcriptPath": compact.to_string_lossy() + }); + + add_transcript_observation(&mut payload, "transcriptPath"); + + let observed = &payload["_bt_transcript_observation"]; + assert_eq!(observed["path"], compact.to_string_lossy().as_ref()); + assert_eq!(observed["observed_bytes"], 8); + assert_eq!(observed["full_path"], full.to_string_lossy().as_ref()); + assert_eq!(observed["full_observed_bytes"], 16); + } + #[test] fn import_destination_without_session_config_fails_fast() { let mut config = None; diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 97a0a38..72e6a68 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -209,7 +209,6 @@ impl BraintrustSink { .span_type(map_span_type(row.span_type)) .span_id(row.span_id.clone()) .row_id(row.span_id.clone()) - .project_name(project) .parent_info(parent) .span_origin( SpanOrigin::new() @@ -231,9 +230,6 @@ impl BraintrustSink { self.ensure_handle(client, row)?; let handle = self.open.get(&row.span_id).expect("just inserted"); handle.log(build_log(row)?); - if let Some(end) = row.end_ms { - handle.end_with_time(ms_to_secs(end)); - } Ok(()) } } @@ -455,14 +451,22 @@ fn build_log(row: &SpanRow) -> anyhow::Result { if let Some(Value::Object(md)) = &row.metadata { lb = lb.metadata(md.clone()); } - if let Some(Value::Object(metrics)) = &row.metrics { - let hm: HashMap = metrics + let mut metrics = match &row.metrics { + Some(Value::Object(metrics)) => metrics .iter() .filter_map(|(k, v)| v.as_f64().map(|f| (k.clone(), f))) - .collect(); - if !hm.is_empty() { - lb = lb.metrics(hm); - } + .collect(), + _ => HashMap::new(), + }; + if let Some(end) = row.end_ms { + // `SpanHandle::end_with_time` intentionally keeps the first end time. + // Coding-agent sessions can be resumed after an idle stop, so encode + // end as a regular mergeable metric and allow a later lifecycle event + // to extend the same deterministic span. + metrics.insert("end".to_string(), ms_to_secs(end)); + } + if !metrics.is_empty() { + lb = lb.metrics(metrics); } if let Some(err) = &row.error { lb = lb.error(Value::String(err.clone())); diff --git a/bt-daemon/src/translate/antigravity.rs b/bt-daemon/src/translate/antigravity.rs new file mode 100644 index 0000000..c5ec613 --- /dev/null +++ b/bt-daemon/src/translate/antigravity.rs @@ -0,0 +1,859 @@ +//! Google Antigravity hook and transcript translator. +//! +//! Native hooks own lifecycle timing and correlation (`invocationNum` and +//! `stepIdx`). The append-only full transcript supplies the actual user/model +//! messages and tool details. Hook capture records transcript byte boundaries, +//! so replay observes exactly the records that existed when each hook fired. + +use super::git::GitMetadataCache; +use super::{AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory}; +use crate::ids; +use crate::wire::Envelope; +use serde_json::{json, Map, Value}; +use std::collections::HashMap; +use std::io::{BufRead, Seek, SeekFrom}; +use std::path::Path; +use std::sync::Arc; + +pub struct AntigravityTranslatorFactory { + git: Arc, +} + +impl AntigravityTranslatorFactory { + pub(super) fn new(git: Arc) -> Self { + Self { git } + } +} + +impl TranslatorFactory for AntigravityTranslatorFactory { + fn source(&self) -> &str { + "antigravity" + } + + fn create(&self, session_id: &str) -> Box { + Box::new(AntigravityTranslator::new(session_id, self.git.clone())) + } +} + +struct Turn { + span_id: String, + number: u32, + start_ms: i64, + last_output: Option, +} + +struct Invocation { + span_id: String, + parent_span_id: String, + history_start: usize, + record_start: usize, +} + +struct PendingTool { + span_id: String, + parent_span_id: String, + name: String, +} + +struct AntigravityTranslator { + session_id: String, + session_span_id: String, + root_span_id: String, + root_open: bool, + root_ended: bool, + turn: Option, + turn_count: u32, + transcript_offsets: HashMap, + records: Vec, + records_by_step: HashMap, + history: Vec, + invocations: HashMap, + tools: HashMap, + last_ts_ms: i64, + git: Arc, +} + +impl AntigravityTranslator { + fn new(session_id: &str, git: Arc) -> Self { + let root = ids::span_id(session_id, "root"); + Self { + session_id: session_id.to_string(), + session_span_id: root.clone(), + root_span_id: root, + root_open: false, + root_ended: false, + turn: None, + turn_count: 0, + transcript_offsets: HashMap::new(), + records: Vec::new(), + records_by_step: HashMap::new(), + history: Vec::new(), + invocations: HashMap::new(), + tools: HashMap::new(), + last_ts_ms: 0, + git, + } + } + + fn ensure_root(&mut self, event: &Envelope, ctx: &SessionCtx, ops: &mut Vec) { + if self.root_open { + return; + } + self.root_open = true; + let (parent_span_id, external_root_span_id) = ctx + .config + .as_ref() + .map(|config| config.attached_span_ids()) + .unwrap_or_default(); + if let Some(external_root) = external_root_span_id { + self.root_span_id = external_root; + } + + let workspace = event + .payload + .get("workspacePaths") + .and_then(Value::as_array) + .and_then(|paths| paths.first()) + .and_then(Value::as_str); + let mut metadata = ctx + .config + .as_ref() + .and_then(|config| config.additional_metadata.clone()) + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + metadata.retain(|key, _| !key.starts_with("_bt_")); + metadata.insert("session_id".into(), json!(self.session_id)); + metadata.insert("conversation_id".into(), json!(self.session_id)); + metadata.insert("source".into(), json!("antigravity")); + if let Some(model) = string_field(&event.payload, "modelName") { + metadata.insert("model".into(), json!(model)); + } + if let Some(workspaces) = event.payload.get("workspacePaths") { + metadata.insert("workspace_paths".into(), workspaces.clone()); + } + if let Some(path) = string_field(&event.payload, "artifactDirectoryPath") { + metadata.insert("artifact_directory_path".into(), json!(path)); + } + if let Some(version) = &event.source_version { + metadata.insert("antigravity_version".into(), json!(version)); + } + + let label = workspace + .and_then(|path| Path::new(path).file_name()) + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .unwrap_or("session"); + ops.push(SpanOp::Insert(SpanRow { + span_id: self.session_span_id.clone(), + root_span_id: self.root_span_id.clone(), + parent_span_ids: parent_span_id.into_iter().collect(), + name: format!("Antigravity: {label}"), + span_type: SpanType::Task, + start_ms: Some(event.ts_ms), + metadata: Some(Value::Object(metadata)), + ..Default::default() + })); + } + + fn tail_transcript(&mut self, event: &Envelope, ops: &mut Vec) { + let Some(source) = transcript_source(event) else { + return; + }; + let offset = self + .transcript_offsets + .entry(source.path.clone()) + .or_default(); + let records = if let Some(contents) = source.snapshot { + read_snapshot_records(contents, offset, source.through) + } else { + read_file_records(&source.path, offset, source.through) + }; + for record in records { + self.observe_record(record, event.ts_ms, ops); + } + } + + fn observe_record(&mut self, record: Value, ts_ms: i64, ops: &mut Vec) { + let record_type = normalized_record_type(&record); + let source = string_field(&record, "source").unwrap_or_default(); + if let Some(step) = + integer_field(&record, "step_index").or_else(|| integer_field(&record, "stepIndex")) + { + self.records_by_step.insert(step, record.clone()); + } + + if record_type == "USER_INPUT" + && (source.is_empty() || source == "USER_EXPLICIT" || source == "USER") + { + self.start_turn(clean_user_input(record_content(&record)), ts_ms, ops); + } + + if let Some(message) = transcript_message(&record, &record_type, &source) { + if record_type == "PLANNER_RESPONSE" { + if let Some(turn) = &mut self.turn { + turn.last_output = message.get("content").cloned(); + } + } + self.history.push(message); + } + self.records.push(record); + } + + fn start_turn(&mut self, input: Value, ts_ms: i64, ops: &mut Vec) { + self.close_turn(ts_ms, None, ops); + self.turn_count += 1; + let span_id = ids::span_id(&self.session_id, &format!("turn:{}", self.turn_count)); + ops.push(SpanOp::Insert(SpanRow { + span_id: span_id.clone(), + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![self.session_span_id.clone()], + name: format!("Turn {}", self.turn_count), + span_type: SpanType::Task, + start_ms: Some(ts_ms), + input: nonempty_value(input), + metadata: Some(json!({"turn_number": self.turn_count})), + ..Default::default() + })); + self.turn = Some(Turn { + span_id, + number: self.turn_count, + start_ms: ts_ms, + last_output: None, + }); + } + + fn ensure_turn(&mut self, ts_ms: i64, ops: &mut Vec) -> String { + if self.turn.is_none() { + self.start_turn(Value::Null, ts_ms, ops); + } + self.turn + .as_ref() + .map(|turn| turn.span_id.clone()) + .unwrap_or_else(|| self.session_span_id.clone()) + } + + fn pre_invocation(&mut self, event: &Envelope, ops: &mut Vec) { + let Some(invocation_num) = integer_field(&event.payload, "invocationNum") else { + return; + }; + if self.invocations.contains_key(&invocation_num) { + return; + } + let parent = self.ensure_turn(event.ts_ms, ops); + // Antigravity's invocation counter is process-local and resets to zero + // when `--conversation` resumes an existing conversation. Include the + // stable transcript-derived turn number so resumed invocations do not + // merge into an earlier turn's LLM span. + let turn_number = self.turn.as_ref().map(|turn| turn.number).unwrap_or(0); + let span_id = ids::span_id( + &self.session_id, + &format!("turn:{turn_number}:invocation:{invocation_num}"), + ); + let model = string_field(&event.payload, "modelName") + .unwrap_or_else(|| "Antigravity model".to_string()); + let mut metadata = json!({ + "invocation_num": invocation_num, + "initial_num_steps": integer_field(&event.payload, "initialNumSteps"), + "turn_number": turn_number, + "model": model + }); + remove_null_fields(&mut metadata); + ops.push(SpanOp::Insert(SpanRow { + span_id: span_id.clone(), + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![parent.clone()], + name: model, + span_type: SpanType::Llm, + start_ms: Some(event.ts_ms), + input: Some(Value::Array(self.history.clone())), + metadata: Some(metadata), + ..Default::default() + })); + self.invocations.insert( + invocation_num, + Invocation { + span_id, + parent_span_id: parent, + history_start: self.history.len(), + record_start: self.records.len(), + }, + ); + } + + fn post_invocation(&mut self, event: &Envelope, ops: &mut Vec) { + let Some(invocation_num) = integer_field(&event.payload, "invocationNum") else { + return; + }; + let Some(invocation) = self.invocations.remove(&invocation_num) else { + return; + }; + let output = self.history[invocation.history_start.min(self.history.len())..] + .iter() + .filter(|message| message.get("role").and_then(Value::as_str) == Some("assistant")) + .cloned() + .collect::>(); + let metrics = + token_metrics(&self.records[invocation.record_start.min(self.records.len())..]); + ops.push(SpanOp::Merge(SpanRow { + span_id: invocation.span_id, + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![invocation.parent_span_id], + end_ms: Some(event.ts_ms), + output: (!output.is_empty()).then_some(Value::Array(output)), + metrics, + ..Default::default() + })); + } + + fn pre_tool(&mut self, event: &Envelope, ops: &mut Vec) { + let Some(step) = integer_field(&event.payload, "stepIdx") else { + return; + }; + if self.tools.contains_key(&step) { + return; + } + let parent = self.ensure_turn(event.ts_ms, ops); + let call = event.payload.get("toolCall").unwrap_or(&Value::Null); + let name = string_field(call, "name").unwrap_or_else(|| format!("Tool step {step}")); + let input = call.get("args").cloned(); + let span_id = ids::span_id(&self.session_id, &format!("tool:{step}")); + ops.push(SpanOp::Insert(SpanRow { + span_id: span_id.clone(), + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![parent.clone()], + name: name.clone(), + span_type: SpanType::Tool, + start_ms: Some(event.ts_ms), + input, + metadata: Some(json!({"step_index": step, "tool_name": name})), + ..Default::default() + })); + self.tools.insert( + step, + PendingTool { + span_id, + parent_span_id: parent, + name, + }, + ); + } + + fn post_tool(&mut self, event: &Envelope, ops: &mut Vec) { + let Some(step) = integer_field(&event.payload, "stepIdx") else { + return; + }; + let transcript = self.records_by_step.get(&step); + let mut details = transcript.map(tool_details).unwrap_or_default(); + if let Some(call) = event.payload.get("toolCall") { + if let Some(name) = string_field(call, "name") { + details.name = Some(name); + } + if let Some(input) = call.get("args") { + details.input = Some(input.clone()); + } + } + let recovered_start_ms = planned_tool_start(&self.records, details.name.as_deref(), step) + .map(|start| { + self.turn + .as_ref() + .map(|turn| start.max(turn.start_ms)) + .unwrap_or(start) + }); + let pending = self.tools.remove(&step); + let was_pending = pending.is_some(); + let parent = pending + .as_ref() + .map(|tool| tool.parent_span_id.clone()) + .unwrap_or_else(|| self.ensure_turn(event.ts_ms, ops)); + let name = pending + .as_ref() + .map(|tool| tool.name.clone()) + .or(details.name) + .unwrap_or_else(|| format!("Tool step {step}")); + let span_id = pending + .map(|tool| tool.span_id) + .unwrap_or_else(|| ids::span_id(&self.session_id, &format!("tool:{step}"))); + let error = string_field(&event.payload, "error") + .filter(|error| !error.is_empty()) + .or(details.error); + let outcome = if error.is_some() { "error" } else { "success" }; + if !was_pending { + ops.push(SpanOp::Insert(SpanRow { + span_id: span_id.clone(), + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![parent.clone()], + name: name.clone(), + span_type: SpanType::Tool, + start_ms: Some(recovered_start_ms.unwrap_or(event.ts_ms)), + input: details.input.clone(), + metadata: Some(json!({ + "step_index": step, + "tool_name": name, + "recovered_from_transcript": true + })), + ..Default::default() + })); + } + ops.push(SpanOp::Merge(SpanRow { + span_id, + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![parent], + name, + span_type: SpanType::Tool, + end_ms: Some(event.ts_ms), + input: details.input, + output: details.output, + metadata: Some(json!({ + "step_index": step, + "tool_outcome": outcome + })), + error, + ..Default::default() + })); + } + + fn stop(&mut self, event: &Envelope, ops: &mut Vec) { + let error = string_field(&event.payload, "error").filter(|error| !error.is_empty()); + self.close_pending(event.ts_ms, error.clone(), ops); + self.close_turn(event.ts_ms, error.clone(), ops); + if event + .payload + .get("fullyIdle") + .and_then(Value::as_bool) + .unwrap_or(true) + { + self.close_root(event.ts_ms, error, ops); + } + } + + fn close_pending(&mut self, ts_ms: i64, error: Option, ops: &mut Vec) { + for (_, invocation) in self.invocations.drain() { + ops.push(SpanOp::Merge(SpanRow { + span_id: invocation.span_id, + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![invocation.parent_span_id], + end_ms: Some(ts_ms), + output: self.turn.as_ref().and_then(|turn| { + turn.last_output + .clone() + .map(|content| json!([{"role":"assistant","content":content}])) + }), + error: error.clone(), + ..Default::default() + })); + } + for (step, tool) in self.tools.drain() { + ops.push(SpanOp::Merge(SpanRow { + span_id: tool.span_id, + root_span_id: self.root_span_id.clone(), + parent_span_ids: vec![tool.parent_span_id], + name: tool.name, + span_type: SpanType::Tool, + end_ms: Some(ts_ms), + metadata: Some(json!({"step_index":step,"tool_outcome":"unknown"})), + error: error.clone(), + ..Default::default() + })); + } + } + + fn close_turn(&mut self, ts_ms: i64, error: Option, ops: &mut Vec) { + if let Some(turn) = self.turn.take() { + ops.push(SpanOp::Merge(SpanRow { + span_id: turn.span_id, + root_span_id: self.root_span_id.clone(), + end_ms: Some(ts_ms), + output: turn.last_output, + metadata: Some(json!({"turn_number":turn.number})), + error, + ..Default::default() + })); + } + } + + fn close_root(&mut self, ts_ms: i64, error: Option, ops: &mut Vec) { + if self.root_ended { + return; + } + self.root_ended = true; + ops.push(SpanOp::Merge(SpanRow { + span_id: self.session_span_id.clone(), + root_span_id: self.root_span_id.clone(), + end_ms: Some(ts_ms), + error, + ..Default::default() + })); + } +} + +impl AgentTranslator for AntigravityTranslator { + fn handle(&mut self, event: &Envelope, ctx: &SessionCtx) -> anyhow::Result> { + self.last_ts_ms = self.last_ts_ms.max(event.ts_ms); + let mut ops = Vec::new(); + // A later Antigravity process can resume the same conversation after a + // fully-idle Stop. Reopen the logical root so the resumed Stop extends + // its duration through all subsequent turns. + if self.root_ended && event.event != "Stop" { + self.root_ended = false; + } + self.ensure_root(event, ctx, &mut ops); + self.tail_transcript(event, &mut ops); + match event.event.as_str() { + "PreInvocation" => self.pre_invocation(event, &mut ops), + "PostInvocation" => self.post_invocation(event, &mut ops), + "PreToolUse" => self.pre_tool(event, &mut ops), + "PostToolUse" => self.post_tool(event, &mut ops), + "Stop" => self.stop(event, &mut ops), + _ => {} + } + let cwd = event + .payload + .get("workspacePaths") + .and_then(Value::as_array) + .and_then(|paths| paths.first()) + .and_then(Value::as_str); + self.git.enrich_rows(cwd, &mut ops); + Ok(ops) + } + + fn flush(&mut self, _ctx: &SessionCtx) -> anyhow::Result> { + let mut ops = Vec::new(); + self.close_pending(self.last_ts_ms, None, &mut ops); + self.close_turn(self.last_ts_ms, None, &mut ops); + self.close_root(self.last_ts_ms, None, &mut ops); + Ok(ops) + } +} + +struct TranscriptSource<'a> { + path: String, + through: u64, + snapshot: Option<&'a str>, +} + +fn transcript_source(event: &Envelope) -> Option> { + let observation = event.payload.get("_bt_transcript_observation"); + let compact_path = observation + .and_then(|value| value.get("path")) + .and_then(Value::as_str) + .or_else(|| event.payload.get("transcriptPath").and_then(Value::as_str))?; + let full_path = observation + .and_then(|value| value.get("full_path")) + .and_then(Value::as_str); + let (path, through) = if let (Some(path), Some(through)) = ( + full_path, + observation + .and_then(|value| value.get("full_observed_bytes")) + .and_then(Value::as_u64), + ) { + (path, through) + } else { + let through = observation + .and_then(|value| value.get("observed_bytes")) + .and_then(Value::as_u64) + .or_else(|| { + std::fs::metadata(compact_path) + .ok() + .map(|metadata| metadata.len()) + })?; + (compact_path, through) + }; + let snapshot = event + .payload + .get("_bt_transcript_snapshot") + .filter(|value| value.get("path").and_then(Value::as_str) == Some(path)) + .and_then(|value| value.get("contents")) + .and_then(Value::as_str); + Some(TranscriptSource { + path: path.to_string(), + through, + snapshot, + }) +} + +fn read_file_records(path: &str, offset: &mut u64, through: u64) -> Vec { + let Ok(mut file) = std::fs::File::open(path) else { + return Vec::new(); + }; + let len = file.metadata().map(|metadata| metadata.len()).unwrap_or(0); + if *offset > len { + *offset = 0; + } + if file.seek(SeekFrom::Start(*offset)).is_err() { + return Vec::new(); + } + read_records(&mut std::io::BufReader::new(file), offset, through.min(len)) +} + +fn read_snapshot_records(contents: &str, offset: &mut u64, through: u64) -> Vec { + if *offset > contents.len() as u64 { + *offset = 0; + } + let mut reader = std::io::BufReader::new(std::io::Cursor::new(contents.as_bytes())); + if reader.seek(SeekFrom::Start(*offset)).is_err() { + return Vec::new(); + } + read_records(&mut reader, offset, through.min(contents.len() as u64)) +} + +fn read_records( + reader: &mut std::io::BufReader, + offset: &mut u64, + through: u64, +) -> Vec { + let mut records = Vec::new(); + let mut line = String::new(); + while *offset < through { + line.clear(); + let start = *offset; + let Ok(read) = reader.read_line(&mut line) else { + break; + }; + if read == 0 || start + read as u64 > through { + break; + } + *offset += read as u64; + if let Ok(value) = serde_json::from_str::(line.trim()) { + records.push(value); + } + } + records +} + +fn normalized_record_type(record: &Value) -> String { + let record_type = string_field(record, "type").unwrap_or_default(); + record_type + .strip_prefix("CORTEX_STEP_TYPE_") + .unwrap_or(&record_type) + .to_string() +} + +fn transcript_message(record: &Value, record_type: &str, source: &str) -> Option { + let role = match (record_type, source) { + ("USER_INPUT", _) | (_, "USER_EXPLICIT") | (_, "USER") => "user", + ("PLANNER_RESPONSE", _) => "assistant", + (_, "SYSTEM") => "system", + _ if record.get("content").is_some() => "tool", + _ => return None, + }; + let mut message = Map::new(); + let content = match role { + "user" => clean_user_input(record_content(record)), + "tool" => clean_tool_content(record_content(record)), + _ => record_content(record), + }; + if role == "system" && content.is_null() { + return None; + } + message.insert("role".into(), json!(role)); + message.insert("content".into(), content); + if let Some(tool_calls) = record.get("tool_calls").or_else(|| record.get("toolCalls")) { + message.insert("tool_calls".into(), tool_calls.clone()); + } + if role == "tool" { + message.insert("name".into(), json!(record_type.to_ascii_lowercase())); + } + message.insert("step_type".into(), json!(record_type)); + Some(Value::Object(message)) +} + +fn record_content(record: &Value) -> Value { + record + .get("content") + .or_else(|| record.get("message")) + .or_else(|| record.get("text")) + .cloned() + .unwrap_or(Value::Null) +} + +fn clean_user_input(content: Value) -> Value { + let Value::String(text) = content else { + return content; + }; + let Some(start) = text.find("") else { + return Value::String(text); + }; + let body_start = start + "".len(); + let Some(relative_end) = text[body_start..].find("") else { + return Value::String(text); + }; + Value::String( + text[body_start..body_start + relative_end] + .trim() + .to_string(), + ) +} + +#[derive(Default)] +struct ToolDetails { + name: Option, + input: Option, + output: Option, + error: Option, +} + +fn tool_details(record: &Value) -> ToolDetails { + let call = record + .get("tool_calls") + .or_else(|| record.get("toolCalls")) + .and_then(|calls| { + calls + .as_array() + .and_then(|calls| calls.first()) + .or(Some(calls)) + }); + let record_type = normalized_record_type(record); + let name = call + .and_then(|call| { + string_field(call, "name") + .or_else(|| string_field(call, "tool_name")) + .or_else(|| string_field(call, "toolName")) + }) + .or_else(|| { + (!matches!(record_type.as_str(), "USER_INPUT" | "PLANNER_RESPONSE")) + .then(|| record_type.to_ascii_lowercase()) + }); + let input = call.and_then(|call| { + call.get("args") + .or_else(|| call.get("arguments")) + .or_else(|| call.get("input")) + .cloned() + }); + let output = call + .and_then(|call| call.get("output").or_else(|| call.get("result"))) + .cloned() + .or_else(|| nonempty_value(clean_tool_content(record_content(record)))); + let error = string_field(record, "error") + .filter(|error| !error.is_empty()) + .or_else(|| { + (record_type == "ERROR_MESSAGE") + .then(|| record_content(record)) + .and_then(|content| content.as_str().map(str::to_owned)) + }) + .or_else(|| { + string_field(record, "status") + .filter(|status| matches!(status.to_ascii_lowercase().as_str(), "error" | "failed")) + }); + ToolDetails { + name, + input, + output, + error, + } +} + +fn clean_tool_content(content: Value) -> Value { + let Value::String(text) = content else { + return content; + }; + let mut lines = text.lines(); + let first = lines.next(); + let second = lines.next(); + if first.is_some_and(|line| line.starts_with("Created At:")) + && second.is_some_and(|line| line.starts_with("Completed At:")) + { + return Value::String(lines.collect::>().join("\n")); + } + Value::String(text) +} + +fn planned_tool_start(records: &[Value], name: Option<&str>, step: i64) -> Option { + let name = name?; + records.iter().rev().find_map(|record| { + let record_step = + integer_field(record, "step_index").or_else(|| integer_field(record, "stepIndex"))?; + if record_step >= step { + return None; + } + let calls = record + .get("tool_calls") + .or_else(|| record.get("toolCalls"))? + .as_array()?; + calls + .iter() + .any(|call| string_field(call, "name").as_deref() == Some(name)) + .then(|| parse_created_at(record)) + .flatten() + }) +} + +fn parse_created_at(record: &Value) -> Option { + let timestamp = + string_field(record, "created_at").or_else(|| string_field(record, "createdAt"))?; + chrono::DateTime::parse_from_rfc3339(×tamp) + .ok() + .map(|timestamp| timestamp.timestamp_millis()) +} + +fn token_metrics(records: &[Value]) -> Option { + let mut found = HashMap::<&'static str, f64>::new(); + for record in records { + collect_token_metrics(record, &mut found); + } + if found.is_empty() { + return None; + } + let mut metrics = Map::new(); + for (key, value) in found { + metrics.insert(key.to_string(), json!(value)); + } + Some(Value::Object(metrics)) +} + +fn collect_token_metrics(value: &Value, found: &mut HashMap<&'static str, f64>) { + match value { + Value::Object(object) => { + for (key, value) in object { + if let Some(number) = value.as_f64() { + let metric = match key.as_str() { + "input_tokens" | "inputTokens" | "prompt_tokens" | "promptTokens" => { + Some("prompt_tokens") + } + "output_tokens" | "outputTokens" | "completion_tokens" + | "completionTokens" => Some("completion_tokens"), + "total_tokens" | "totalTokens" => Some("tokens"), + "thinking_tokens" | "thinkingTokens" => Some("thinking_tokens"), + "cache_read_tokens" | "cacheReadTokens" => Some("prompt_cached_tokens"), + _ => None, + }; + if let Some(metric) = metric { + found.insert(metric, number); + } + } + collect_token_metrics(value, found); + } + } + Value::Array(values) => { + for value in values { + collect_token_metrics(value, found); + } + } + _ => {} + } +} + +fn string_field(value: &Value, field: &str) -> Option { + value.get(field).and_then(Value::as_str).map(str::to_owned) +} + +fn integer_field(value: &Value, field: &str) -> Option { + value + .get(field) + .and_then(|value| value.as_i64().or_else(|| value.as_str()?.parse().ok())) +} + +fn nonempty_value(value: Value) -> Option { + match &value { + Value::Null => None, + Value::String(text) if text.is_empty() => None, + Value::Array(values) if values.is_empty() => None, + Value::Object(object) if object.is_empty() => None, + _ => Some(value), + } +} + +fn remove_null_fields(value: &mut Value) { + if let Some(object) = value.as_object_mut() { + object.retain(|_, value| !value.is_null()); + } +} diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index d182161..eb4a2ec 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -7,6 +7,7 @@ //! whole pipeline be exercised with a debug sink and makes translators unit- //! testable without any network. +mod antigravity; mod claude; mod codex; mod debug; @@ -14,6 +15,7 @@ mod git; mod opencode; mod pi; +pub use antigravity::AntigravityTranslatorFactory; pub use claude::ClaudeTranslatorFactory; pub use codex::CodexTranslatorFactory; pub use debug::DebugTranslatorFactory; @@ -125,6 +127,7 @@ impl Registry { pub fn default_agents() -> Self { let mut r = Registry::debug_only(); let git = Arc::new(git::GitMetadataCache::default()); + r.register(Box::new(AntigravityTranslatorFactory::new(git.clone()))); r.register(Box::new(ClaudeTranslatorFactory::new(git.clone()))); r.register(Box::new(CodexTranslatorFactory::new(git.clone()))); r.register(Box::new(OpenCodeTranslatorFactory::new(git.clone()))); diff --git a/bt-daemon/tests/antigravity_translator.rs b/bt-daemon/tests/antigravity_translator.rs new file mode 100644 index 0000000..0c91c1c --- /dev/null +++ b/bt-daemon/tests/antigravity_translator.rs @@ -0,0 +1,692 @@ +use bt_daemon::wire::Envelope; +use bt_daemon::{Registry, SessionCtx, SpanOp, SpanRow, SpanType}; +use serde_json::{json, Value}; +use std::collections::HashMap; + +fn jsonl(records: &[Value]) -> (String, Vec) { + let mut contents = String::new(); + let mut boundaries = Vec::new(); + for record in records { + contents.push_str(&serde_json::to_string(record).unwrap()); + contents.push('\n'); + boundaries.push(contents.len() as u64); + } + (contents, boundaries) +} + +fn event( + name: &str, + ts_ms: i64, + transcript_path: &str, + transcript: &str, + through: u64, + extra: Value, +) -> Envelope { + let full_path = transcript_path.replace("transcript.jsonl", "transcript_full.jsonl"); + let mut payload = json!({ + "conversationId": "conversation-1", + "workspacePaths": ["/workspace/demo"], + "transcriptPath": transcript_path, + "artifactDirectoryPath": "/tmp/artifacts", + "modelName": "gemini-3.1-pro", + "_bt_transcript_observation": { + "path": transcript_path, + "observed_bytes": through, + "full_path": full_path, + "full_observed_bytes": through + }, + "_bt_transcript_snapshot": { + "path": full_path, + "contents": transcript + } + }); + if let (Value::Object(payload), Value::Object(extra)) = (&mut payload, extra) { + payload.extend(extra); + } + Envelope { + source: "antigravity".into(), + source_version: Some("1.1.12".into()), + plugin_version: None, + session_id: "conversation-1".into(), + event: name.into(), + ts_ms, + payload, + route: None, + config: None, + } +} + +fn reduce(ops: Vec) -> HashMap { + let mut rows: HashMap = HashMap::new(); + for op in ops { + match op { + SpanOp::Insert(row) => { + rows.insert(row.span_id.clone(), row); + } + SpanOp::Merge(row) => { + let existing = rows.entry(row.span_id.clone()).or_default(); + if !row.root_span_id.is_empty() { + existing.root_span_id = row.root_span_id; + } + if !row.parent_span_ids.is_empty() { + existing.parent_span_ids = row.parent_span_ids; + } + if !row.name.is_empty() { + existing.name = row.name; + } + if row.start_ms.is_some() { + existing.start_ms = row.start_ms; + } + if row.end_ms.is_some() { + existing.end_ms = row.end_ms; + } + if row.input.is_some() { + existing.input = row.input; + } + if row.output.is_some() { + existing.output = row.output; + } + if row.metrics.is_some() { + existing.metrics = row.metrics; + } + if row.error.is_some() { + existing.error = row.error; + } + if let Some(Value::Object(incoming)) = row.metadata { + let mut metadata = existing + .metadata + .take() + .and_then(|value| value.as_object().cloned()) + .unwrap_or_default(); + metadata.extend(incoming); + existing.metadata = Some(Value::Object(metadata)); + } + } + } + } + rows +} + +#[test] +fn hooks_and_full_transcript_build_model_and_tool_spans() { + let records = vec![ + json!({ + "step_index": 0, + "source": "USER_EXPLICIT", + "type": "USER_INPUT", + "status": "COMPLETED", + "content": "List the files" + }), + json!({ + "step_index": 1, + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "COMPLETED", + "content": "I'll inspect the directory.", + "tool_calls": [{"name":"run_command","args":{"CommandLine":"ls"}}], + "usage": {"inputTokens": 12, "outputTokens": 7, "totalTokens": 19} + }), + json!({ + "step_index": 2, + "source": "MODEL", + "type": "CORTEX_STEP_TYPE_RUN_COMMAND", + "status": "COMPLETED", + "content": "README.md\nsrc", + "tool_calls": [{ + "name": "run_command", + "args": {"CommandLine":"ls"}, + "result": "README.md\nsrc" + }] + }), + ]; + let (transcript, boundary) = jsonl(&records); + let path = "/tmp/conversation/transcript.jsonl"; + let registry = Registry::default_agents(); + let mut translator = registry.create("antigravity", "conversation-1"); + let ctx = SessionCtx { + session_id: "conversation-1".into(), + config: None, + }; + let mut ops = Vec::new(); + ops.extend( + translator + .handle( + &event( + "PreInvocation", + 100, + path, + &transcript, + boundary[0], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "PostInvocation", + 200, + path, + &transcript, + boundary[1], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + &ctx, + ) + .unwrap(), + ); + // The safe default plugin uses PostToolUse only: transcript stepIdx + // recovery provides the tool name, arguments, and output without changing + // Antigravity's permission behavior via a PreToolUse decision. + ops.extend( + translator + .handle( + &event( + "PostToolUse", + 300, + path, + &transcript, + boundary[2], + json!({"stepIdx":2,"error":""}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "Stop", + 400, + path, + &transcript, + boundary[2], + json!({ + "executionNum": 0, + "terminationReason": "model_stop", + "error": "", + "fullyIdle": true + }), + ), + &ctx, + ) + .unwrap(), + ); + + let rows = reduce(ops); + let root = rows + .values() + .find(|row| row.name == "Antigravity: demo") + .unwrap(); + assert_eq!(root.metadata.as_ref().unwrap()["source"], "antigravity"); + assert!(root.end_ms.is_some()); + + let turn = rows + .values() + .find(|row| row.span_type == SpanType::Task && row.name == "Turn 1") + .unwrap(); + assert_eq!(turn.input, Some(json!("List the files"))); + assert_eq!(turn.output, Some(json!("I'll inspect the directory."))); + assert_eq!(turn.parent_span_ids, vec![root.span_id.clone()]); + + let llm = rows + .values() + .find(|row| row.span_type == SpanType::Llm) + .unwrap(); + assert_eq!(llm.name, "gemini-3.1-pro"); + assert_eq!(llm.parent_span_ids, vec![turn.span_id.clone()]); + assert_eq!(llm.input.as_ref().unwrap()[0]["role"], "user"); + assert_eq!(llm.output.as_ref().unwrap()[0]["role"], "assistant"); + assert_eq!(llm.metrics.as_ref().unwrap()["prompt_tokens"], 12.0); + assert_eq!(llm.metrics.as_ref().unwrap()["completion_tokens"], 7.0); + assert_eq!(llm.metrics.as_ref().unwrap()["tokens"], 19.0); + + let tool = rows + .values() + .find(|row| row.span_type == SpanType::Tool) + .unwrap(); + assert_eq!(tool.name, "run_command"); + assert_eq!(tool.input.as_ref().unwrap()["CommandLine"], "ls"); + assert_eq!(tool.output, Some(json!("README.md\nsrc"))); + assert_eq!(tool.metadata.as_ref().unwrap()["tool_outcome"], "success"); + assert_eq!(tool.parent_span_ids, vec![turn.span_id.clone()]); +} + +#[test] +fn pre_tool_pair_preserves_start_time_and_reports_failure() { + let records = vec![json!({ + "step_index": 4, + "source": "MODEL", + "type": "CORTEX_STEP_TYPE_RUN_COMMAND", + "status": "FAILED", + "content": "permission denied" + })]; + let (transcript, boundary) = jsonl(&records); + let path = "/tmp/conversation/transcript.jsonl"; + let registry = Registry::default_agents(); + let mut translator = registry.create("antigravity", "conversation-2"); + let ctx = SessionCtx { + session_id: "conversation-2".into(), + config: None, + }; + let mut ops = translator + .handle( + &event( + "PreToolUse", + 10, + path, + &transcript, + 0, + json!({ + "stepIdx": 4, + "toolCall": {"name":"run_command","args":{"CommandLine":"secret"}} + }), + ), + &ctx, + ) + .unwrap(); + ops.extend( + translator + .handle( + &event( + "PostToolUse", + 20, + path, + &transcript, + boundary[0], + json!({"stepIdx":4,"error":"permission denied"}), + ), + &ctx, + ) + .unwrap(), + ); + let rows = reduce(ops); + let tool = rows + .values() + .find(|row| row.span_type == SpanType::Tool) + .unwrap(); + assert_eq!(tool.start_ms, Some(10)); + assert_eq!(tool.end_ms, Some(20)); + assert_eq!(tool.error.as_deref(), Some("permission denied")); + assert_eq!(tool.metadata.as_ref().unwrap()["tool_outcome"], "error"); +} + +#[test] +fn real_cli_schema_recovers_messages_and_post_only_tool() { + let records = vec![ + json!({ + "step_index": 0, + "created_at": "2026-08-11T17:46:20Z", + "source": "USER_EXPLICIT", + "type": "USER_INPUT", + "status": "DONE", + "content": "\nList /tmp\n\n\nignored\n" + }), + json!({ + "step_index": 1, + "created_at": "2026-08-11T17:46:20Z", + "source": "SYSTEM", + "type": "CONVERSATION_HISTORY", + "status": "DONE" + }), + json!({ + "step_index": 2, + "created_at": "2026-08-11T17:46:20Z", + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "DONE", + "tool_calls": [{ + "name": "list_dir", + "args": {"DirectoryPath":"/tmp","toolSummary":"List /tmp"} + }] + }), + json!({ + "step_index": 3, + "created_at": "2026-08-11T17:46:21Z", + "source": "MODEL", + "type": "LIST_DIRECTORY", + "status": "DONE", + "content": "Created At: 2026-08-12T01:46:21+08:00\nCompleted At: 2026-08-12T01:46:21+08:00\n{\"name\":\"sample\"}" + }), + json!({ + "step_index": 4, + "created_at": "2026-08-11T17:46:21Z", + "source": "SYSTEM", + "type": "CHECKPOINT", + "status": "DONE", + "content": "checkpoint" + }), + json!({ + "step_index": 5, + "created_at": "2026-08-11T17:46:21Z", + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "DONE", + "content": "done" + }), + ]; + let (transcript, boundary) = jsonl(&records); + let path = "/tmp/conversation/transcript_full.jsonl"; + let registry = Registry::default_agents(); + let mut translator = registry.create("antigravity", "real-conversation"); + let ctx = SessionCtx { + session_id: "real-conversation".into(), + config: None, + }; + let mut ops = Vec::new(); + ops.extend( + translator + .handle( + &event( + "PreInvocation", + 100, + path, + &transcript, + boundary[1], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "PostInvocation", + 200, + path, + &transcript, + boundary[3], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "PostToolUse", + 210, + path, + &transcript, + boundary[3], + json!({ + "stepIdx":3, + "error":"", + "toolCall":{"name":"list_dir","args":{"DirectoryPath":"/tmp"}} + }), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "PreInvocation", + 220, + path, + &transcript, + boundary[3], + json!({"invocationNum":1,"initialNumSteps":5}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "PostInvocation", + 300, + path, + &transcript, + boundary[5], + json!({"invocationNum":1,"initialNumSteps":5}), + ), + &ctx, + ) + .unwrap(), + ); + ops.extend( + translator + .handle( + &event( + "Stop", + 310, + path, + &transcript, + boundary[5], + json!({"executionNum":0,"terminationReason":"NO_TOOL_CALL","fullyIdle":true}), + ), + &ctx, + ) + .unwrap(), + ); + + let rows = reduce(ops); + let turn = rows.values().find(|row| row.name == "Turn 1").unwrap(); + assert_eq!(turn.input, Some(json!("List /tmp"))); + assert_eq!(turn.output, Some(json!("done"))); + + let mut llms = rows + .values() + .filter(|row| row.span_type == SpanType::Llm) + .collect::>(); + llms.sort_by_key(|row| row.start_ms); + assert_eq!(llms.len(), 2); + assert_eq!( + llms[0].output.as_ref().unwrap()[0]["tool_calls"][0]["name"], + "list_dir" + ); + assert!(llms[1] + .input + .as_ref() + .unwrap() + .as_array() + .unwrap() + .iter() + .any(|message| message["role"] == "tool")); + assert_eq!(llms[1].output.as_ref().unwrap()[0]["content"], "done"); + + let tool = rows + .values() + .find(|row| row.span_type == SpanType::Tool) + .unwrap(); + assert_eq!(tool.name, "list_dir"); + assert_eq!(tool.start_ms, Some(1_786_470_380_000)); + assert_eq!(tool.output, Some(json!("{\"name\":\"sample\"}"))); + assert_eq!(tool.metadata.as_ref().unwrap()["tool_outcome"], "success"); +} + +#[test] +fn resumed_process_reuses_invocation_zero_without_reparenting_to_turn_one() { + let records = vec![ + json!({ + "step_index": 0, + "source": "USER_EXPLICIT", + "type": "USER_INPUT", + "status": "DONE", + "content": "Use a tool" + }), + json!({ + "step_index": 1, + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "DONE", + "tool_calls": [{ + "name": "list_dir", + "args": {"DirectoryPath":"/workspace/demo"} + }] + }), + json!({ + "step_index": 2, + "source": "MODEL", + "type": "LIST_DIRECTORY", + "status": "DONE", + "content": "README.md" + }), + json!({ + "step_index": 3, + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "DONE", + "content": "first-turn-complete" + }), + json!({ + "step_index": 4, + "source": "USER_EXPLICIT", + "type": "USER_INPUT", + "status": "DONE", + "content": "This is turn two" + }), + json!({ + "step_index": 5, + "source": "MODEL", + "type": "PLANNER_RESPONSE", + "status": "DONE", + "content": "second-turn-complete" + }), + ]; + let (transcript, boundary) = jsonl(&records); + let path = "/tmp/resumed-conversation/transcript_full.jsonl"; + let registry = Registry::default_agents(); + let mut translator = registry.create("antigravity", "resumed-conversation"); + let ctx = SessionCtx { + session_id: "resumed-conversation".into(), + config: None, + }; + let mut ops = Vec::new(); + + // First process: invocation zero asks for a tool, invocation one consumes + // its result, then a fully-idle Stop closes the root. + for hook in [ + event( + "PreInvocation", + 100, + path, + &transcript, + boundary[0], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + event( + "PostInvocation", + 200, + path, + &transcript, + boundary[1], + json!({"invocationNum":0,"initialNumSteps":1}), + ), + event( + "PostToolUse", + 210, + path, + &transcript, + boundary[2], + json!({"stepIdx":2,"error":""}), + ), + event( + "PreInvocation", + 220, + path, + &transcript, + boundary[2], + json!({"invocationNum":1,"initialNumSteps":3}), + ), + event( + "PostInvocation", + 300, + path, + &transcript, + boundary[3], + json!({"invocationNum":1,"initialNumSteps":3}), + ), + event( + "Stop", + 310, + path, + &transcript, + boundary[3], + json!({"executionNum":0,"fullyIdle":true}), + ), + // A new process resumes the conversation. Antigravity resets + // invocationNum to zero while the transcript step index continues. + event( + "PreInvocation", + 400, + path, + &transcript, + boundary[4], + json!({"invocationNum":0,"initialNumSteps":5}), + ), + event( + "PostInvocation", + 500, + path, + &transcript, + boundary[5], + json!({"invocationNum":0,"initialNumSteps":5}), + ), + event( + "Stop", + 510, + path, + &transcript, + boundary[5], + json!({"executionNum":0,"fullyIdle":true}), + ), + ] { + ops.extend(translator.handle(&hook, &ctx).unwrap()); + } + + let rows = reduce(ops); + let root = rows + .values() + .find(|row| row.name == "Antigravity: demo") + .unwrap(); + assert_eq!(root.end_ms, Some(510)); + + let turn_one = rows.values().find(|row| row.name == "Turn 1").unwrap(); + let turn_two = rows.values().find(|row| row.name == "Turn 2").unwrap(); + assert_eq!(turn_one.output, Some(json!("first-turn-complete"))); + assert_eq!(turn_two.input, Some(json!("This is turn two"))); + assert_eq!(turn_two.output, Some(json!("second-turn-complete"))); + + let llms = rows + .values() + .filter(|row| row.span_type == SpanType::Llm) + .collect::>(); + assert_eq!(llms.len(), 3); + let resumed = llms + .iter() + .find(|row| row.metadata.as_ref().unwrap()["turn_number"] == 2) + .unwrap(); + assert_eq!(resumed.parent_span_ids, vec![turn_two.span_id.clone()]); + assert_eq!( + resumed.output.as_ref().unwrap()[0]["content"], + "second-turn-complete" + ); + let invocation_zero = llms + .iter() + .filter(|row| row.metadata.as_ref().unwrap()["invocation_num"] == 0) + .collect::>(); + assert_eq!(invocation_zero.len(), 2); + assert_ne!(invocation_zero[0].span_id, invocation_zero[1].span_id); + + let tool = rows + .values() + .find(|row| row.span_type == SpanType::Tool) + .unwrap(); + assert_eq!(tool.name, "list_directory"); + assert_eq!(tool.parent_span_ids, vec![turn_one.span_id.clone()]); + assert_eq!(tool.metadata.as_ref().unwrap()["tool_outcome"], "success"); +} diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index d3aeee6..5944a17 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -90,6 +90,66 @@ async fn logs3_bodies(server: &MockServer) -> String { .join("\n") } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn explicit_project_id_does_not_register_a_project_name() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .respond_with(ResponseTemplate::new(403)) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + + let base = server.uri(); + let factory = BraintrustSinkFactory::new(BraintrustSinkConfig { + api_url: Some(base.clone()), + app_url: Some(base.clone()), + version: "test".into(), + }); + let mut sink = factory.create("sess-id", "antigravity", None).unwrap(); + let mut config = session_config(&base); + config.destination = Some(TraceDestination::ProjectLogs { + project_id: Some("proj-existing".into()), + project_name: None, + }); + sink.configure(&config); + sink.emit(&[SpanOp::Insert(row( + "root-id", + "root-id", + &[], + "Antigravity", + SpanType::Task, + 1, + Some(2), + ))]) + .await + .unwrap(); + sink.flush().await.unwrap(); + + let requests = server.received_requests().await.unwrap(); + assert!( + requests + .iter() + .any(|request| request.url.path() == "/logs3"), + "expected delivery to the existing project" + ); + assert!( + !requests + .iter() + .any(|request| request.url.path() == "/api/project/register"), + "an explicit project id must not trigger project registration" + ); +} + /// Two sessions on two different backend URLs, from one factory, each deliver /// only to their own collector — the per-`(api_url, app_url)` client cache. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -176,6 +236,61 @@ async fn merge_with_empty_name_does_not_clobber_the_original_name() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_resumed_session_can_extend_an_already_ended_root() { + let server = mock_backend().await; + let base = server.uri(); + let factory = BraintrustSinkFactory::new(BraintrustSinkConfig { + api_url: Some(base.clone()), + app_url: Some(base.clone()), + version: "test".into(), + }); + let mut sink = factory.create("sess-resumed", "antigravity", None).unwrap(); + sink.configure(&session_config(&base)); + + let root = row( + "resumed-root", + "resumed-root", + &[], + "Antigravity", + SpanType::Task, + 1_000, + None, + ); + let first_stop = row( + "resumed-root", + "resumed-root", + &[], + "", + SpanType::Task, + 1_000, + Some(2_000), + ); + let resumed_stop = row( + "resumed-root", + "resumed-root", + &[], + "", + SpanType::Task, + 1_000, + Some(5_000), + ); + sink.emit(&[ + SpanOp::Insert(root), + SpanOp::Merge(first_stop), + SpanOp::Merge(resumed_stop), + ]) + .await + .unwrap(); + sink.flush().await.unwrap(); + + let bodies = logs3_bodies(&server).await; + assert!( + bodies.contains("\"end\":5.0"), + "later stop did not extend the root end time: {bodies}" + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn attached_trace_children_keep_the_external_root() { let server = mock_backend().await; diff --git a/src/plugins/antigravity/build.sh b/src/plugins/antigravity/build.sh new file mode 100755 index 0000000..3e093f8 --- /dev/null +++ b/src/plugins/antigravity/build.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +set -euo pipefail + +TARGET_DIR="${1:?usage: build.sh }" +SRC_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +mkdir -p "$TARGET_DIR" +rsync -a --delete --exclude '.git' "$SRC_DIR/content/" "$TARGET_DIR/" +echo "Built antigravity dist into $TARGET_DIR." diff --git a/src/plugins/antigravity/content/README.md b/src/plugins/antigravity/content/README.md new file mode 100644 index 0000000..4453a07 --- /dev/null +++ b/src/plugins/antigravity/content/README.md @@ -0,0 +1,20 @@ +# Braintrust tracing for Google Antigravity + +This Antigravity plugin forwards native lifecycle hooks to the Braintrust +daemon. The daemon combines exact model and tool boundaries from hooks with +the conversation's full JSONL transcript to construct a session, turn, model, +and tool span tree. + +The hook adapter is synchronous, credential-free, and fail-open. Braintrust +authentication and destination routing remain owned by the `bt` CLI. + +The initial implementation captures `PreInvocation`, `PostInvocation`, +`PostToolUse`, and `Stop`. It intentionally does not register `PreToolUse`: +Antigravity requires that hook to return a permission decision. Live testing +confirmed that an empty decision is handled as a denial, while `allow` would +bypass normal permission checks and `ask` could add prompts. + +This package currently requires a `bt` CLI that exposes `bt trace hook` and a +Unix-compatible `sh`. Persistent setup, managed-run injection, transcript +import/attach, Windows support, and production distribution are not yet part of +this feasibility implementation. diff --git a/src/plugins/antigravity/content/bin/antigravity-hook.sh b/src/plugins/antigravity/content/bin/antigravity-hook.sh new file mode 100755 index 0000000..6b0d81d --- /dev/null +++ b/src/plugins/antigravity/content/bin/antigravity-hook.sh @@ -0,0 +1,25 @@ +#!/bin/sh +# Thin, credential-free Antigravity hook adapter. Antigravity runs hooks from +# the directory containing hooks.json and requires a JSON response on stdout. +# Tracing is deliberately fail-open: a missing or unhealthy bt CLI must never +# interrupt the coding-agent loop. + +event=${1:-} +bt_bin=${BT_BIN:-bt} + +if [ -n "$event" ] && command -v "$bt_bin" >/dev/null 2>&1; then + "$bt_bin" trace hook \ + --source antigravity \ + --session-id-field conversationId \ + --event "$event" \ + --transcript-path-field transcriptPath \ + --flush-on-turn-end \ + >/dev/null 2>&1 || : +fi + +case "$event" in + Stop) printf '{"decision":""}\n' ;; + *) printf '{}\n' ;; +esac + +exit 0 diff --git a/src/plugins/antigravity/content/hooks.json b/src/plugins/antigravity/content/hooks.json new file mode 100644 index 0000000..ebd6f16 --- /dev/null +++ b/src/plugins/antigravity/content/hooks.json @@ -0,0 +1,37 @@ +{ + "braintrust-tracing": { + "PostToolUse": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "sh \"./bin/antigravity-hook.sh\" PostToolUse", + "timeout": 30 + } + ] + } + ], + "PreInvocation": [ + { + "type": "command", + "command": "sh \"./bin/antigravity-hook.sh\" PreInvocation", + "timeout": 30 + } + ], + "PostInvocation": [ + { + "type": "command", + "command": "sh \"./bin/antigravity-hook.sh\" PostInvocation", + "timeout": 30 + } + ], + "Stop": [ + { + "type": "command", + "command": "sh \"./bin/antigravity-hook.sh\" Stop", + "timeout": 30 + } + ] + } +} diff --git a/src/plugins/antigravity/content/plugin.json b/src/plugins/antigravity/content/plugin.json new file mode 100644 index 0000000..57ad84e --- /dev/null +++ b/src/plugins/antigravity/content/plugin.json @@ -0,0 +1,3 @@ +{ + "name": "braintrust-antigravity-tracing" +} diff --git a/src/plugins/antigravity/test/bt-standalone-wrapper.sh b/src/plugins/antigravity/test/bt-standalone-wrapper.sh new file mode 100755 index 0000000..931e073 --- /dev/null +++ b/src/plugins/antigravity/test/bt-standalone-wrapper.sh @@ -0,0 +1,13 @@ +#!/bin/sh +# Adapt the production `bt trace hook ...` invocation to the standalone +# bt-daemon test binary while retaining the exact plugin command contract. +set -eu + +: "${BT_DAEMON_BIN:?set BT_DAEMON_BIN}" +: "${BT_DAEMON_SOCKET:?set BT_DAEMON_SOCKET}" + +[ "${1:-}" = "trace" ] +[ "${2:-}" = "hook" ] +shift 2 + +exec "$BT_DAEMON_BIN" hook "$@" --socket "$BT_DAEMON_SOCKET" --no-spawn diff --git a/src/plugins/antigravity/test/capture-bt.sh b/src/plugins/antigravity/test/capture-bt.sh new file mode 100755 index 0000000..7b0570d --- /dev/null +++ b/src/plugins/antigravity/test/capture-bt.sh @@ -0,0 +1,15 @@ +#!/bin/sh +# Test-only bt stub for recording the exact payloads emitted by a real +# Antigravity session. The production plugin never ships this file. +set -eu + +capture_dir=${ANTIGRAVITY_CAPTURE_DIR:?set ANTIGRAVITY_CAPTURE_DIR} +mkdir -p "$capture_dir" +payload=$(mktemp "$capture_dir/payload.XXXXXX") +cp /dev/stdin "$payload" +printf '%s\n' "$@" > "$payload.args" +transcript=$(jq -r '.transcriptPath // empty' "$payload" 2>/dev/null || true) +if [ -n "$transcript" ] && [ -f "$transcript" ]; then + wc -c < "$transcript" > "$payload.transcript-bytes" + cp "$transcript" "$payload.transcript.jsonl" +fi diff --git a/src/plugins/antigravity/test/test_hook.sh b/src/plugins/antigravity/test/test_hook.sh new file mode 100755 index 0000000..5499028 --- /dev/null +++ b/src/plugins/antigravity/test/test_hook.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +set -euo pipefail + +HOOK="${1:?usage: test_hook.sh }" +TMP="$(mktemp -d)" +trap 'rm -rf "$TMP"' EXIT + +BT_STUB="$TMP/bt" +cp /dev/stdin "$BT_STUB" <<'STUB' +#!/bin/sh +printf '%s\n' "$@" > "$BT_STUB_ARGS" +cp /dev/stdin "$BT_STUB_STDIN" +exit "${BT_STUB_EXIT:-0}" +STUB +chmod +x "$BT_STUB" + +export BT_STUB_ARGS="$TMP/args" +export BT_STUB_STDIN="$TMP/stdin" +payload='{"conversationId":"test","transcriptPath":"/tmp/transcript.jsonl"}' + +response=$(printf '%s' "$payload" | BT_BIN="$BT_STUB" "$HOOK" PostInvocation) +[[ "$response" == '{}' ]] +cmp -s "$TMP/stdin" <(printf '%s' "$payload") +grep -Fx -- 'trace' "$TMP/args" >/dev/null +grep -Fx -- 'antigravity' "$TMP/args" >/dev/null +grep -Fx -- 'conversationId' "$TMP/args" >/dev/null +grep -Fx -- 'transcriptPath' "$TMP/args" >/dev/null + +response=$(printf '%s' "$payload" | BT_STUB_EXIT=1 BT_BIN="$BT_STUB" "$HOOK" Stop) +[[ "$response" == '{"decision":""}' ]] + +echo "test: antigravity hook adapter OK" diff --git a/src/plugins/antigravity/validate.sh b/src/plugins/antigravity/validate.sh new file mode 100755 index 0000000..143e312 --- /dev/null +++ b/src/plugins/antigravity/validate.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +set -euo pipefail + +TARGET_DIR="${1:?usage: validate.sh }" +fail() { echo "validate: $*" >&2; exit 1; } + +for file in plugin.json hooks.json bin/antigravity-hook.sh README.md; do + [[ -f "$TARGET_DIR/$file" ]] || fail "missing $file" +done +[[ -x "$TARGET_DIR/bin/antigravity-hook.sh" ]] || fail "hook adapter is not executable" + +if command -v jq >/dev/null 2>&1; then + jq empty "$TARGET_DIR/plugin.json" "$TARGET_DIR/hooks.json" >/dev/null \ + || fail "invalid JSON" +else + python3 -m json.tool "$TARGET_DIR/plugin.json" >/dev/null || fail "invalid plugin.json" + python3 -m json.tool "$TARGET_DIR/hooks.json" >/dev/null || fail "invalid hooks.json" +fi + +"$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/test/test_hook.sh" \ + "$TARGET_DIR/bin/antigravity-hook.sh" +echo "validate: antigravity dist OK ($TARGET_DIR)"