use crate::language_model_message::LanguageModelPrompt;
use crate::options::CallOptions;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet, VecDeque};
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, SyncSender, sync_channel};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use ts_rs::TS;
pub const RECORDING_SCHEMA: u32 = 1;
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct Recording {
pub schema: u32,
pub call_id: String,
pub recorded_at: String,
pub input: InputRecord,
pub provider: ProviderRecord,
pub exchanges: Vec<HttpExchange>,
pub outcome: OutcomeRecord,
#[serde(default)]
pub complete: bool,
#[serde(default)]
pub transport_closed: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub step: Option<u32>,
}
impl Recording {
#[must_use]
pub fn new(call_id: &str, input: InputRecord, provider: ProviderRecord) -> Self {
Self {
schema: RECORDING_SCHEMA,
call_id: call_id.to_string(),
recorded_at: iso8601_now(),
input,
provider,
exchanges: Vec::new(),
outcome: OutcomeRecord::default(),
complete: false,
transport_closed: false,
session_id: None,
step: None,
}
}
fn ready(&self) -> bool {
self.transport_closed
&& self.outcome.status != OutcomeStatus::Pending
&& !self.input.prompt.is_empty()
&& self.exchanges.iter().all(|e| e.finalized)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct InputRecord {
pub prompt: LanguageModelPrompt,
pub options: serde_json::Value,
}
impl InputRecord {
#[must_use]
pub fn from_call_options(options: &CallOptions) -> Self {
let value = serde_json::to_value(options).unwrap_or(serde_json::Value::Null);
Self {
prompt: options.prompt.clone(),
options: redact_json(value),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct ProviderRecord {
pub provider: String,
pub model_id: String,
pub base_url: Option<String>,
pub api_key_source: String,
pub profile: Option<serde_json::Value>,
pub provider_options: Option<serde_json::Value>,
}
impl ProviderRecord {
#[must_use]
pub fn minimal(provider: &str, model_id: &str) -> Self {
Self {
provider: provider.to_string(),
model_id: model_id.to_string(),
base_url: None,
api_key_source: "unknown".to_string(),
profile: None,
provider_options: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct HttpExchange {
pub attempt: u32,
pub request: HttpRecord,
pub response: Option<ResponseRecord>,
pub timing: TimingRecord,
pub error: Option<String>,
#[serde(default = "default_finalized")]
pub finalized: bool,
}
fn default_finalized() -> bool {
true
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct HttpRecord {
pub method: String,
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct ResponseRecord {
pub status: u16,
pub headers: Vec<(String, String)>,
pub body: Option<String>,
pub stream_chunks: Option<usize>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(type = "number | null")]
pub ttfb_ms: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct TimingRecord {
#[ts(type = "number")]
pub latency_ms: u64,
#[ts(type = "number | null")]
pub ttfb_ms: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export)]
pub enum OutcomeStatus {
#[default]
Pending,
Success,
Error,
Incomplete,
Cancelled,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct OutcomeRecord {
pub status: OutcomeStatus,
pub finish_reason: Option<String>,
pub error: Option<String>,
pub usage: Option<serde_json::Value>,
}
impl OutcomeRecord {
#[must_use]
pub fn from_generate_result(r: &crate::result::GenerateResult) -> Self {
Self {
status: OutcomeStatus::Success,
finish_reason: serde_json::to_value(r.finish_reason.unified)
.ok()
.and_then(|v| v.as_str().map(std::string::ToString::to_string)),
error: None,
usage: serde_json::to_value(&r.usage).ok(),
}
}
#[must_use]
pub fn from_error(e: &crate::error::AiMuxError) -> Self {
Self {
status: OutcomeStatus::Error,
finish_reason: None,
error: Some(e.to_string()),
usage: None,
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum RecordingError {
#[error("recording init failed for {path}: {source}")]
Init {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to open recording file {path}: {source}")]
OpenFile {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to spawn recording writer thread: {source}")]
Spawn {
#[source]
source: std::io::Error,
},
#[error("recording writer thread unavailable")]
WriterGone,
#[error("recording flush timed out")]
FlushTimeout,
#[error("recording write failed: {0}")]
Write(String),
}
pub trait Recorder: Send + Sync {
fn record_input(&self, call_id: &str, options: &CallOptions, provider: &str, model_id: &str);
fn record_session(&self, _call_id: &str, _session_id: &str, _step: u32) {}
fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord);
fn record_exchange(&self, call_id: &str, exchange: &HttpExchange);
fn record_exchange_update(
&self,
call_id: &str,
attempt: u32,
response: &ResponseRecord,
error: Option<String>,
);
fn record_outcome(&self, call_id: &str, outcome: &OutcomeRecord);
fn record_transport_closed(&self, call_id: &str) {
let _ = call_id;
}
fn flush(&self);
fn try_flush(&self) -> Result<(), RecordingError> {
Ok(())
}
}
static RECORDER: RwLock<Option<Arc<dyn Recorder>>> = RwLock::new(None);
#[derive(Clone)]
pub struct RecordingContext {
pub call_id: String,
pub recorder: Arc<dyn Recorder>,
}
impl std::fmt::Debug for RecordingContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RecordingContext")
.field("call_id", &self.call_id)
.field("recorder", &"..")
.finish()
}
}
pub fn init_recording(recorder: Option<Arc<dyn Recorder>>) {
if let Ok(mut g) = RECORDER.write() {
*g = recorder;
}
}
#[must_use]
pub fn init_recording_from_env() -> bool {
if std::env::var("AIMUX_RECORD").as_deref() != Ok("1") {
return false;
}
let dir = std::env::var("AIMUX_RECORD_DIR").unwrap_or_else(|_| "./recordings".to_string());
init_recording(Some(Arc::new(JsonlRecorder::new(dir))));
true
}
pub fn context(call_id: impl Into<String>) -> Option<RecordingContext> {
recorder().map(|recorder| RecordingContext {
call_id: call_id.into(),
recorder,
})
}
pub fn recorder() -> Option<Arc<dyn Recorder>> {
RECORDER.read().ok()?.clone()
}
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
pub fn new_call_id() -> String {
let ns = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("call-{ns}-{}", CALL_SEQ.fetch_add(1, Ordering::Relaxed))
}
#[must_use]
pub fn is_sensitive_key(name: &str) -> bool {
let n = name.to_ascii_lowercase();
n == "cookie"
|| n == "set-cookie"
|| n == "key"
|| n == "x-amz-security-token"
|| n.contains("authorization")
|| n.contains("api-key")
|| n.contains("api_key")
|| n.contains("apikey")
}
pub(crate) fn redact_json(value: serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Object(map) => {
let mut out = serde_json::Map::with_capacity(map.len());
for (k, v) in map {
if is_sensitive_key(&k) {
out.insert(k, serde_json::Value::String("[REDACTED]".into()));
} else {
out.insert(k, redact_json(v));
}
}
serde_json::Value::Object(out)
}
serde_json::Value::Array(items) => {
serde_json::Value::Array(items.into_iter().map(redact_json).collect())
}
other => other,
}
}
enum RecordEvent {
Input {
call_id: String,
input: InputRecord,
provider: ProviderRecord,
},
Session {
call_id: String,
session_id: String,
step: u32,
},
Provider {
call_id: String,
provider: ProviderRecord,
},
Exchange {
call_id: String,
exchange: HttpExchange,
},
ExchangeUpdate {
call_id: String,
attempt: u32,
response: ResponseRecord,
error: Option<String>,
},
Outcome {
call_id: String,
outcome: OutcomeRecord,
},
TransportClosed { call_id: String },
Flush {
ack: SyncSender<Result<(), RecordingError>>,
},
}
const JSONL_CHANNEL_CAPACITY: usize = 1024;
pub struct JsonlRecorder {
tx: Option<SyncSender<RecordEvent>>,
dir: PathBuf,
thread: Option<std::thread::JoinHandle<()>>,
dropped: AtomicU64,
inconsistent: Arc<Mutex<HashSet<String>>>,
write_error: Arc<Mutex<Option<String>>>,
}
impl JsonlRecorder {
pub fn try_new(dir: impl Into<PathBuf>) -> Result<Self, RecordingError> {
let dir = dir.into();
std::fs::create_dir_all(&dir).map_err(|source| RecordingError::Init {
path: dir.clone(),
source,
})?;
let path = dir.join("recordings.jsonl");
let file = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&path)
.map_err(|source| RecordingError::OpenFile {
path: path.clone(),
source,
})?;
let w = BufWriter::new(file);
let (tx, rx) = sync_channel::<RecordEvent>(JSONL_CHANNEL_CAPACITY);
let inconsistent = Arc::new(Mutex::new(HashSet::new()));
let writer_inconsistent = inconsistent.clone();
let write_error = Arc::new(Mutex::new(None));
let writer_write_error = write_error.clone();
let handle = std::thread::Builder::new()
.name("aimux-recording".into())
.spawn(move || writer_loop(rx, w, writer_inconsistent, writer_write_error))
.map_err(|source| RecordingError::Spawn { source })?;
Ok(Self {
tx: Some(tx),
dir,
thread: Some(handle),
dropped: AtomicU64::new(0),
inconsistent,
write_error,
})
}
pub fn new(dir: impl Into<PathBuf>) -> Self {
let dir = dir.into();
Self::try_new(dir.clone()).unwrap_or_else(|_| Self::disabled(dir))
}
fn disabled(dir: PathBuf) -> Self {
Self {
tx: None,
dir,
thread: None,
dropped: AtomicU64::new(0),
inconsistent: Arc::new(Mutex::new(HashSet::new())),
write_error: Arc::new(Mutex::new(None)),
}
}
fn send_ev(&self, ev: RecordEvent) {
if let Some(tx) = &self.tx {
send_or_drop(tx, ev, &self.dropped);
}
}
pub fn dropped_count(&self) -> u64 {
self.dropped.load(Ordering::Relaxed)
}
pub fn inconsistent_call_ids(&self) -> Vec<String> {
let mut v: Vec<String> = self.inconsistent.lock().unwrap().iter().cloned().collect();
v.sort();
v
}
pub fn try_flush(&self) -> Result<(), RecordingError> {
let Some(tx) = &self.tx else {
return Err(self.write_error.lock().unwrap().as_ref().map_or_else(
|| RecordingError::WriterGone,
|msg| RecordingError::Write(msg.clone()),
));
};
let (ack_tx, ack_rx) = sync_channel::<Result<(), RecordingError>>(0);
tx.send(RecordEvent::Flush { ack: ack_tx })
.map_err(|_| RecordingError::WriterGone)?;
match ack_rx.recv_timeout(Duration::from_secs(30)) {
Ok(status) => status,
Err(_) => Err(RecordingError::FlushTimeout),
}
}
pub fn last_write_error(&self) -> Option<String> {
self.write_error.lock().unwrap().clone()
}
pub fn path(&self) -> PathBuf {
self.dir.join("recordings.jsonl")
}
}
impl Drop for JsonlRecorder {
fn drop(&mut self) {
if let Some(tx) = self.tx.take() {
drop(tx);
}
if let Some(handle) = self.thread.take() {
let _ = handle.join();
}
}
}
impl Recorder for JsonlRecorder {
fn record_input(&self, call_id: &str, options: &CallOptions, provider: &str, model_id: &str) {
self.send_ev(RecordEvent::Input {
call_id: call_id.to_string(),
input: InputRecord::from_call_options(options),
provider: ProviderRecord::minimal(provider, model_id),
});
}
fn record_session(&self, call_id: &str, session_id: &str, step: u32) {
self.send_ev(RecordEvent::Session {
call_id: call_id.to_string(),
session_id: session_id.to_string(),
step,
});
}
fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord) {
let mut snap = snapshot.clone();
snap.provider_options = snap.provider_options.take().map(redact_json);
snap.profile = snap.profile.take().map(redact_json);
self.send_ev(RecordEvent::Provider {
call_id: call_id.to_string(),
provider: snap,
});
}
fn record_exchange(&self, call_id: &str, exchange: &HttpExchange) {
self.send_ev(RecordEvent::Exchange {
call_id: call_id.to_string(),
exchange: exchange.clone(),
});
}
fn record_exchange_update(
&self,
call_id: &str,
attempt: u32,
response: &ResponseRecord,
error: Option<String>,
) {
self.send_ev(RecordEvent::ExchangeUpdate {
call_id: call_id.to_string(),
attempt,
response: response.clone(),
error,
});
}
fn record_outcome(&self, call_id: &str, outcome: &OutcomeRecord) {
self.send_ev(RecordEvent::Outcome {
call_id: call_id.to_string(),
outcome: outcome.clone(),
});
}
fn record_transport_closed(&self, call_id: &str) {
self.send_ev(RecordEvent::TransportClosed {
call_id: call_id.to_string(),
});
}
fn flush(&self) {
let _ = JsonlRecorder::try_flush(self);
}
fn try_flush(&self) -> Result<(), RecordingError> {
JsonlRecorder::try_flush(self)
}
}
fn entry_or_init<'a>(
pending: &'a mut HashMap<String, Recording>,
call_id: &str,
) -> &'a mut Recording {
pending.entry(call_id.to_string()).or_insert_with(|| {
Recording::new(
call_id,
InputRecord {
prompt: Vec::new(),
options: serde_json::Value::Null,
},
ProviderRecord::minimal("", ""),
)
})
}
fn writer_loop(
rx: Receiver<RecordEvent>,
mut w: impl Write,
inconsistent: Arc<Mutex<HashSet<String>>>,
write_error: Arc<Mutex<Option<String>>>,
) {
let mut pending: HashMap<String, Recording> = HashMap::new();
while let Ok(ev) = rx.recv() {
match ev {
RecordEvent::Input {
call_id,
input,
provider,
} => {
let rec = pending
.entry(call_id.clone())
.or_insert_with(|| Recording::new(&call_id, input.clone(), provider.clone()));
rec.input = input;
if rec.provider.provider.is_empty() {
rec.provider = provider;
}
}
RecordEvent::Session {
call_id,
session_id,
step,
} => {
let rec = entry_or_init(&mut pending, &call_id);
rec.session_id = Some(session_id);
rec.step = Some(step);
try_finalize(&mut w, &mut pending, &call_id, &write_error);
}
RecordEvent::Provider { call_id, provider } => {
entry_or_init(&mut pending, &call_id).provider = provider;
}
RecordEvent::Exchange { call_id, exchange } => {
let rec = entry_or_init(&mut pending, &call_id);
if insert_exchange(rec, exchange) {
mark_inconsistent(&inconsistent, &call_id);
}
try_finalize(&mut w, &mut pending, &call_id, &write_error);
}
RecordEvent::ExchangeUpdate {
call_id,
attempt,
response,
error,
} => {
let rec = entry_or_init(&mut pending, &call_id);
if !matches!(
apply_exchange_update(rec, attempt, response, error),
UpdateMatch::Patched
) {
mark_inconsistent(&inconsistent, &call_id);
}
try_finalize(&mut w, &mut pending, &call_id, &write_error);
}
RecordEvent::TransportClosed { call_id } => {
entry_or_init(&mut pending, &call_id).transport_closed = true;
try_finalize(&mut w, &mut pending, &call_id, &write_error);
}
RecordEvent::Outcome { call_id, outcome } => {
entry_or_init(&mut pending, &call_id).outcome = outcome;
try_finalize(&mut w, &mut pending, &call_id, &write_error);
}
RecordEvent::Flush { ack } => {
write_ready_all(&mut w, &mut pending, &write_error);
if let Err(e) = w.flush() {
note_write_error(&write_error, &e);
}
let status = match write_error.lock().unwrap().as_ref() {
Some(msg) => Err(RecordingError::Write(msg.clone())),
None => Ok(()),
};
let _ = ack.send(status);
}
}
}
for (_, rec) in pending.drain() {
let mut rec = rec;
rec.complete = false;
if rec.outcome.status == OutcomeStatus::Pending {
rec.outcome.status = OutcomeStatus::Incomplete;
}
write_line(&mut w, rec, &write_error);
}
if let Err(e) = w.flush() {
note_write_error(&write_error, &e);
}
}
fn try_finalize(
w: &mut impl Write,
pending: &mut HashMap<String, Recording>,
call_id: &str,
write_error: &Mutex<Option<String>>,
) {
if pending.get(call_id).map(Recording::ready).unwrap_or(false)
&& let Some(mut rec) = pending.remove(call_id)
{
if !rec.exchanges.is_empty() || rec.transport_closed {
rec.complete = true;
}
write_line(w, rec, write_error);
}
}
fn write_ready_all(
w: &mut impl Write,
pending: &mut HashMap<String, Recording>,
write_error: &Mutex<Option<String>>,
) {
let ids: Vec<String> = pending
.iter()
.filter(|(_, r)| r.ready())
.map(|(id, _)| id.clone())
.collect();
for id in ids {
if let Some(mut rec) = pending.remove(&id) {
rec.complete = true;
write_line(w, rec, write_error);
}
}
}
fn write_line(w: &mut impl Write, rec: Recording, write_error: &Mutex<Option<String>>) {
match serde_json::to_string(&rec) {
Ok(line) => {
if let Err(e) = writeln!(w, "{line}").and_then(|()| w.flush()) {
note_write_error(write_error, &e);
}
}
Err(e) => note_write_error(write_error, &e),
}
}
fn note_write_error(slot: &Mutex<Option<String>>, e: &impl std::fmt::Display) {
let mut guard = slot.lock().unwrap();
if guard.is_none() {
*guard = Some(e.to_string());
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UpdateMatch {
Patched,
NotFound,
Ambiguous,
}
fn insert_exchange(rec: &mut Recording, exchange: HttpExchange) -> bool {
if let Some(existing) = rec
.exchanges
.iter_mut()
.find(|e| e.attempt == exchange.attempt)
{
merge_exchange(existing, &exchange);
true
} else {
rec.exchanges.push(exchange);
false
}
}
fn merge_exchange(existing: &mut HttpExchange, new: &HttpExchange) {
existing.request = new.request.clone();
existing.timing = new.timing.clone();
if new.response.is_some() {
existing.response = new.response.clone();
}
if new.error.is_some() {
existing.error = new.error.clone();
}
if new.finalized {
existing.finalized = true;
}
}
fn apply_exchange_update(
rec: &mut Recording,
attempt: u32,
response: ResponseRecord,
error: Option<String>,
) -> UpdateMatch {
let matches: Vec<usize> = rec
.exchanges
.iter()
.enumerate()
.filter(|(_, e)| e.attempt == attempt)
.map(|(i, _)| i)
.collect();
match matches.len() {
0 => UpdateMatch::NotFound,
1 => {
let ex = &mut rec.exchanges[matches[0]];
ex.response = Some(response);
if let Some(err) = error {
ex.error = Some(err);
}
ex.finalized = true;
UpdateMatch::Patched
}
_ => UpdateMatch::Ambiguous,
}
}
fn mark_inconsistent(set: &Mutex<HashSet<String>>, call_id: &str) {
if let Ok(mut g) = set.lock() {
g.insert(call_id.to_string());
}
}
fn send_or_drop(tx: &SyncSender<RecordEvent>, ev: RecordEvent, dropped: &AtomicU64) {
if let Err(err) = tx.try_send(ev) {
drop(err);
dropped.fetch_add(1, Ordering::Relaxed);
}
}
fn iso8601_now() -> String {
let d = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default();
format_rfc3339_utc(d)
}
fn format_rfc3339_utc(d: Duration) -> String {
let secs = d.as_secs() as i64;
let millis = d.subsec_millis();
let (year, month, day) = civil_from_days(secs.div_euclid(86_400));
let hms = secs.rem_euclid(86_400);
format!(
"{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{millis:03}Z",
hms / 3600,
(hms % 3600) / 60,
hms % 60,
)
}
fn civil_from_days(z: i64) -> (i64, u32, u32) {
let z = z + 719_468;
let era = (if z >= 0 { z } else { z - 146_096 }) / 146_097;
let doe = (z - era * 146_097) as u64;
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
let y = yoe as i64 + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
(if m <= 2 { y + 1 } else { y }, m, d)
}
pub const DEFAULT_RING_CAPACITY: usize = 2048;
pub struct RingRecorder {
inner: Mutex<RingInner>,
}
fn ring_pending_max(cap: usize) -> usize {
cap.saturating_mul(2).max(2)
}
struct RingInner {
cap: usize,
pending_max: usize,
pending_order: VecDeque<String>,
pending: HashMap<String, Recording>,
completed: VecDeque<Recording>,
dropped: u64,
pending_evicted: u64,
inconsistent: HashSet<String>,
}
impl RingRecorder {
#[must_use]
pub fn new() -> Self {
Self::with_capacity(DEFAULT_RING_CAPACITY)
}
#[must_use]
pub fn with_capacity(cap: usize) -> Self {
assert!(cap > 0, "RingRecorder capacity must be > 0");
Self {
inner: Mutex::new(RingInner {
cap,
pending_max: ring_pending_max(cap),
pending_order: VecDeque::new(),
pending: HashMap::new(),
completed: VecDeque::with_capacity(cap),
dropped: 0,
pending_evicted: 0,
inconsistent: HashSet::new(),
}),
}
}
pub fn len(&self) -> usize {
self.inner.lock().unwrap().completed.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn pending_count(&self) -> usize {
self.inner.lock().unwrap().pending.len()
}
pub fn dropped_count(&self) -> u64 {
self.inner.lock().unwrap().dropped
}
pub fn pending_evicted_count(&self) -> u64 {
self.inner.lock().unwrap().pending_evicted
}
pub fn inconsistent_call_ids(&self) -> Vec<String> {
let mut v: Vec<String> = self
.inner
.lock()
.unwrap()
.inconsistent
.iter()
.cloned()
.collect();
v.sort();
v
}
pub fn clear(&self) {
let mut inner = self.inner.lock().unwrap();
inner.pending.clear();
inner.pending_order.clear();
inner.completed.clear();
inner.inconsistent.clear();
inner.dropped = 0;
inner.pending_evicted = 0;
}
pub fn export_jsonl(&self, w: &mut impl Write) -> std::io::Result<()> {
let inner = self.inner.lock().unwrap();
for rec in &inner.completed {
let line = serde_json::to_string(rec).map_err(std::io::Error::other)?;
writeln!(w, "{line}")?;
}
Ok(())
}
pub fn completed(&self) -> Vec<Recording> {
self.inner
.lock()
.unwrap()
.completed
.iter()
.cloned()
.collect()
}
pub fn get(&self, call_id: &str) -> Option<Recording> {
self.inner
.lock()
.unwrap()
.completed
.iter()
.find(|r| r.call_id == call_id)
.cloned()
}
}
impl RingInner {
fn entry_or_init_bounded(&mut self, call_id: &str) -> &mut Recording {
if !self.pending.contains_key(call_id) {
self.evict_pending_if_needed();
self.pending_order.push_back(call_id.to_string());
self.pending.insert(
call_id.to_string(),
Recording::new(
call_id,
InputRecord {
prompt: Vec::new(),
options: serde_json::Value::Null,
},
ProviderRecord::minimal("", ""),
),
);
}
self.pending
.get_mut(call_id)
.expect("just inserted or present")
}
fn evict_pending_if_needed(&mut self) {
while self
.pending_order
.front()
.is_some_and(|id| !self.pending.contains_key(id))
{
self.pending_order.pop_front();
}
while self.pending.len() >= self.pending_max {
let Some(old) = self.pending_order.pop_front() else {
break;
};
if let Some(mut rec) = self.pending.remove(&old) {
rec.complete = false;
if rec.outcome.status == OutcomeStatus::Pending {
rec.outcome.status = OutcomeStatus::Incomplete;
}
self.pending_evicted += 1;
self.push(rec);
}
}
}
fn finalize(&mut self, call_id: &str) {
if self.pending.get(call_id).is_some_and(Recording::ready)
&& let Some(mut rec) = self.pending.remove(call_id)
{
if !rec.exchanges.is_empty() || rec.transport_closed {
rec.complete = true;
}
self.push(rec);
}
}
fn push(&mut self, rec: Recording) {
if self.completed.len() >= self.cap {
self.completed.pop_front();
self.dropped += 1;
}
self.completed.push_back(rec);
}
fn mark_inconsistent(&mut self, call_id: &str) {
self.inconsistent.insert(call_id.to_string());
}
}
impl Default for RingRecorder {
fn default() -> Self {
Self::new()
}
}
impl Recorder for RingRecorder {
fn record_input(&self, call_id: &str, options: &CallOptions, provider: &str, model_id: &str) {
let mut inner = self.inner.lock().unwrap();
let rec = inner.entry_or_init_bounded(call_id);
rec.input = InputRecord::from_call_options(options);
if rec.provider.provider.is_empty() {
rec.provider = ProviderRecord::minimal(provider, model_id);
}
}
fn record_session(&self, call_id: &str, session_id: &str, step: u32) {
let mut inner = self.inner.lock().unwrap();
let rec = inner.entry_or_init_bounded(call_id);
rec.session_id = Some(session_id.to_string());
rec.step = Some(step);
}
fn record_provider(&self, call_id: &str, snapshot: &ProviderRecord) {
let mut snap = snapshot.clone();
snap.provider_options = snap.provider_options.take().map(redact_json);
snap.profile = snap.profile.take().map(redact_json);
let mut inner = self.inner.lock().unwrap();
inner.entry_or_init_bounded(call_id).provider = snap;
}
fn record_exchange(&self, call_id: &str, exchange: &HttpExchange) {
let mut inner = self.inner.lock().unwrap();
let rec = inner.entry_or_init_bounded(call_id);
if insert_exchange(rec, exchange.clone()) {
inner.mark_inconsistent(call_id);
}
inner.finalize(call_id);
}
fn record_exchange_update(
&self,
call_id: &str,
attempt: u32,
response: &ResponseRecord,
error: Option<String>,
) {
let mut inner = self.inner.lock().unwrap();
let rec = inner.entry_or_init_bounded(call_id);
if !matches!(
apply_exchange_update(rec, attempt, response.clone(), error),
UpdateMatch::Patched
) {
inner.mark_inconsistent(call_id);
}
inner.finalize(call_id);
}
fn record_outcome(&self, call_id: &str, outcome: &OutcomeRecord) {
let mut inner = self.inner.lock().unwrap();
inner.entry_or_init_bounded(call_id).outcome = outcome.clone();
inner.finalize(call_id);
}
fn record_transport_closed(&self, call_id: &str) {
let mut inner = self.inner.lock().unwrap();
inner.entry_or_init_bounded(call_id).transport_closed = true;
inner.finalize(call_id);
}
fn flush(&self) {
let mut inner = self.inner.lock().unwrap();
let ids: Vec<String> = inner
.pending
.iter()
.filter(|(_, r)| r.ready())
.map(|(id, _)| id.clone())
.collect();
for id in ids {
if let Some(mut rec) = inner.pending.remove(&id) {
rec.complete = true;
inner.push(rec);
}
}
}
fn try_flush(&self) -> Result<(), RecordingError> {
RingRecorder::flush(self);
Ok(())
}
}
pub struct RecordingOutcomeStream<S> {
inner: S,
recorder: Option<Arc<dyn Recorder>>,
call_id: String,
recorded: bool,
}
impl<S> RecordingOutcomeStream<S> {
pub fn new(inner: S, recorder: Option<Arc<dyn Recorder>>, call_id: impl Into<String>) -> Self {
Self {
inner,
recorder,
call_id: call_id.into(),
recorded: false,
}
}
fn record(&mut self, outcome: &OutcomeRecord) {
if !self.recorded {
self.recorded = true;
if let Some(rec) = &self.recorder {
rec.record_outcome(&self.call_id, outcome);
rec.record_transport_closed(&self.call_id);
}
}
}
}
impl<S, E> futures::Stream for RecordingOutcomeStream<S>
where
S: futures::Stream<Item = Result<crate::stream_part::StreamPart, E>> + Unpin,
E: std::fmt::Display,
{
type Item = Result<crate::stream_part::StreamPart, E>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
let this = self.as_mut().get_mut();
let inner = std::pin::Pin::new(&mut this.inner);
match inner.poll_next(cx) {
std::task::Poll::Ready(Some(Ok(part))) => {
match &part {
crate::stream_part::StreamPart::Finish {
finish_reason,
usage,
..
} => {
let outcome = OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: serde_json::to_value(finish_reason.unified)
.ok()
.and_then(|v| v.as_str().map(std::string::ToString::to_string)),
error: None,
usage: serde_json::to_value(usage).ok(),
};
self.as_mut().get_mut().record(&outcome);
}
crate::stream_part::StreamPart::Error { error } => {
let outcome = OutcomeRecord {
status: OutcomeStatus::Error,
finish_reason: None,
error: Some(error.to_string()),
usage: None,
};
self.as_mut().get_mut().record(&outcome);
}
_ => {}
}
std::task::Poll::Ready(Some(Ok(part)))
}
std::task::Poll::Ready(Some(Err(e))) => {
let outcome = OutcomeRecord {
status: OutcomeStatus::Error,
finish_reason: None,
error: Some(e.to_string()),
usage: None,
};
this.record(&outcome);
std::task::Poll::Ready(Some(Err(e)))
}
std::task::Poll::Ready(None) => {
let outcome = OutcomeRecord {
status: OutcomeStatus::Incomplete,
finish_reason: None,
error: None,
usage: None,
};
this.record(&outcome);
std::task::Poll::Ready(None)
}
std::task::Poll::Pending => std::task::Poll::Pending,
}
}
}
impl<S> Drop for RecordingOutcomeStream<S> {
fn drop(&mut self) {
if !self.recorded {
let outcome = OutcomeRecord {
status: OutcomeStatus::Cancelled,
finish_reason: None,
error: None,
usage: None,
};
self.record(&outcome);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::generate::GenerateTextOptions;
fn sample_options() -> CallOptions {
GenerateTextOptions {
temperature: Some(0.7),
..Default::default()
}
.into_call_options(
crate::language_model_message::convert_to_language_model_prompt(
&[crate::message::ModelMessage::user("ping")],
None,
),
)
}
#[test]
fn gate_is_off_by_default() {
assert!(recorder().is_none());
}
#[test]
fn call_id_is_unique() {
let a = new_call_id();
let b = new_call_id();
assert_ne!(a, b);
assert!(a.starts_with("call-"));
}
#[test]
fn input_record_serializes_options_without_internal_fields() {
let mut opts = sample_options();
opts.call_id = Some("call-test".into());
opts.abort_signal = None;
let rec = InputRecord::from_call_options(&opts);
let v = &rec.options;
assert!(v.get("call_id").is_none());
assert!(v.get("abort_signal").is_none());
}
#[test]
fn redact_json_hides_sensitive_values_everywhere() {
let v = serde_json::json!({
"headers": {
"Authorization": "Bearer sk-abc",
"X-Key": "ok",
"x-goog-api-key": "gkey",
"x-amz-security-token": "sts-tok"
},
"provider_options": { "headers": { "Cookie": "s=1" } },
"body_overrides": { "api_key": "sk-secret" },
"usage": {
"max_output_tokens": 4096,
"prompt_tokens": 10,
"completion_tokens": 20
},
"temperature": 0.5,
});
let r = redact_json(v);
assert_eq!(r["headers"]["Authorization"], "[REDACTED]");
assert_eq!(r["headers"]["x-goog-api-key"], "[REDACTED]");
assert_eq!(r["headers"]["x-amz-security-token"], "[REDACTED]");
assert_eq!(r["headers"]["X-Key"], "ok");
assert_eq!(r["provider_options"]["headers"]["Cookie"], "[REDACTED]");
assert_eq!(r["body_overrides"]["api_key"], "[REDACTED]");
assert_eq!(r["usage"]["max_output_tokens"], 4096);
assert_eq!(r["usage"]["prompt_tokens"], 10);
assert_eq!(r["usage"]["completion_tokens"], 20);
assert_eq!(r["temperature"], 0.5);
}
#[test]
fn recording_round_trips() {
let rec = Recording::new(
"call-1",
InputRecord::from_call_options(&sample_options()),
ProviderRecord::minimal("openai", "gpt-4o"),
);
let json = serde_json::to_string(&rec).unwrap();
let back: Recording = serde_json::from_str(&json).unwrap();
assert_eq!(back.call_id, "call-1");
assert_eq!(back.schema, RECORDING_SCHEMA);
assert_eq!(back.provider.provider, "openai");
}
#[test]
fn jsonl_flush_writes_complete_line_blocking() {
let dir = std::env::temp_dir().join(format!("aimux-rec-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-flush-1", &sample_options(), "openai", "gpt-4o");
rec.record_outcome(
"call-flush-1",
&OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
},
);
rec.record_transport_closed("call-flush-1");
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(parsed.call_id, "call-flush-1");
assert_eq!(parsed.outcome.status, OutcomeStatus::Success);
assert!(parsed.exchanges.is_empty());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn barrier_requires_input_record() {
let mut rec = Recording::new(
"call-x",
InputRecord {
prompt: Vec::new(),
options: serde_json::Value::Null,
},
ProviderRecord::minimal("openai", "gpt-4o"),
);
rec.transport_closed = true;
rec.outcome = OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
};
assert!(!rec.ready(), "empty-prompt placeholder must not finalize");
rec.input = InputRecord::from_call_options(&sample_options());
assert!(rec.ready(), "real input satisfies the barrier");
}
#[test]
fn missing_input_writes_incomplete_on_shutdown() {
let dir = std::env::temp_dir().join(format!("aimux-rec-noinput-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_outcome(
"call-noinput",
&OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
},
);
rec.record_transport_closed("call-noinput");
rec.flush();
drop(rec);
let content = std::fs::read_to_string(dir.join("recordings.jsonl")).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert!(
!parsed.complete,
"must not be marked complete without input"
);
assert!(parsed.input.prompt.is_empty());
assert_eq!(parsed.outcome.status, OutcomeStatus::Success);
let _ = std::fs::remove_dir_all(&dir);
}
struct FailingWriter;
impl std::io::Write for FailingWriter {
fn write(&mut self, _buf: &[u8]) -> std::io::Result<usize> {
Err(std::io::Error::other("simulated disk full"))
}
fn flush(&mut self) -> std::io::Result<()> {
Err(std::io::Error::other("simulated disk full"))
}
}
#[test]
fn try_flush_reports_writer_io_error() {
let (tx, rx) = sync_channel(16);
let write_error = Arc::new(Mutex::new(None));
let handle = {
let inconsistent = Arc::new(Mutex::new(HashSet::new()));
let w_err = write_error.clone();
std::thread::spawn(move || writer_loop(rx, FailingWriter, inconsistent, w_err))
};
tx.send(RecordEvent::Input {
call_id: "c1".into(),
input: InputRecord::from_call_options(&sample_options()),
provider: ProviderRecord::minimal("openai", "gpt-4o"),
})
.unwrap();
tx.send(RecordEvent::TransportClosed {
call_id: "c1".into(),
})
.unwrap();
tx.send(RecordEvent::Outcome {
call_id: "c1".into(),
outcome: OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
},
})
.unwrap();
let (ack_tx, ack_rx) = sync_channel(0);
tx.send(RecordEvent::Flush { ack: ack_tx }).unwrap();
let status = ack_rx
.recv_timeout(Duration::from_secs(5))
.expect("flush ack within timeout");
assert!(
matches!(&status, Err(RecordingError::Write(m)) if m.contains("disk full")),
"expected Err(Write) carrying the io error, got {status:?}"
);
drop(tx);
handle.join().unwrap();
assert!(write_error.lock().unwrap().is_some());
}
#[test]
fn record_session_writes_session_fields_to_jsonl() {
let dir = std::env::temp_dir().join(format!("aimux-rec-sess-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-sess-1", &sample_options(), "openai", "gpt-4o");
rec.record_session("call-sess-1", "sess-abc", 3);
rec.record_outcome(
"call-sess-1",
&OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
},
);
rec.record_transport_closed("call-sess-1");
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(parsed.session_id.as_deref(), Some("sess-abc"));
assert_eq!(parsed.step, Some(3));
assert!(content.contains("\"session_id\":\"sess-abc\""));
assert!(content.contains("\"step\":3"));
}
#[test]
fn ungrouped_recording_omits_session_fields_and_old_jsonl_parses() {
let rec = Recording::new(
"call-1",
InputRecord::from_call_options(&sample_options()),
ProviderRecord::minimal("openai", "gpt-4o"),
);
let json = serde_json::to_string(&rec).unwrap();
assert!(!json.contains("session_id"), "ungrouped must omit: {json}");
assert!(!json.contains("\"step\""), "ungrouped must omit: {json}");
let old_json = serde_json::to_string(&{
let mut r = serde_json::to_value(&rec).unwrap();
r.as_object_mut()
.unwrap()
.insert("transport_closed".into(), serde_json::json!(true));
r
})
.unwrap();
let back: Recording = serde_json::from_str(&old_json).unwrap();
assert_eq!(back.session_id, None);
assert_eq!(back.step, None);
}
#[test]
fn exchange_update_is_attached_to_correct_attempt() {
let dir = std::env::temp_dir().join(format!("aimux-rec-test2-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-x2", &sample_options(), "openai", "gpt-4o");
rec.record_exchange(
"call-x2",
&HttpExchange {
attempt: 0,
request: HttpRecord {
method: "post".into(),
url: "u".into(),
headers: vec![("authorization".into(), "[REDACTED]".into())],
body: Some("{}".into()),
},
response: None,
timing: TimingRecord {
latency_ms: 1,
ttfb_ms: None,
},
error: None,
finalized: false,
},
);
rec.record_exchange_update(
"call-x2",
0,
&ResponseRecord {
status: 200,
headers: vec![],
body: Some("{\"ok\":true}".into()),
stream_chunks: Some(1),
ttfb_ms: None,
},
None,
);
rec.record_outcome(
"call-x2",
&OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: None,
error: None,
usage: None,
},
);
rec.record_transport_closed("call-x2");
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(parsed.exchanges.len(), 1);
assert_eq!(parsed.exchanges[0].response.as_ref().unwrap().status, 200);
assert!(parsed.exchanges[0].finalized);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn error_patches_existing_exchange_not_duplicate() {
let dir = std::env::temp_dir().join(format!("aimux-rec-errpatch-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-errp", &sample_options(), "openai", "gpt-4o");
rec.record_exchange(
"call-errp",
&HttpExchange {
attempt: 3,
request: HttpRecord {
method: "post".into(),
url: "u".into(),
headers: vec![],
body: Some("{}".into()),
},
response: Some(ResponseRecord {
status: 200,
headers: vec![],
body: None,
stream_chunks: None,
ttfb_ms: None,
}),
timing: TimingRecord {
latency_ms: 1,
ttfb_ms: None,
},
error: None,
finalized: false,
},
);
rec.record_exchange_update(
"call-errp",
3,
&ResponseRecord {
status: 200,
headers: vec![],
body: Some(r#"{"partial":true}"#.into()),
stream_chunks: Some(1),
ttfb_ms: Some(12),
},
Some("mid-stream error".to_string()),
);
rec.record_transport_closed("call-errp");
rec.record_outcome(
"call-errp",
&OutcomeRecord {
status: OutcomeStatus::Error,
finish_reason: None,
error: Some("mid-stream".into()),
usage: None,
},
);
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(
parsed.exchanges.len(),
1,
"error must patch same attempt, not duplicate"
);
assert_eq!(parsed.exchanges[0].attempt, 3);
let ex = &parsed.exchanges[0];
assert!(ex.finalized);
assert_eq!(ex.error.as_deref(), Some("mid-stream error"));
assert!(
ex.response
.as_ref()
.unwrap()
.body
.as_deref()
.unwrap()
.contains("partial")
);
assert_eq!(ex.response.as_ref().unwrap().ttfb_ms, Some(12));
assert!(parsed.complete, "finalized exchange + closed -> complete");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn barrier_waits_for_transport_closed_and_marks_complete() {
let dir = std::env::temp_dir().join(format!("aimux-rec-btc-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-btc-1", &sample_options(), "openai", "gpt-4o");
rec.record_outcome(
"call-btc-1",
&OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: None,
error: None,
usage: None,
},
);
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
assert!(
content.trim().is_empty(),
"must not write before transport closed"
);
rec.record_transport_closed("call-btc-1");
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(parsed.call_id, "call-btc-1");
assert!(
parsed.complete,
"fully-finalized recording must be complete"
);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn incomplete_outcome_flushes_on_shutdown() {
let dir = std::env::temp_dir().join(format!("aimux-rec-test3-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
{
let rec = JsonlRecorder::new(&dir);
rec.record_input("call-orphan", &sample_options(), "openai", "gpt-4o");
}
let content = std::fs::read_to_string(dir.join("recordings.jsonl"))
.expect("disconnected fallback must write the incomplete line");
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(parsed.outcome.status, OutcomeStatus::Incomplete);
assert!(
!parsed.complete,
"fallback line must not be marked complete"
);
let _ = std::fs::remove_dir_all(&dir);
}
fn sample_exchange() -> HttpExchange {
HttpExchange {
attempt: 0,
request: HttpRecord {
method: "post".into(),
url: "https://api.openai.com/v1/chat/completions".into(),
headers: vec![],
body: Some("{}".into()),
},
response: Some(ResponseRecord {
status: 200,
headers: vec![],
body: Some("{\"ok\":true}".into()),
stream_chunks: None,
ttfb_ms: None,
}),
timing: TimingRecord {
latency_ms: 1,
ttfb_ms: None,
},
error: None,
finalized: true,
}
}
fn success_outcome() -> OutcomeRecord {
OutcomeRecord {
status: OutcomeStatus::Success,
finish_reason: Some("stop".into()),
error: None,
usage: None,
}
}
fn drive_full_call(rec: &RingRecorder, call_id: &str) {
rec.record_input(call_id, &sample_options(), "openai", "gpt-4o");
rec.record_exchange(call_id, &sample_exchange());
rec.record_outcome(call_id, &success_outcome());
rec.record_transport_closed(call_id);
}
#[test]
fn ring_merges_shards_and_finalizes_on_barrier() {
let ring = RingRecorder::new();
drive_full_call(&ring, "call-a");
assert_eq!(ring.len(), 1);
assert_eq!(ring.pending_count(), 0);
let mut buf = Vec::new();
ring.export_jsonl(&mut buf).unwrap();
let parsed: Recording =
serde_json::from_str(std::str::from_utf8(&buf).unwrap().trim()).unwrap();
assert_eq!(parsed.call_id, "call-a");
assert_eq!(parsed.exchanges.len(), 1);
assert_eq!(parsed.outcome.status, OutcomeStatus::Success);
assert!(parsed.complete);
}
#[test]
fn ring_records_session_fields_and_get_by_call_id() {
let ring = RingRecorder::new();
ring.record_input("call-s", &sample_options(), "openai", "gpt-4o");
ring.record_session("call-s", "sess-1", 3);
ring.record_outcome("call-s", &success_outcome());
ring.record_transport_closed("call-s");
ring.flush();
let rec = ring
.get("call-s")
.expect("recording retrievable by call_id");
assert_eq!(rec.session_id.as_deref(), Some("sess-1"));
assert_eq!(rec.step, Some(3));
let all = ring.completed();
assert_eq!(all.len(), 1);
assert_eq!(all[0].call_id, "call-s");
assert!(ring.get("missing").is_none());
}
#[test]
fn ring_keeps_incomplete_in_pending_until_closed() {
let ring = RingRecorder::new();
ring.record_input("call-x", &sample_options(), "openai", "gpt-4o");
ring.record_outcome("call-x", &success_outcome());
ring.flush();
assert_eq!(ring.len(), 0);
assert_eq!(ring.pending_count(), 1);
ring.record_transport_closed("call-x");
ring.flush();
assert_eq!(ring.len(), 1);
assert_eq!(ring.pending_count(), 0);
}
#[test]
fn ring_evicts_oldest_and_counts_drops() {
let ring = RingRecorder::with_capacity(2);
drive_full_call(&ring, "call-1");
drive_full_call(&ring, "call-2");
drive_full_call(&ring, "call-3");
assert_eq!(ring.len(), 2, "ring is bounded");
assert_eq!(ring.dropped_count(), 1, "oldest evicted and counted");
let mut buf = Vec::new();
ring.export_jsonl(&mut buf).unwrap();
let lines: Vec<&str> = std::str::from_utf8(&buf)
.unwrap()
.lines()
.filter(|l| !l.trim().is_empty())
.collect();
assert_eq!(lines.len(), 2);
assert!(lines[0].contains("call-2"));
assert!(lines[1].contains("call-3"));
}
#[test]
fn ring_flush_finalizes_ready_pending() {
let ring = RingRecorder::new();
ring.record_input("call-f", &sample_options(), "openai", "gpt-4o");
ring.record_outcome("call-f", &success_outcome());
ring.record_transport_closed("call-f");
ring.flush();
ring.flush();
assert_eq!(ring.len(), 1);
}
#[test]
fn ring_clear_resets_everything() {
let ring = RingRecorder::with_capacity(2);
drive_full_call(&ring, "call-1");
ring.record_input("call-p", &sample_options(), "openai", "gpt-4o"); assert_eq!(ring.len(), 1);
assert_eq!(ring.pending_count(), 1);
ring.clear();
assert_eq!(ring.len(), 0);
assert_eq!(ring.pending_count(), 0);
assert_eq!(ring.dropped_count(), 0);
}
fn days_from_civil(y: i64, m: u32, d: u32) -> i64 {
let y = if m <= 2 { y - 1 } else { y };
let era = (if y >= 0 { y } else { y - 399 }) / 400;
let yoe = (y - era * 400) as u64;
let mp = if m > 2 { m - 3 } else { m + 9 } as u64;
let doy = (153 * mp + 2) / 5 + d as u64 - 1;
let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
era * 146_097 + doe as i64 - 719_468
}
fn parse_rfc3339_secs(s: &str) -> Option<(i64, u32)> {
let b = s.as_bytes();
if s.len() < 24
|| b[4] != b'-'
|| b[7] != b'-'
|| b[10] != b'T'
|| b[13] != b':'
|| b[16] != b':'
|| b[19] != b'.'
|| !s.ends_with('Z')
{
return None;
}
let year: i64 = s[0..4].parse().ok()?;
let month: u32 = s[5..7].parse().ok()?;
let day: u32 = s[8..10].parse().ok()?;
let hour: u64 = s[11..13].parse().ok()?;
let min: u64 = s[14..16].parse().ok()?;
let sec: u64 = s[17..19].parse().ok()?;
let millis: u32 = s[20..s.len() - 1].parse().ok()?;
let total =
days_from_civil(year, month, day) * 86_400 + (hour * 3600 + min * 60 + sec) as i64;
Some((total, millis))
}
#[test]
fn format_rfc3339_utc_is_valid_rfc3339() {
assert_eq!(
format_rfc3339_utc(Duration::ZERO),
"1970-01-01T00:00:00.000Z"
);
let d = Duration::from_secs(1_700_000_000);
let s = format_rfc3339_utc(d);
assert_eq!(s.len(), 24);
assert!(s.ends_with('Z'));
let (secs, millis) = parse_rfc3339_secs(&s).expect("must parse");
assert_eq!(secs, 1_700_000_000);
assert_eq!(millis, 0);
let (s2, m2) =
parse_rfc3339_secs(&format_rfc3339_utc(Duration::from_millis(12_345))).unwrap();
assert_eq!((s2, m2), (12, 345));
let now = iso8601_now();
assert!(parse_rfc3339_secs(&now).is_some(), "now not rfc3339: {now}");
}
#[test]
fn send_or_drop_counts_overflow() {
let (tx, rx) = sync_channel::<RecordEvent>(1);
let dropped = AtomicU64::new(0);
let ev = || RecordEvent::TransportClosed {
call_id: "c".into(),
};
send_or_drop(&tx, ev(), &dropped);
send_or_drop(&tx, ev(), &dropped);
send_or_drop(&tx, ev(), &dropped);
assert_eq!(dropped.load(Ordering::Relaxed), 2);
assert!(rx.try_recv().is_ok()); assert!(rx.try_recv().is_err()); }
#[test]
fn jsonl_dropped_count_starts_zero() {
let dir = std::env::temp_dir().join(format!("aimux-rec-drop-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
assert_eq!(rec.dropped_count(), 0);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn jsonl_try_new_fails_on_invalid_dir() {
let blocker =
std::env::temp_dir().join(format!("aimux-rec-blocker-{}", std::process::id()));
let _ = std::fs::remove_file(&blocker);
std::fs::write(&blocker, b"x").unwrap();
let res = JsonlRecorder::try_new(blocker.join("sub"));
assert!(
matches!(res, Err(RecordingError::Init { .. })),
"expected Init, got {:?}",
res.as_ref().err()
);
let _ = std::fs::remove_file(&blocker);
}
#[test]
fn jsonl_try_flush_ok_and_disabled_err() {
let dir = std::env::temp_dir().join(format!("aimux-rec-tf-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("c", &sample_options(), "openai", "gpt-4o");
rec.record_outcome("c", &success_outcome());
rec.record_transport_closed("c");
assert!(rec.try_flush().is_ok());
let content = std::fs::read_to_string(rec.path()).unwrap();
assert!(content.contains("\"call_id\":\"c\""));
let _ = std::fs::remove_dir_all(&dir);
let blocker =
std::env::temp_dir().join(format!("aimux-rec-blocker2-{}", std::process::id()));
let _ = std::fs::remove_file(&blocker);
std::fs::write(&blocker, b"x").unwrap();
let disabled = JsonlRecorder::new(blocker.join("sub"));
assert!(matches!(
disabled.try_flush(),
Err(RecordingError::WriterGone)
));
let _ = std::fs::remove_file(&blocker);
}
#[test]
fn ring_pending_evicts_oldest_when_bounded() {
let ring = RingRecorder::with_capacity(2); for i in 1..=5u32 {
ring.record_input(&format!("call-{i}"), &sample_options(), "openai", "gpt-4o");
}
assert_eq!(ring.pending_count(), 4, "pending bounded to 4");
assert_eq!(ring.len(), 1, "oldest evicted into ring as incomplete");
assert_eq!(ring.pending_evicted_count(), 1);
let mut buf = Vec::new();
ring.export_jsonl(&mut buf).unwrap();
let parsed: Recording =
serde_json::from_str(std::str::from_utf8(&buf).unwrap().trim()).unwrap();
assert_eq!(parsed.call_id, "call-1", "oldest (call-1) evicted");
assert!(!parsed.complete, "evicted pending → incomplete");
assert_eq!(parsed.outcome.status, OutcomeStatus::Incomplete);
}
fn full_provider() -> ProviderRecord {
ProviderRecord {
provider: "openai".into(),
model_id: "gpt-4o".into(),
base_url: Some("https://api.openai.com".into()),
api_key_source: "env:OPENAI_API_KEY".into(),
profile: None,
provider_options: None,
}
}
#[test]
fn ring_provider_before_input_preserved() {
let ring = RingRecorder::new();
ring.record_provider("c", &full_provider()); ring.record_input("c", &sample_options(), "minimal", "minimal-model");
ring.record_outcome("c", &success_outcome());
ring.record_transport_closed("c");
let mut buf = Vec::new();
ring.export_jsonl(&mut buf).unwrap();
let parsed: Recording =
serde_json::from_str(std::str::from_utf8(&buf).unwrap().trim()).unwrap();
assert_eq!(parsed.provider.provider, "openai");
assert_eq!(
parsed.provider.base_url.as_deref(),
Some("https://api.openai.com"),
"full snapshot preserved, not overwritten by minimal"
);
assert_eq!(parsed.provider.model_id, "gpt-4o");
}
#[test]
fn jsonl_provider_before_input_preserved() {
let dir = std::env::temp_dir().join(format!("aimux-rec-pbi-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_provider("c", &full_provider()); rec.record_input("c", &sample_options(), "minimal", "minimal-model");
rec.record_outcome("c", &success_outcome());
rec.record_transport_closed("c");
rec.flush();
let content = std::fs::read_to_string(rec.path()).unwrap();
let parsed: Recording = serde_json::from_str(content.trim()).unwrap();
assert_eq!(
parsed.provider.base_url.as_deref(),
Some("https://api.openai.com")
);
assert_eq!(parsed.provider.provider, "openai");
let _ = std::fs::remove_dir_all(&dir);
}
fn resp_status(status: u16) -> ResponseRecord {
ResponseRecord {
status,
headers: vec![],
body: None,
stream_chunks: None,
ttfb_ms: None,
}
}
fn skeleton_at(attempt: u32) -> HttpExchange {
HttpExchange {
attempt,
request: HttpRecord {
method: "post".into(),
url: "u".into(),
headers: vec![],
body: None,
},
response: None,
timing: TimingRecord {
latency_ms: 1,
ttfb_ms: None,
},
error: None,
finalized: false,
}
}
fn empty_recording(call_id: &str) -> Recording {
Recording::new(
call_id,
InputRecord {
prompt: Vec::new(),
options: serde_json::Value::Null,
},
ProviderRecord::minimal("", ""),
)
}
#[test]
fn insert_exchange_and_apply_update_match_semantics() {
let mut rec = empty_recording("c");
assert!(!insert_exchange(&mut rec, skeleton_at(0)));
assert!(!insert_exchange(&mut rec, skeleton_at(1)));
assert_eq!(rec.exchanges.len(), 2);
let mut dup = skeleton_at(0);
dup.response = Some(resp_status(500));
dup.error = Some("dup".into());
dup.finalized = true;
assert!(insert_exchange(&mut rec, dup));
assert_eq!(rec.exchanges.len(), 2);
let e0 = rec.exchanges.iter().find(|e| e.attempt == 0).unwrap();
assert_eq!(e0.response.as_ref().unwrap().status, 500);
assert!(e0.finalized);
assert_eq!(e0.error.as_deref(), Some("dup"));
assert_eq!(
apply_exchange_update(&mut rec, 1, resp_status(200), None),
UpdateMatch::Patched
);
assert_eq!(
rec.exchanges
.iter()
.find(|e| e.attempt == 1)
.unwrap()
.response
.as_ref()
.unwrap()
.status,
200
);
assert_eq!(
apply_exchange_update(&mut rec, 9, resp_status(200), None),
UpdateMatch::NotFound
);
rec.exchanges.push(skeleton_at(1));
assert_eq!(
apply_exchange_update(&mut rec, 1, resp_status(503), Some("amb".into())),
UpdateMatch::Ambiguous
);
assert!(
rec.exchanges
.iter()
.filter(|e| e.attempt == 1)
.all(|e| e.response.as_ref().is_none_or(|r| r.status != 503))
);
}
#[test]
fn ring_duplicate_attempt_marks_inconsistent_and_merges() {
let ring = RingRecorder::new();
ring.record_input("c", &sample_options(), "openai", "gpt-4o");
ring.record_exchange("c", &sample_exchange()); ring.record_exchange("c", &sample_exchange()); assert!(ring.inconsistent_call_ids().contains(&"c".to_string()));
ring.record_outcome("c", &success_outcome());
ring.record_transport_closed("c");
let mut buf = Vec::new();
ring.export_jsonl(&mut buf).unwrap();
let parsed: Recording =
serde_json::from_str(std::str::from_utf8(&buf).unwrap().trim()).unwrap();
assert_eq!(parsed.exchanges.len(), 1, "duplicate merged, not appended");
assert_eq!(parsed.exchanges[0].attempt, 0);
}
#[test]
fn ring_update_not_found_marks_inconsistent() {
let ring = RingRecorder::new();
ring.record_input("c", &sample_options(), "openai", "gpt-4o");
ring.record_exchange_update("c", 0, &resp_status(200), None);
assert!(ring.inconsistent_call_ids().contains(&"c".to_string()));
assert_eq!(
ring.pending_count(),
1,
"call still pending, no exchange created"
);
}
#[test]
fn jsonl_duplicate_attempt_marks_inconsistent() {
let dir = std::env::temp_dir().join(format!("aimux-rec-dup-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
let rec = JsonlRecorder::new(&dir);
rec.record_input("c", &sample_options(), "openai", "gpt-4o");
rec.record_exchange("c", &sample_exchange());
rec.record_exchange("c", &sample_exchange()); rec.flush();
assert!(rec.inconsistent_call_ids().contains(&"c".to_string()));
let _ = std::fs::remove_dir_all(&dir);
}
}