From 67bc8f7029c26685030914830253449644200f11 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 18 Aug 2026 18:32:05 +0800 Subject: [PATCH 1/2] fix: bound Codex catch-up memory --- bt-daemon/src/dispatch.rs | 101 ++++++++---- bt-daemon/src/lib.rs | 14 +- bt-daemon/src/translate/codex.rs | 241 +++++++++++++++++++++++----- bt-daemon/src/translate/mod.rs | 7 + bt-daemon/tests/codex_translator.rs | 96 +++++++++++ bt-daemon/tests/replay.rs | 34 ++++ 6 files changed, 421 insertions(+), 72 deletions(-) diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index 49a30bf..d065898 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -266,15 +266,16 @@ impl SessionActor { ctx.config = Some(cfg.clone()); self.refresh_permalink(sink.as_ref()); } - match translator.handle(&env, &ctx) { - Ok(ops) => match sink.emit(&ops).await { - Ok(n) => { - self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); - } - Err(e) => self.set_error(format!("sink emit failed: {e}")), - }, - Err(e) => self.set_error(format!("translate failed: {e}")), - } + let translated = translator.handle(&env, &ctx); + self.emit_translator_batches( + &mut translator, + &mut sink, + &ctx, + translated, + "translate failed", + "sink emit failed", + ) + .await; self.counters.queued.fetch_sub(1, Ordering::Relaxed); } SessionMsg::Configure(config, reply) => { @@ -296,6 +297,44 @@ impl SessionActor { } } + /// Emit a translator result and every bounded continuation it schedules. + /// A continuation may be empty (for example, irrelevant rollout rows), so + /// only `None` signals that the translator is fully caught up. + async fn emit_translator_batches( + &self, + translator: &mut Box, + sink: &mut Box, + ctx: &SessionCtx, + first: anyhow::Result>, + translate_error: &str, + emit_error: &str, + ) { + let mut next = match first { + Ok(ops) => Some(ops), + Err(e) => { + self.set_error(format!("{translate_error}: {e}")); + return; + } + }; + while let Some(ops) = next { + if !ops.is_empty() { + match sink.emit(&ops).await { + Ok(n) => { + self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); + } + Err(e) => self.set_error(format!("{emit_error}: {e}")), + } + } + next = match translator.drain_pending(ctx) { + Ok(next) => next, + Err(e) => { + self.set_error(format!("{translate_error}: {e}")); + None + } + }; + } + } + /// Stream the journal through the translator, emitting each entry's spans /// as they are produced. Nothing is accumulated across entries: peak /// memory is one journal entry and the ops it yields, so recovering a @@ -336,22 +375,16 @@ impl SessionActor { continue; } let env = crate::journal::envelope_from_redacted(entry); - let ops = match translator.handle(&env, ctx) { - Ok(ops) => ops, - Err(e) => { - self.set_error(format!("journal replay failed: {e}")); - continue; - } - }; - if ops.is_empty() { - continue; - } - match sink.emit(&ops).await { - Ok(n) => { - self.counters.spans_emitted.fetch_add(n, Ordering::Relaxed); - } - Err(e) => self.set_error(format!("sink replay emit failed: {e}")), - } + let translated = translator.handle(&env, ctx); + self.emit_translator_batches( + translator, + sink, + ctx, + translated, + "journal replay failed", + "sink replay emit failed", + ) + .await; } } @@ -361,14 +394,16 @@ impl SessionActor { sink: &mut Box, ctx: &SessionCtx, ) { - match translator.flush(ctx) { - Ok(ops) => { - if let Err(e) = sink.emit(&ops).await { - self.set_error(format!("sink emit (flush) failed: {e}")); - } - } - Err(e) => self.set_error(format!("translate flush failed: {e}")), - } + let translated = translator.flush(ctx); + self.emit_translator_batches( + translator, + sink, + ctx, + translated, + "translate flush failed", + "sink emit (flush) failed", + ) + .await; if let Err(e) = sink.flush().await { self.set_error(format!("sink flush failed: {e}")); } diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 1ec5532..4a7afe7 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -822,6 +822,17 @@ impl ImportProcessor { live.ctx.config = Some(cfg.clone()); } let ops = live.translator.handle(&env, &live.ctx)?; + Self::emit_translator_batches(live, ops).await?; + } + Ok(()) + } + + async fn emit_translator_batches( + live: &mut ImportLive, + first: Vec, + ) -> anyhow::Result<()> { + let mut next = Some(first); + while let Some(ops) = next { // Imports can contain tens of thousands of SDK log commands. Bound the // number queued between drains without serializing one network flush // for every native turn boundary. @@ -834,6 +845,7 @@ impl ImportProcessor { live.pending_ops = 0; } } + next = live.translator.drain_pending(&live.ctx)?; } Ok(()) } @@ -841,7 +853,7 @@ impl ImportProcessor { async fn finish(self) -> anyhow::Result<()> { for (_sid, mut live) in self.sessions { let ops = live.translator.flush(&live.ctx)?; - live.sink.emit(&ops).await?; + Self::emit_translator_batches(&mut live, ops).await?; live.sink.flush().await?; } Ok(()) diff --git a/bt-daemon/src/translate/codex.rs b/bt-daemon/src/translate/codex.rs index 31c97e3..bc94451 100644 --- a/bt-daemon/src/translate/codex.rs +++ b/bt-daemon/src/translate/codex.rs @@ -29,6 +29,9 @@ use std::sync::{Arc, OnceLock}; const SPAWN_AGENT_TOOL: &str = "spawn_agent"; const MISSING_TOOL_OUTPUT_ERROR: &str = "Tool output missing before turn ended"; +/// A hook can arrive after a daemon restart or against an existing rollout. +/// Keep a single translator batch small even when the unread suffix is large. +const CATCH_UP_BYTE_BUDGET: usize = 64 * 1024; pub struct CodexTranslatorFactory { git: Arc, @@ -63,6 +66,7 @@ impl TranslatorFactory for CodexTranslatorFactory { spawn_turn_by_agent_id: HashMap::new(), compaction_trigger_by_turn: HashMap::new(), compaction_spans: HashSet::new(), + pending: None, git: self.git.clone(), }) } @@ -115,6 +119,37 @@ struct Scope { subagent_ended: bool, } +enum DeferredHook { + None, + PostToolUse(Value), + SubagentStop { + path: Option, + ts: i64, + }, + Stop { + is_main: bool, + payload: Value, + ts: i64, + }, + PostCompact { + payload: Value, + ts: i64, + }, +} + +enum PendingWork { + Hook { + path: String, + hook_ts: i64, + through_ms: Option, + after: DeferredHook, + }, + Flush { + paths: Vec, + next_path: usize, + }, +} + struct CodexTranslator { session_id: String, root_span_id: String, @@ -132,11 +167,16 @@ struct CodexTranslator { spawn_turn_by_agent_id: HashMap, compaction_trigger_by_turn: HashMap, compaction_spans: HashSet, + pending: Option, git: Arc, } impl AgentTranslator for CodexTranslator { fn handle(&mut self, event: &Envelope, ctx: &SessionCtx) -> anyhow::Result> { + anyhow::ensure!( + self.pending.is_none(), + "Codex translator has pending catch-up work; drain it before handling another event" + ); let payload = &event.payload; let mut ops = Vec::new(); @@ -177,45 +217,88 @@ impl AgentTranslator for CodexTranslator { self.ensure_main_scope(&path); } let import_through_ms = payload.get("_bt_import_through_ms").and_then(Value::as_i64); - self.catch_up(&path, event.ts_ms, import_through_ms, &mut ops); + let after = self.deferred_hook(event, agent_id.is_none()); + if self.catch_up_chunk(&path, event.ts_ms, import_through_ms, &mut ops) { + self.finish_deferred_hook(after, &mut ops); + } else { + self.pending = Some(PendingWork::Hook { + path, + hook_ts: event.ts_ms, + through_ms: import_through_ms, + after, + }); + } + } else { + self.finish_deferred_hook(self.deferred_hook(event, agent_id.is_none()), &mut ops); } - // --- hook-specific handling (after catch-up) --- - match event.event.as_str() { - // Catch up first: this same hook may be the first observation of - // the spawn_agent transcript record that establishes call -> turn. - "PostToolUse" => self.record_spawned_agent(payload), - "SubagentStop" => { - if let Some(p) = str_field(payload, "agent_transcript_path") { - self.close_subagent(&p, event.ts_ms, &mut ops); + Ok(ops) + } + + fn drain_pending(&mut self, _ctx: &SessionCtx) -> anyhow::Result>> { + let Some(pending) = self.pending.take() else { + return Ok(None); + }; + let mut ops = Vec::new(); + match pending { + PendingWork::Hook { + path, + hook_ts, + through_ms, + after, + } => { + if self.catch_up_chunk(&path, hook_ts, through_ms, &mut ops) { + self.finish_deferred_hook(after, &mut ops); + } else { + self.pending = Some(PendingWork::Hook { + path, + hook_ts, + through_ms, + after, + }); } } - "Stop" if agent_id.is_none() => { - // Codex writes task_complete slightly after the Stop hook in - // real sessions. Close the active turn from the hook payload - // now so a short-lived process cannot flush an open turn. - self.close_main_turn(payload, event.ts_ms, &mut ops); - self.end_main_root(event.ts_ms, &mut ops); + PendingWork::Flush { + paths, + mut next_path, + } => { + while next_path < paths.len() { + let path = &paths[next_path]; + if self.catch_up_chunk(path, 0, None, &mut ops) { + if let Some(mut scope) = self.scopes.remove(path) { + self.close_dangling(&mut scope, None, &mut ops); + self.scopes.insert(path.clone(), scope); + } + next_path += 1; + } + // Return after any completed scope or a bounded partial read. + // This keeps a flush over many scopes bounded as well. + if !ops.is_empty() || next_path < paths.len() { + self.pending = (next_path < paths.len()) + .then_some(PendingWork::Flush { paths, next_path }); + break; + } + } } - "PostCompact" => self.close_compaction_turn(payload, event.ts_ms, &mut ops), - _ => {} } - - Ok(ops) + Ok(Some(ops)) } - fn flush(&mut self, _ctx: &SessionCtx) -> anyhow::Result> { - let mut ops = Vec::new(); + fn flush(&mut self, ctx: &SessionCtx) -> anyhow::Result> { + anyhow::ensure!( + self.pending.is_none(), + "Codex translator has pending catch-up work; drain it before flushing" + ); // Re-read each scope to catch a late task_complete, then close dangling. let paths: Vec = self.scopes.keys().cloned().collect(); - for path in paths { - self.catch_up(&path, 0, None, &mut ops); - if let Some(mut scope) = self.scopes.remove(&path) { - self.close_dangling(&mut scope, None, &mut ops); - self.scopes.insert(path, scope); - } + if paths.is_empty() { + return Ok(Vec::new()); } - Ok(ops) + self.pending = Some(PendingWork::Flush { + paths, + next_path: 0, + }); + Ok(self.drain_pending(ctx)?.unwrap_or_default()) } } @@ -294,24 +377,79 @@ impl CodexTranslator { self.scopes.insert(path, scope); } - /// Read new transcript lines for `path` and process them against its scope. - fn catch_up( + fn deferred_hook(&self, event: &Envelope, is_main_scope: bool) -> DeferredHook { + match event.event.as_str() { + // Catch up first: this same hook may be the first observation of + // the spawn_agent transcript record that establishes call -> turn. + "PostToolUse" => DeferredHook::PostToolUse(event.payload.clone()), + "SubagentStop" => DeferredHook::SubagentStop { + path: str_field(&event.payload, "agent_transcript_path"), + ts: event.ts_ms, + }, + // Codex writes task_complete slightly after the Stop hook in real + // sessions. Do this only after the entire bounded catch-up finishes. + "Stop" => DeferredHook::Stop { + is_main: is_main_scope, + payload: event.payload.clone(), + ts: event.ts_ms, + }, + "PostCompact" => DeferredHook::PostCompact { + payload: event.payload.clone(), + ts: event.ts_ms, + }, + _ => DeferredHook::None, + } + } + + fn finish_deferred_hook(&mut self, after: DeferredHook, ops: &mut Vec) { + match after { + DeferredHook::None => {} + DeferredHook::PostToolUse(payload) => self.record_spawned_agent(&payload), + DeferredHook::SubagentStop { path, ts } => { + if let Some(path) = path { + self.close_subagent(&path, ts, ops); + } + } + DeferredHook::Stop { + is_main: true, + payload, + ts, + } => { + self.close_main_turn(&payload, ts, ops); + self.end_main_root(ts, ops); + } + DeferredHook::Stop { is_main: false, .. } => {} + DeferredHook::PostCompact { payload, ts } => { + self.close_compaction_turn(&payload, ts, ops) + } + } + } + + /// Read a bounded batch of transcript lines for `path` and process them + /// against its scope. `true` means the requested catch-up is complete. + fn catch_up_chunk( &mut self, path: &str, hook_ts: i64, through_ms: Option, ops: &mut Vec, - ) { + ) -> bool { let Some(mut scope) = self.scopes.remove(path) else { - return; + return true; }; - let lines = read_new_lines(&scope.path, &mut scope.offset, through_ms); - for line in lines { + let read = read_new_lines( + &scope.path, + &mut scope.offset, + through_ms, + CATCH_UP_BYTE_BUDGET, + ); + for line in read.lines { if let Ok(rec) = serde_json::from_str::(&line) { self.process_record(&mut scope, &rec, hook_ts, ops); } } self.scopes.insert(path.to_string(), scope); + read.complete } fn process_record( @@ -1454,10 +1592,23 @@ fn parse_ts(rec: &Value) -> Option { .map(|dt| dt.timestamp_millis()) } -fn read_new_lines(path: &str, offset: &mut u64, through_ms: Option) -> Vec { +struct ReadLines { + lines: Vec, + complete: bool, +} + +fn read_new_lines( + path: &str, + offset: &mut u64, + through_ms: Option, + byte_budget: usize, +) -> ReadLines { use std::io::{BufRead, BufReader, Seek, SeekFrom}; let Ok(mut f) = std::fs::File::open(path) else { - return Vec::new(); + return ReadLines { + lines: Vec::new(), + complete: true, + }; }; if let Ok(meta) = f.metadata() { if *offset > meta.len() { @@ -1465,11 +1616,21 @@ fn read_new_lines(path: &str, offset: &mut u64, through_ms: Option) -> Vec< } } if f.seek(SeekFrom::Start(*offset)).is_err() { - return Vec::new(); + return ReadLines { + lines: Vec::new(), + complete: true, + }; } let mut reader = BufReader::new(f); let mut lines = Vec::new(); + let mut consumed = 0usize; loop { + if consumed >= byte_budget && !lines.is_empty() { + return ReadLines { + lines, + complete: false, + }; + } let mut line = String::new(); let Ok(bytes) = reader.read_line(&mut line) else { break; @@ -1488,11 +1649,15 @@ fn read_new_lines(path: &str, offset: &mut u64, through_ms: Option) -> Vec< } } *offset += bytes as u64; + consumed += bytes; if !trimmed.trim().is_empty() { lines.push(trimmed.to_string()); } } - lines + ReadLines { + lines, + complete: true, + } } fn token_metrics(usage: &Value) -> Map { diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index d182161..8e6e10a 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -89,6 +89,13 @@ pub trait AgentTranslator: Send { /// Handle one event, returning span ops to emit. fn handle(&mut self, event: &Envelope, ctx: &SessionCtx) -> anyhow::Result>; + /// Continue bounded work started by [`Self::handle`] or [`Self::flush`]. + /// `Some` means the caller must emit this batch and call again; `None` + /// means the translator is fully caught up. + fn drain_pending(&mut self, _ctx: &SessionCtx) -> anyhow::Result>> { + Ok(None) + } + /// Emit any pending spans (e.g. close dangling turns) at flush/shutdown. fn flush(&mut self, ctx: &SessionCtx) -> anyhow::Result> { let _ = ctx; diff --git a/bt-daemon/tests/codex_translator.rs b/bt-daemon/tests/codex_translator.rs index 8f94d1a..955d574 100644 --- a/bt-daemon/tests/codex_translator.rs +++ b/bt-daemon/tests/codex_translator.rs @@ -273,6 +273,102 @@ fn codex_incremental_reads_advance_offset() { } } +/// Reproduction for a large rollout first observed through one `SessionStart` +/// hook. The source transcript stays modest, but each LLM span snapshots the +/// growing conversation history. The translator must emit bounded batches +/// instead of retaining every clone from the catch-up at once. +#[test] +fn codex_large_catch_up_emits_bounded_batches() { + const CALLS: usize = 256; + const USER_BYTES: usize = 8 * 1024; + + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("large-rollout.jsonl"); + let tpath = transcript.to_str().unwrap(); + let mut file = std::fs::File::create(&transcript).unwrap(); + for record in [ + json!({ "timestamp": "2026-01-01T00:00:01Z", "type": "session_meta", "payload": { "id": "large", "cwd": "/x/app" } }), + json!({ "timestamp": "2026-01-01T00:00:02Z", "type": "turn_context", "payload": { "model": "gpt-test" } }), + json!({ "timestamp": "2026-01-01T00:00:03Z", "type": "event_msg", "payload": { "type": "task_started", "turn_id": "large-turn" } }), + ] { + writeln!(file, "{}", line(record)).unwrap(); + } + let user_text = "x".repeat(USER_BYTES); + for index in 0..CALLS { + let timestamp = |offset: usize| { + let seconds = 4 + index * 3 + offset; + format!("2026-01-01T00:{:02}:{:02}Z", seconds / 60, seconds % 60) + }; + writeln!(file, "{}", line(json!({ + "timestamp": timestamp(0), + "type": "response_item", + "payload": { "type": "message", "role": "user", "content": [{ "type": "output_text", "text": user_text }] } + }))).unwrap(); + writeln!(file, "{}", line(json!({ + "timestamp": timestamp(1), + "type": "response_item", + "payload": { "type": "message", "role": "assistant", "content": [{ "type": "output_text", "text": "ok" }] } + }))).unwrap(); + writeln!(file, "{}", line(json!({ + "timestamp": timestamp(2), + "type": "event_msg", + "payload": { "type": "token_count", "info": { "last_token_usage": { "input_tokens": 1, "output_tokens": 1 } } } + }))).unwrap(); + } + file.flush().unwrap(); + + let reg = Registry::default_agents(); + let mut translator = reg.create("codex", "large"); + let ctx = SessionCtx { + session_id: "large".into(), + config: None, + }; + let mut next = Some( + translator + .handle(&envelope("large", "SessionStart", tpath, json!({})), &ctx) + .unwrap(), + ); + let mut batches = 0; + let mut total_llms = 0; + let mut total_llm_inputs = 0; + let mut largest_batch_inputs = 0; + while let Some(ops) = next { + batches += 1; + let batch_inputs = ops + .iter() + .filter_map(|op| match op { + SpanOp::Insert(row) if row.span_type == SpanType::Llm => row.input.as_ref(), + _ => None, + }) + .map(|input| serde_json::to_vec(input).unwrap().len()) + .sum::(); + total_llms += ops + .iter() + .filter(|op| matches!(op, SpanOp::Insert(row) if row.span_type == SpanType::Llm)) + .count(); + total_llm_inputs += batch_inputs; + largest_batch_inputs = largest_batch_inputs.max(batch_inputs); + next = translator.drain_pending(&ctx).unwrap(); + } + + eprintln!( + "reproduction: rollout={} bytes, {batches} batches, total llm inputs={} bytes, largest batch={} bytes", + std::fs::metadata(&transcript).unwrap().len(), + total_llm_inputs, + largest_batch_inputs, + ); + assert_eq!(total_llms, CALLS); + assert!( + total_llm_inputs > 200 * 1024 * 1024, + "expected substantial total output" + ); + assert!(batches > 1, "large rollout must be drained incrementally"); + assert!( + largest_batch_inputs < 40 * 1024 * 1024, + "each emitted batch must stay bounded" + ); +} + #[test] fn codex_import_checkpoints_preserve_native_turn_boundaries() { let tmp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/tests/replay.rs b/bt-daemon/tests/replay.rs index 2b2d79b..4ac4120 100644 --- a/bt-daemon/tests/replay.rs +++ b/bt-daemon/tests/replay.rs @@ -12,6 +12,40 @@ fn write_jsonl(path: &std::path::Path, records: &[Value]) { } } +#[tokio::test] +async fn importing_large_codex_rollout_drains_translator_continuations() { + const CALLS: usize = 32; + let tmp = tempfile::tempdir().unwrap(); + let transcript = tmp.path().join("large-codex.jsonl"); + let mut records = vec![ + json!({"timestamp":"2026-01-01T00:00:01Z","type":"session_meta","payload":{"id":"codex-large-import","cwd":"/tmp/demo"}}), + json!({"timestamp":"2026-01-01T00:00:02Z","type":"turn_context","payload":{"model":"gpt-test"}}), + json!({"timestamp":"2026-01-01T00:00:03Z","type":"event_msg","payload":{"type":"task_started","turn_id":"large-turn"}}), + ]; + let prompt = "x".repeat(4 * 1024); + for _ in 0..CALLS { + records.push(json!({"timestamp":"2026-01-01T00:00:04Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"output_text","text":prompt}]}})); + records.push(json!({"timestamp":"2026-01-01T00:00:05Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}})); + records.push(json!({"timestamp":"2026-01-01T00:00:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":1,"output_tokens":1}}}})); + } + records.push(json!({"timestamp":"2026-01-01T00:01:00Z","type":"event_msg","payload":{"type":"task_complete","last_agent_message":"done"}})); + write_jsonl(&transcript, &records); + + let output = tmp.path().join("spans"); + import_transcript( + &transcript, + ImportSource::Codex, + options(&output), + None, + false, + ) + .await + .unwrap(); + + let output_rows = rows(&output.join("codex-large-import.ndjson")); + assert_eq!(inserted(&output_rows, "llm"), CALLS); +} + fn options(output: &std::path::Path) -> ServeOptions { ServeOptions { version: "test".into(), From b3966b3cecc4e5813e0029a8ceb63dd0a7232808 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 18 Aug 2026 19:47:12 +0800 Subject: [PATCH 2/2] fix: bound translator and sink memory --- bt-daemon/src/lib.rs | 30 +++- bt-daemon/src/sink/braintrust.rs | 212 +++++++++++++++++++++---- bt-daemon/src/transcript_import/mod.rs | 45 ++++-- bt-daemon/src/translate/claude.rs | 196 +++++++++++++++++++---- bt-daemon/src/translate/codex.rs | 43 ++--- bt-daemon/src/translate/mod.rs | 1 + bt-daemon/src/translate/opencode.rs | 22 ++- bt-daemon/src/translate/pi.rs | 9 +- bt-daemon/src/translate/recent.rs | 132 +++++++++++++++ bt-daemon/tests/braintrust_sink.rs | 44 +++++ bt-daemon/tests/claude_translator.rs | 97 +++++++++++ 11 files changed, 734 insertions(+), 97 deletions(-) create mode 100644 bt-daemon/src/translate/recent.rs diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 4a7afe7..901732f 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -73,7 +73,7 @@ pub struct ServeArgs { /// Retire a session's in-memory state after this many seconds without /// traffic. A later event rebuilds it from the journal. 0 disables /// retirement. - #[arg(long, default_value_t = 300)] + #[arg(long, default_value_t = 30)] pub session_idle_timeout_secs: u64, } @@ -762,7 +762,14 @@ pub async fn import_transcripts( let entries = tail .poll(true) .with_context(|| format!("import transcript {}", file.display()))?; + let session_ids = entries + .iter() + .map(|entry| entry.session_id.clone()) + .collect::>(); processor.process(entries).await?; + for session_id in session_ids { + processor.finish_session(&session_id).await?; + } } processor.finish().await } @@ -858,6 +865,15 @@ impl ImportProcessor { } Ok(()) } + + async fn finish_session(&mut self, session_id: &str) -> anyhow::Result<()> { + let Some(mut live) = self.sessions.remove(session_id) else { + return Ok(()); + }; + let ops = live.translator.flush(&live.ctx)?; + Self::emit_translator_batches(&mut live, ops).await?; + live.sink.flush().await + } } /// Build a Phase-1 debug [`ServeOptions`]: debug translator registry + a debug @@ -928,6 +944,18 @@ mod tests { args: ImportArgs, } + #[derive(Debug, Parser)] + struct ServeCli { + #[command(flatten)] + args: ServeArgs, + } + + #[test] + fn serve_defaults_to_short_journal_backed_session_retirement() { + let args = ServeCli::try_parse_from(["test"]).unwrap().args; + assert_eq!(args.session_idle_timeout_secs, 30); + } + #[test] fn import_args_accept_multiple_sessions_or_all() { let explicit = ImportCli::try_parse_from([ diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 97a0a38..b5bc08b 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -13,11 +13,12 @@ use super::{Sink, SinkFactory}; use crate::translate::{SpanOp, SpanRow, SpanType}; use crate::wire::{SessionConfig, TraceDestination}; use braintrust_sdk_rust::{ - BraintrustClient, ParentSpanInfo, SpanHandle, SpanLog, SpanObjectType, SpanOrigin, - SpanType as SdkSpanType, DEFAULT_API_URL, DEFAULT_APP_URL, + BraintrustClient, ParentSpanInfo, SpanComponents, SpanHandle, SpanLog, SpanObjectType, + SpanOrigin, SpanType as SdkSpanType, DEFAULT_API_URL, DEFAULT_APP_URL, }; use serde_json::{Map, Value}; -use std::collections::HashMap; +use std::collections::{HashMap, VecDeque}; +use std::io::Write; use std::sync::Arc; use tokio::sync::Mutex as AsyncMutex; @@ -33,14 +34,78 @@ pub struct BraintrustSinkConfig { /// Lazily-built, shared-by-URL client pool. struct ClientCache { - clients: AsyncMutex>>, + state: AsyncMutex, version: String, } +const CLIENT_CACHE_CAPACITY: usize = 16; +const QUEUE_FLUSH_BYTE_BUDGET: usize = 8 * 1024 * 1024; + +#[derive(Default)] +struct ClientCacheState { + clients: HashMap<(String, String), Arc>, + order: VecDeque<(String, String)>, +} + +struct CachedClient { + client: BraintrustClient, + queued_bytes: AsyncMutex, +} + +impl CachedClient { + async fn account_and_flush(&self, bytes: usize) -> anyhow::Result<()> { + let mut queued_bytes = self.queued_bytes.lock().await; + *queued_bytes = queued_bytes.saturating_add(bytes); + if *queued_bytes < QUEUE_FLUSH_BYTE_BUDGET { + return Ok(()); + } + self.client + .flush() + .await + .map_err(|error| anyhow::anyhow!("braintrust flush failed: {error}"))?; + *queued_bytes = 0; + Ok(()) + } + + async fn flush(&self) -> anyhow::Result<()> { + let mut queued_bytes = self.queued_bytes.lock().await; + self.client + .flush() + .await + .map_err(|error| anyhow::anyhow!("braintrust flush failed: {error}"))?; + *queued_bytes = 0; + Ok(()) + } +} + +impl ClientCacheState { + fn touch(&mut self, key: &(String, String)) { + self.order.retain(|candidate| candidate != key); + self.order.push_back(key.clone()); + } + + fn trim_inactive(&mut self) { + while self.clients.len() >= CLIENT_CACHE_CAPACITY { + let Some(index) = self.order.iter().position(|key| { + self.clients + .get(key) + .is_some_and(|client| Arc::strong_count(client) == 1) + }) else { + // Every client is still referenced by a live session; those are + // active working state and cannot be evicted safely. + break; + }; + if let Some(key) = self.order.remove(index) { + self.clients.remove(&key); + } + } + } +} + impl ClientCache { fn new(version: String) -> Self { Self { - clients: AsyncMutex::new(HashMap::new()), + state: AsyncMutex::new(ClientCacheState::default()), version, } } @@ -49,13 +114,14 @@ impl ClientCache { &self, api_url: &str, app_url: &str, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let key = (api_url.to_string(), app_url.to_string()); // Hold the lock across build so two sessions on a new URL don't build // duplicate clients. Build is cheap (skip_login: no network). - let mut map = self.clients.lock().await; - if let Some(c) = map.get(&key) { - return Ok(c.clone()); + let mut state = self.state.lock().await; + if let Some(c) = state.clients.get(&key).cloned() { + state.touch(&key); + return Ok(c); } let client = BraintrustClient::builder() .skip_login(true) @@ -65,8 +131,13 @@ impl ClientCache { .build() .await .map_err(|e| anyhow::anyhow!("braintrust client build failed: {e}"))?; - let arc = Arc::new(client); - map.insert(key, arc.clone()); + let arc = Arc::new(CachedClient { + client, + queued_bytes: AsyncMutex::new(0), + }); + state.trim_inactive(); + state.clients.insert(key.clone(), arc.clone()); + state.touch(&key); Ok(arc) } } @@ -139,7 +210,7 @@ struct BraintrustSink { /// Resolved `(api_url, app_url)` for this session, from its config. urls: Option<(String, String)>, /// The client for `urls`, obtained from the cache on first emit. - client: Option>, + client: Option>, /// Live span handles keyed by deterministic span id, so a later op (e.g. /// setting `end`) merges onto the same row the SDK already knows. open: HashMap>, @@ -180,7 +251,7 @@ impl BraintrustSink { } } - async fn ensure_client(&mut self) -> anyhow::Result> { + async fn ensure_client(&mut self) -> anyhow::Result> { if let Some(c) = &self.client { return Ok(c.clone()); } @@ -227,15 +298,63 @@ impl BraintrustSink { Ok(()) } - fn upsert(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { + 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)?); if let Some(end) = row.end_ms { handle.end_with_time(ms_to_secs(end)); + // SpanHandle retains the complete accumulated input/output. Once a + // span is terminal, use stateless SDK merges for any unusually late + // update instead of pinning that payload for the session lifetime. + self.open.remove(&row.span_id); } Ok(()) } + + fn merge_closed(&self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { + let creds = self + .creds + .as_ref() + .ok_or_else(|| anyhow::anyhow!("session has no credentials/config yet"))?; + let project = self.project(creds); + let components = self.span_components(row, creds, &project).to_str(); + client + .update_span_with_credentials( + creds.token.clone(), + creds.org_id.clone(), + &components, + build_log(row)?, + ) + .map_err(|error| anyhow::anyhow!("braintrust span merge failed: {error}")) + } + + fn span_components(&self, row: &SpanRow, creds: &Creds, project: &str) -> SpanComponents { + let destination = creds + .destination + .as_ref() + .map(|destination| destination_components(destination, project)) + .unwrap_or_else(|| { + let mut args = Map::new(); + args.insert("project_name".into(), Value::String(project.to_string())); + DestinationComponents { + object_type: SpanObjectType::ProjectLogs, + object_id: None, + compute_object_metadata_args: Some(args), + propagated_event: None, + } + }); + SpanComponents { + object_type: destination.object_type, + object_id: destination.object_id, + compute_object_metadata_args: destination.compute_object_metadata_args, + row_id: Some(row.span_id.clone()), + span_id: Some(row.span_id.clone()), + root_span_id: Some(destination_root(creds).unwrap_or_else(|| row.root_span_id.clone())), + span_parents: (!row.parent_span_ids.is_empty()).then(|| row.parent_span_ids.clone()), + propagated_event: destination.propagated_event, + } + } } #[async_trait::async_trait] @@ -292,21 +411,24 @@ impl Sink for BraintrustSink { let client = self.ensure_client().await?; let mut n = 0u64; for op in ops { - let row = match op { - SpanOp::Insert(r) | SpanOp::Merge(r) => r, - }; - self.upsert(&client, row)?; + match op { + SpanOp::Insert(row) => self.update_open(&client.client, row)?, + SpanOp::Merge(row) if self.open.contains_key(&row.span_id) => { + self.update_open(&client.client, row)? + } + SpanOp::Merge(row) => self.merge_closed(&client.client, row)?, + } n += 1; + client + .account_and_flush(serialized_len(op).unwrap_or(0)) + .await?; } Ok(n) } async fn flush(&mut self) -> anyhow::Result<()> { match &self.client { - Some(client) => client - .flush() - .await - .map_err(|e| anyhow::anyhow!("braintrust flush failed: {e}")), + Some(client) => client.flush().await, None => Ok(()), } } @@ -455,14 +577,24 @@ 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 - .iter() - .filter_map(|(k, v)| v.as_f64().map(|f| (k.clone(), f))) - .collect(); - if !hm.is_empty() { - lb = lb.metrics(hm); - } + let mut metrics = row + .metrics + .as_ref() + .and_then(Value::as_object) + .map(|metrics| { + metrics + .iter() + .filter_map(|(k, v)| v.as_f64().map(|f| (k.clone(), f))) + .collect::>() + }) + .unwrap_or_default(); + // Stateless updates do not have a SpanHandle on which to call `end`, so + // carry terminal timing in the merge payload itself. + if let Some(end) = row.end_ms { + metrics.insert("end".into(), 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())); @@ -475,3 +607,23 @@ fn build_log(row: &SpanRow) -> anyhow::Result { lb.build() .map_err(|e| anyhow::anyhow!("span log build failed: {e}")) } + +#[derive(Default)] +struct ByteCounter(usize); + +impl Write for ByteCounter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0 = self.0.saturating_add(bytes.len()); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +fn serialized_len(op: &SpanOp) -> serde_json::Result { + let mut counter = ByteCounter::default(); + serde_json::to_writer(&mut counter, op)?; + Ok(counter.0) +} diff --git a/bt-daemon/src/transcript_import/mod.rs b/bt-daemon/src/transcript_import/mod.rs index e43f5e1..d546fee 100644 --- a/bt-daemon/src/transcript_import/mod.rs +++ b/bt-daemon/src/transcript_import/mod.rs @@ -3,6 +3,7 @@ use crate::ImportSource; use anyhow::{bail, Context}; use serde_json::Value; use std::collections::BTreeMap; +use std::io::BufRead; use std::path::{Path, PathBuf}; mod claude; @@ -191,31 +192,33 @@ pub(crate) fn transcript_envelopes( path: &Path, source: ImportSource, ) -> anyhow::Result> { - let contents = std::fs::read_to_string(path) - .with_context(|| format!("read transcript {}", path.display()))?; + let file = + std::fs::File::open(path).with_context(|| format!("read transcript {}", path.display()))?; + let mut reader = std::io::BufReader::new(file); let mut records = Vec::new(); let mut record_end_offsets = Vec::new(); let mut offset = 0_u64; - for (index, line) in contents.split_inclusive('\n').enumerate() { + let mut line = String::new(); + let mut index = 0_usize; + while reader.read_line(&mut line)? != 0 { + index += 1; offset += line.len() as u64; - if line.trim().is_empty() { - continue; + if !line.trim().is_empty() { + records.push( + serde_json::from_str(&line).with_context(|| { + format!("parse transcript {} line {}", path.display(), index) + })?, + ); + record_end_offsets.push(offset); } - records.push( - serde_json::from_str(line).with_context(|| { - format!("parse transcript {} line {}", path.display(), index + 1) - })?, - ); - record_end_offsets.push(offset); + line.clear(); } if records.is_empty() { bail!("transcript {} is empty", path.display()); } match source { ImportSource::Codex => codex::envelopes(path, &records), - ImportSource::Claude => { - claude::envelopes(path, &records, &record_end_offsets, contents.len() as u64) - } + ImportSource::Claude => claude::envelopes(path, &records, &record_end_offsets, offset), } } @@ -226,6 +229,7 @@ pub(crate) struct TranscriptTail { path: PathBuf, source: ImportSource, state: TailState, + observed_file: Option<(u64, Option)>, } enum TailState { @@ -242,16 +246,27 @@ impl TranscriptTail { ImportSource::Codex => TailState::Codex(codex::Tail::default()), ImportSource::Claude => TailState::Claude(claude::Tail::default()), }, + observed_file: None, } } pub(crate) fn poll(&mut self, finalize: bool) -> anyhow::Result> { + let metadata = match std::fs::metadata(&self.path) { + Ok(metadata) => metadata, + Err(_) if !finalize => return Ok(Vec::new()), + Err(error) => return Err(error.into()), + }; + let len = metadata.len(); + let observed_file = (len, metadata.modified().ok()); + if !finalize && self.observed_file == Some(observed_file) { + return Ok(Vec::new()); + } let events = match transcript_envelopes(&self.path, self.source) { Ok(events) => events, Err(_) if !finalize => return Ok(Vec::new()), Err(error) => return Err(error), }; - let len = std::fs::metadata(&self.path)?.len(); + self.observed_file = Some(observed_file); match &mut self.state { TailState::Codex(state) => state.poll(events, len, finalize), TailState::Claude(state) => state.poll(events, len, finalize), diff --git a/bt-daemon/src/translate/claude.rs b/bt-daemon/src/translate/claude.rs index 979552e..9e5bf63 100644 --- a/bt-daemon/src/translate/claude.rs +++ b/bt-daemon/src/translate/claude.rs @@ -7,11 +7,12 @@ //! transcript that already contains the completed session. use super::git::GitMetadataCache; +use super::recent::RecentSet; 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, HashSet}; +use std::collections::{HashMap, VecDeque}; use std::io::{BufRead, Seek, SeekFrom}; use std::process::Command; use std::sync::Arc; @@ -58,6 +59,20 @@ struct PendingTool { parent_id: String, } +enum PendingHistory { + Main, + Owned(Vec), +} + +struct PendingEmission { + segments: VecDeque>, + history: PendingHistory, + scope: String, + parent: String, + cwd: Option, + clear_after: bool, +} + struct ClaudeTranslator { session_id: String, session_span_id: String, @@ -71,11 +86,12 @@ struct ClaudeTranslator { main_transcript: Option, transcripts: HashMap, main_history: Vec, - emitted_requests: HashSet, - emitted_tools: HashSet, + emitted_requests: RecentSet, + emitted_tools: RecentSet, pending_tools: HashMap, subagents: HashMap, pending_skills: Vec, + pending_emission: Option, claude_version: Option, claude_version_logged: bool, git: Arc, @@ -99,11 +115,12 @@ impl ClaudeTranslator { main_transcript: None, transcripts: HashMap::new(), main_history: Vec::new(), - emitted_requests: HashSet::new(), - emitted_tools: HashSet::new(), + emitted_requests: RecentSet::default(), + emitted_tools: RecentSet::default(), pending_tools: HashMap::new(), subagents: HashMap::new(), pending_skills: Vec::new(), + pending_emission: None, claude_version: None, claude_version_logged: false, git, @@ -414,7 +431,15 @@ impl ClaudeTranslator { .buffered .extend(read_event_records(event, &path, &mut cursor.offset)); let records = std::mem::take(&mut cursor.buffered); - self.emit_transcript(&records, &format!("subagent:{agent_id}"), &parent, ops); + self.queue_transcript( + records, + PendingHistory::Owned(Vec::new()), + format!("subagent:{agent_id}"), + parent.clone(), + self.current_cwd.clone(), + ops, + ); + self.transcripts.remove(&path); } self.close_pending_tools( &parent, @@ -423,12 +448,13 @@ impl ClaudeTranslator { ops, ); ops.push(SpanOp::Merge(SpanRow { - span_id: parent, + span_id: parent.clone(), root_span_id: self.root_span_id.clone(), end_ms: Some(event.ts_ms), output: event.payload.get("last_assistant_message").cloned(), ..Default::default() })); + self.subagents.remove(&agent_id); } fn emit_main(&mut self, parent: &str, ops: &mut Vec) { @@ -440,25 +466,83 @@ impl ClaudeTranslator { .get_mut(&path) .map(|cursor| std::mem::take(&mut cursor.buffered)) .unwrap_or_default(); - let parsed = parse_transcript(&records, std::mem::take(&mut self.main_history)); - // Moved, not cloned: this used to deep-copy the whole conversation on - // every turn. The history is never trimmed — it is the model input - // these spans report, so dropping any of it would silently corrupt the - // trace. Its footprint is the active session's own conversation, and - // retiring an idle session releases it. - self.main_history = parsed.history; - self.emit_parsed(parsed.calls, parsed.tools, "main", parent, ops); + self.queue_transcript( + records, + PendingHistory::Main, + "main".into(), + parent.into(), + self.current_cwd.clone(), + ops, + ); } - fn emit_transcript( + fn queue_transcript( &mut self, - records: &[Value], - scope: &str, - parent: &str, + records: Vec, + history: PendingHistory, + scope: String, + parent: String, + cwd: Option, ops: &mut Vec, ) { - let parsed = parse_transcript(records, Vec::new()); - self.emit_parsed(parsed.calls, parsed.tools, scope, parent, ops); + debug_assert!(self.pending_emission.is_none()); + let pending = PendingEmission { + segments: transcript_segments(records), + history, + scope, + parent, + cwd, + clear_after: false, + }; + self.pending_emission = self.emit_next_pending(pending, ops); + } + + fn emit_next_pending( + &mut self, + mut pending: PendingEmission, + ops: &mut Vec, + ) -> Option { + let Some(records) = pending.segments.pop_front() else { + if pending.clear_after { + self.release_terminal_state(); + } + return None; + }; + let history = match &mut pending.history { + PendingHistory::Main => std::mem::take(&mut self.main_history), + PendingHistory::Owned(history) => std::mem::take(history), + }; + let ParsedTranscript { + calls, + tools, + history, + } = parse_transcript(&records, history); + match &mut pending.history { + PendingHistory::Main => self.main_history = history, + PendingHistory::Owned(pending_history) => *pending_history = history, + } + let op_start = ops.len(); + self.emit_parsed(calls, tools, &pending.scope, &pending.parent, ops); + self.git + .enrich_rows(pending.cwd.as_deref(), &mut ops[op_start..]); + if pending.segments.is_empty() { + if pending.clear_after { + self.release_terminal_state(); + } + None + } else { + Some(pending) + } + } + + fn release_terminal_state(&mut self) { + self.main_history.clear(); + self.transcripts.clear(); + self.subagents.clear(); + self.pending_tools.clear(); + self.pending_skills.clear(); + self.emitted_requests.clear(); + self.emitted_tools.clear(); } fn emit_parsed( @@ -493,7 +577,6 @@ impl ClaudeTranslator { } fn flush_previous_turn_rows(&mut self, ops: &mut Vec) { - let op_start = ops.len(); let (Some(path), Some(parent)) = (self.main_transcript.clone(), self.last_turn_id.clone()) else { return; @@ -508,11 +591,14 @@ impl ClaudeTranslator { .unwrap_or(cursor.buffered.len()); let current_turn_rows = cursor.buffered.split_off(split); let previous_rows = std::mem::replace(&mut cursor.buffered, current_turn_rows); - let parsed = parse_transcript(&previous_rows, std::mem::take(&mut self.main_history)); - self.main_history = parsed.history; - self.emit_parsed(parsed.calls, parsed.tools, "main", &parent, ops); - self.git - .enrich_rows(self.last_turn_cwd.as_deref(), &mut ops[op_start..]); + self.queue_transcript( + previous_rows, + PendingHistory::Main, + "main".into(), + parent, + self.last_turn_cwd.clone(), + ops, + ); } fn stop_turn(&mut self, event: &Envelope, error: Option, ops: &mut Vec) { @@ -604,11 +690,20 @@ impl ClaudeTranslator { ..Default::default() })); } + if let Some(pending) = &mut self.pending_emission { + pending.clear_after = true; + } else { + self.release_terminal_state(); + } } } impl AgentTranslator for ClaudeTranslator { fn handle(&mut self, event: &Envelope, ctx: &SessionCtx) -> anyhow::Result> { + anyhow::ensure!( + self.pending_emission.is_none(), + "Claude translator has pending catch-up work; drain it before handling another event" + ); let mut ops = Vec::new(); if let Some(cwd) = string_field(&event.payload, "cwd") { self.current_cwd = Some(cwd); @@ -674,11 +769,58 @@ impl AgentTranslator for ClaudeTranslator { Ok(ops) } + fn drain_pending(&mut self, _ctx: &SessionCtx) -> anyhow::Result>> { + let Some(pending) = self.pending_emission.take() else { + return Ok(None); + }; + let mut ops = Vec::new(); + self.pending_emission = self.emit_next_pending(pending, &mut ops); + Ok(Some(ops)) + } + fn flush(&mut self, _ctx: &SessionCtx) -> anyhow::Result> { Ok(Vec::new()) } } +fn transcript_segments(records: Vec) -> VecDeque> { + let mut segments = VecDeque::new(); + let mut current = Vec::new(); + let mut request_id: Option = None; + for record in records { + let next_request_id = assistant_request_id(&record); + if next_request_id + .as_ref() + .zip(request_id.as_ref()) + .is_some_and(|(next, current)| next != current) + && !current.is_empty() + { + segments.push_back(std::mem::take(&mut current)); + } + if next_request_id.is_some() { + request_id = next_request_id; + } + current.push(record); + } + if !current.is_empty() { + segments.push_back(current); + } + segments +} + +fn assistant_request_id(record: &Value) -> Option { + (record.get("type").and_then(Value::as_str) == Some("assistant")) + .then(|| { + record + .get("message") + .and_then(|message| string_field(message, "id")) + .or_else(|| string_field(record, "requestId")) + .or_else(|| string_field(record, "uuid")) + }) + .flatten() + .filter(|id| !id.is_empty()) +} + struct ParsedTranscript { calls: Vec, tools: Vec, diff --git a/bt-daemon/src/translate/codex.rs b/bt-daemon/src/translate/codex.rs index bc94451..de53562 100644 --- a/bt-daemon/src/translate/codex.rs +++ b/bt-daemon/src/translate/codex.rs @@ -18,12 +18,13 @@ //! newer turn has already started. use super::git::GitMetadataCache; +use super::recent::{RecentMap, RecentSet}; use super::{AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory}; use crate::ids; use crate::wire::Envelope; use regex::Regex; use serde_json::{json, Map, Value}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::path::Path; use std::sync::{Arc, OnceLock}; @@ -62,10 +63,10 @@ impl TranslatorFactory for CodexTranslatorFactory { main_path: None, // The main scope is created lazily once we learn its transcript path. scopes: HashMap::new(), - spawn_turn_by_call_id: HashMap::new(), - spawn_turn_by_agent_id: HashMap::new(), - compaction_trigger_by_turn: HashMap::new(), - compaction_spans: HashSet::new(), + spawn_turn_by_call_id: RecentMap::default(), + spawn_turn_by_agent_id: RecentMap::default(), + compaction_trigger_by_turn: RecentMap::default(), + compaction_spans: RecentSet::default(), pending: None, git: self.git.clone(), }) @@ -163,10 +164,10 @@ struct CodexTranslator { additional_metadata: Map, main_path: Option, scopes: HashMap, - spawn_turn_by_call_id: HashMap, - spawn_turn_by_agent_id: HashMap, - compaction_trigger_by_turn: HashMap, - compaction_spans: HashSet, + spawn_turn_by_call_id: RecentMap, + spawn_turn_by_agent_id: RecentMap, + compaction_trigger_by_turn: RecentMap, + compaction_spans: RecentSet, pending: Option, git: Arc, } @@ -322,7 +323,8 @@ impl CodexTranslator { self.compaction_trigger_by_turn .insert(turn_id.clone(), trigger.clone()); // Back-fill onto an already-built compaction span. - if self.compaction_spans.contains(&turn_id) { + if self.compaction_spans.remove(&turn_id) { + self.compaction_trigger_by_turn.remove(&turn_id); let span_id = ids::span_id(&self.session_id, &format!("turn:{turn_id}")); ops.push(SpanOp::Merge(SpanRow { span_id, @@ -348,9 +350,8 @@ impl CodexTranslator { _ => None, }); let Some(agent_id) = agent_id else { return }; - if let Some(turn_span) = self.spawn_turn_by_call_id.get(&call_id) { - self.spawn_turn_by_agent_id - .insert(agent_id, turn_span.clone()); + if let Some(turn_span) = self.spawn_turn_by_call_id.remove(&call_id) { + self.spawn_turn_by_agent_id.insert(agent_id, turn_span); } } @@ -366,8 +367,7 @@ impl CodexTranslator { } let parent = self .spawn_turn_by_agent_id - .get(&agent_id) - .cloned() + .remove(&agent_id) .unwrap_or_else(|| self.root_span_id.clone()); let subagent_root = ids::span_id(&self.session_id, &format!("subagent:{agent_id}")); let mut scope = Scope::new(&path, ScopeKind::Subagent, subagent_root); @@ -1082,8 +1082,10 @@ impl CodexTranslator { .get("replacement_history") .and_then(Value::as_array) .cloned(); - let trigger = self.compaction_trigger_by_turn.get(&turn_id).cloned(); - self.compaction_spans.insert(turn_id.clone()); + let trigger = self.compaction_trigger_by_turn.remove(&turn_id); + if trigger.is_none() { + self.compaction_spans.insert(turn_id.clone()); + } // Relabel the turn as a compaction span. ops.push(SpanOp::Merge(SpanRow { @@ -1121,6 +1123,9 @@ impl CodexTranslator { ..Default::default() })); if let Some(replacement) = replacement { + // Native compaction is a semantic memory barrier: future requests + // use only the replacement context, so the pre-compaction Values + // can be dropped as soon as their one compaction span is emitted. scope.conversation_history = replacement; } let _ = (turn_span, name); @@ -1181,7 +1186,9 @@ impl CodexTranslator { ..Default::default() })); } - self.scopes.insert(path.to_string(), scope); + // SubagentStop is terminal for this transcript. Its durable rollout is + // still on disk and journal replay can rebuild it, so retaining the + // closed scope would only pin its entire conversation history. } /// Close any open llm/tool/turn in `scope` (used on subagent stop + flush). diff --git a/bt-daemon/src/translate/mod.rs b/bt-daemon/src/translate/mod.rs index 8e6e10a..fc5d2ee 100644 --- a/bt-daemon/src/translate/mod.rs +++ b/bt-daemon/src/translate/mod.rs @@ -13,6 +13,7 @@ mod debug; mod git; mod opencode; mod pi; +mod recent; pub use claude::ClaudeTranslatorFactory; pub use codex::CodexTranslatorFactory; diff --git a/bt-daemon/src/translate/opencode.rs b/bt-daemon/src/translate/opencode.rs index 62b2b2e..283883b 100644 --- a/bt-daemon/src/translate/opencode.rs +++ b/bt-daemon/src/translate/opencode.rs @@ -2,6 +2,7 @@ //! this module owns span construction, correlation, and recovery. use super::git::GitMetadataCache; +use super::recent::RecentSet; use super::{AgentTranslator, SessionCtx, SpanOp, SpanRow, SpanType, TranslatorFactory}; use crate::ids; use crate::wire::Envelope; @@ -54,7 +55,7 @@ struct NativeSession { tool_errors: HashMap, tool_message_ids: HashMap, denied_tools: HashSet, - completed_messages: HashSet, + completed_messages: RecentSet, } struct OpenCodeTranslator { @@ -224,6 +225,14 @@ impl OpenCodeTranslator { ..Default::default() })); } + // Message fragments are only needed until their completed LLM row is + // emitted. A new turn is also a hard boundary for any incomplete native + // fragments left behind by a missing completion event. + state.output_parts.clear(); + state.reasoning_parts.clear(); + state.tool_calls.clear(); + state.tool_message_ids.clear(); + state.completed_messages.clear(); state.turn_number += 1; let turn_id = ids::span_id( &self.daemon_session_id, @@ -379,13 +388,16 @@ impl OpenCodeTranslator { let prompt = num(info, "/tokens/input") + cache_read + cache_write; let completion = num(info, "/tokens/output"); let reasoning = num(info, "/tokens/reasoning"); - let mut assistant = json!({"role":"assistant","content":state.output_parts.get(mid).cloned().unwrap_or_default()}); - if let Some(calls) = state.tool_calls.get(mid) { - assistant["tool_calls"] = Value::Array(calls.clone()) + let mut assistant = json!({"role":"assistant","content":state.output_parts.remove(mid).unwrap_or_default()}); + if let Some(calls) = state.tool_calls.remove(mid) { + assistant["tool_calls"] = Value::Array(calls) } - if let Some(reason) = state.reasoning_parts.get(mid) { + if let Some(reason) = state.reasoning_parts.remove(mid) { assistant["reasoning"] = json!([{"id":"reasoning","content":reason}]) } + state + .tool_message_ids + .retain(|_, message_id| message_id != mid); let mut input = Vec::new(); if let Some(system) = &state.system_prompt { input.push(json!({"role":"system","content":system})) diff --git a/bt-daemon/src/translate/pi.rs b/bt-daemon/src/translate/pi.rs index 5129aa0..f0b746f 100644 --- a/bt-daemon/src/translate/pi.rs +++ b/bt-daemon/src/translate/pi.rs @@ -472,7 +472,12 @@ impl PiTranslator { })] } fn close_turn(&mut self, ts: i64, error: Option) -> Vec { - let Some((id, _)) = self.turn.take() else { + let turn = self.turn.take(); + // Native context/tool payloads are only correlation state for the active + // turn. The journal can rebuild them if a retired session later resumes. + self.pending_llms.clear(); + self.tools.clear(); + let Some((id, _)) = turn else { return vec![]; }; vec![SpanOp::Merge(SpanRow { @@ -485,6 +490,8 @@ impl PiTranslator { } fn close_root(&mut self, ts: i64) -> SpanOp { self.opened = false; + self.compaction = None; + self.branch_summary = None; SpanOp::Merge(SpanRow { span_id: self.root_span_id.clone(), root_span_id: self.effective_root_span_id.clone(), diff --git a/bt-daemon/src/translate/recent.rs b/bt-daemon/src/translate/recent.rs new file mode 100644 index 0000000..2339fc4 --- /dev/null +++ b/bt-daemon/src/translate/recent.rs @@ -0,0 +1,132 @@ +//! Small bounded caches for replay/deduplication state. +//! +//! Journals and native transcripts are the durable source of truth. Translators +//! only need a recent window of completed identifiers to absorb duplicate or +//! slightly reordered events; keeping every identifier for the whole session +//! makes payload-free bookkeeping grow without bound. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::hash::Hash; + +pub(super) const RECENT_ID_CAPACITY: usize = 4_096; + +pub(super) struct RecentSet { + values: HashSet, + order: VecDeque, + capacity: usize, +} + +impl Default for RecentSet { + fn default() -> Self { + Self { + values: HashSet::new(), + order: VecDeque::new(), + capacity: RECENT_ID_CAPACITY, + } + } +} + +impl RecentSet { + pub(super) fn contains(&self, value: &K) -> bool { + self.values.contains(value) + } + + pub(super) fn insert(&mut self, value: K) -> bool { + if !self.values.insert(value.clone()) { + return false; + } + self.order.push_back(value); + while self.values.len() > self.capacity { + if let Some(oldest) = self.order.pop_front() { + self.values.remove(&oldest); + } + } + true + } + + pub(super) fn remove(&mut self, value: &K) -> bool { + let removed = self.values.remove(value); + if removed { + self.order.retain(|candidate| candidate != value); + } + removed + } + + pub(super) fn clear(&mut self) { + self.values.clear(); + self.order.clear(); + } +} + +pub(super) struct RecentMap { + values: HashMap, + order: VecDeque, + capacity: usize, +} + +impl Default for RecentMap { + fn default() -> Self { + Self { + values: HashMap::new(), + order: VecDeque::new(), + capacity: RECENT_ID_CAPACITY, + } + } +} + +impl RecentMap { + #[cfg(test)] + pub(super) fn get(&self, key: &K) -> Option<&V> { + self.values.get(key) + } + + pub(super) fn insert(&mut self, key: K, value: V) -> Option { + self.order.retain(|candidate| candidate != &key); + self.order.push_back(key.clone()); + let previous = self.values.insert(key, value); + while self.values.len() > self.capacity { + if let Some(oldest) = self.order.pop_front() { + self.values.remove(&oldest); + } + } + previous + } + + pub(super) fn remove(&mut self, key: &K) -> Option { + let removed = self.values.remove(key); + if removed.is_some() { + self.order.retain(|candidate| candidate != key); + } + removed + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recent_collections_evict_oldest_entries() { + let mut set = RecentSet { + capacity: 2, + ..Default::default() + }; + assert!(set.insert("a")); + assert!(set.insert("b")); + assert!(set.insert("c")); + assert!(!set.contains(&"a")); + assert!(set.contains(&"b")); + assert!(set.contains(&"c")); + + let mut map = RecentMap { + capacity: 2, + ..Default::default() + }; + map.insert("a", 1); + map.insert("b", 2); + map.insert("c", 3); + assert!(map.get(&"a").is_none()); + assert_eq!(map.get(&"b"), Some(&2)); + assert_eq!(map.get(&"c"), Some(&3)); + } +} diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index d3aeee6..4cc622a 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -176,6 +176,50 @@ async fn merge_with_empty_name_does_not_clobber_the_original_name() { ); } +/// Completed handles must be releasable without changing the update contract: +/// an update that arrives after terminal insertion is emitted as a stateless +/// merge against the same deterministic span id. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn late_merge_updates_a_completed_span_without_an_open_handle() { + 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-late", "codex", None).unwrap(); + sink.configure(&session_config(&base)); + + sink.emit(&[SpanOp::Insert(row( + "finished", + "finished", + &[], + "original name", + SpanType::Task, + 1, + Some(2), + ))]) + .await + .unwrap(); + let mut late = SpanRow { + span_id: "finished".into(), + root_span_id: "finished".into(), + output: Some(json!({"status":"late"})), + ..Default::default() + }; + late.name.clear(); + sink.emit(&[SpanOp::Merge(late)]).await.unwrap(); + sink.flush().await.unwrap(); + + let bodies = logs3_bodies(&server).await; + assert!( + bodies.contains("original name"), + "initial row absent: {bodies}" + ); + assert!(bodies.contains("late"), "late merge absent: {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/bt-daemon/tests/claude_translator.rs b/bt-daemon/tests/claude_translator.rs index b8e4710..e3956b9 100644 --- a/bt-daemon/tests/claude_translator.rs +++ b/bt-daemon/tests/claude_translator.rs @@ -82,8 +82,14 @@ fn replay_from(name: &str, source: Source) -> Vec { config: None, }; ops.extend(translator.handle(&env, &ctx).unwrap()); + while let Some(batch) = translator.drain_pending(&ctx).unwrap() { + ops.extend(batch); + } } ops.extend(translator.flush(&ctx).unwrap()); + while let Some(batch) = translator.drain_pending(&ctx).unwrap() { + ops.extend(batch); + } ops } @@ -610,3 +616,94 @@ fn claude_groups_streamed_rows_and_reads_late_final_output_at_session_end() { json!("done") ); } + +#[test] +fn claude_large_catch_up_emits_one_historical_snapshot_per_batch() { + const CALLS: usize = 24; + const MESSAGE_BYTES: usize = 64 * 1024; + let transcript = tempfile::NamedTempFile::new().unwrap(); + let transcript_path = transcript.path().to_str().unwrap(); + let mut records = Vec::with_capacity(CALLS * 2); + for index in 0..CALLS { + records.push(json!({ + "type": "user", + "timestamp": "2026-07-28T16:00:00Z", + "message": {"role":"user", "content": "x".repeat(MESSAGE_BYTES)} + })); + records.push(json!({ + "type": "assistant", + "timestamp": "2026-07-28T16:00:01Z", + "message": { + "id": format!("request-{index}"), + "model": "claude-test", + "role": "assistant", + "content": [{"type":"text", "text":format!("answer-{index}")}], + "usage": {"input_tokens":1,"output_tokens":1} + } + })); + } + let contents = records + .iter() + .map(Value::to_string) + .collect::>() + .join("\n"); + std::fs::write(transcript.path(), &contents).unwrap(); + + let registry = Registry::default_agents(); + let mut translator = registry.create("claude-code", "bounded"); + let ctx = SessionCtx { + session_id: "bounded".into(), + config: None, + }; + let envelope = |name: &str, payload: Value| Envelope { + source: "claude-code".into(), + source_version: None, + plugin_version: None, + session_id: "bounded".into(), + event: name.into(), + ts_ms: 2_000_000_000_000, + managed_run_id: None, + payload, + route: None, + config: None, + }; + translator + .handle( + &envelope( + "UserPromptSubmit", + json!({"session_id":"bounded","prompt":"go"}), + ), + &ctx, + ) + .unwrap(); + let mut batch = translator + .handle( + &envelope( + "Stop", + json!({ + "session_id":"bounded", + "transcript_path":transcript_path, + "_bt_transcript_snapshot":{"path":transcript_path,"contents":contents} + }), + ), + &ctx, + ) + .unwrap(); + let mut llm_count = 0; + loop { + let batch_llms = batch + .iter() + .filter(|op| matches!(op, SpanOp::Insert(row) if row.span_type == SpanType::Llm)) + .count(); + assert!( + batch_llms <= 1, + "catch-up batch materialized {batch_llms} LLM inputs" + ); + llm_count += batch_llms; + let Some(next) = translator.drain_pending(&ctx).unwrap() else { + break; + }; + batch = next; + } + assert_eq!(llm_count, CALLS); +}