use std::collections::{BTreeMap, BTreeSet};
use serde::{Deserialize, Serialize};
use crate::Config;
use crate::measured::{Measured, Source};
use crate::normalize::{Api, Counts, Exchange, Msg};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Row {
pub id: u64,
pub api: Api,
pub model: String,
pub client: String,
pub t_start_ms: i64,
pub kind: RowKind,
pub side_call_shape: Option<Shape>,
pub session: u32,
pub conversation: Option<u32>,
pub turn: Option<u32>,
pub counts: Counts,
pub composition: Composition,
pub linked_to: Option<u64>,
pub new_conversation: bool,
pub prefix_tokens: Option<u32>,
pub break_cause: Option<BreakCause>,
pub compaction: bool,
pub predicted_reuse: Option<f64>,
pub observed_reuse: Option<f64>,
pub cause: Option<Cause>,
pub lost_client: Option<Loss>,
pub lost_external: Option<Loss>,
pub side_calls_between: Vec<u64>,
pub other_conversations_between: Vec<u64>,
pub other_domains_between: Vec<u64>,
pub loss_measured: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RowKind {
Turn,
SideCall,
Stateless,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Shape {
Summarise,
Structured,
Oneshot,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BreakCause {
ToolsChanged,
SystemChanged,
HistoryRewritten,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Cause {
Structural,
SideCall,
OtherConversation,
OtherModel,
ServerReset,
Undetermined,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Loss {
pub cause: String,
pub tokens: u32,
pub ms: Option<Measured<f64>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Composition {
pub system: Measured<u32>,
pub tools: Measured<u32>,
pub history: Measured<u32>,
pub new_content: Measured<u32>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolUsage {
pub client: String,
pub name: String,
pub schema_tokens: u32,
pub offered: u32,
pub called: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Client {
pub client: String,
pub user_agent: Option<String>,
pub requests: u32,
pub conversations: u32,
pub tools: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct Table {
pub rows: Vec<Row>,
pub tools: Vec<ToolUsage>,
pub clients: Vec<Client>,
}
const SIDE_CALL_MAX_MESSAGES: usize = 3;
const SUMMARISE_USER_SHARE: f64 = 0.9;
const SUMMARISE_MIN_CHARS: u64 = 20_000;
const STRUCTURED_MAX_TOKENS: u32 = 32;
const STATELESS_MIN_REQUESTS: usize = 3;
const COMPACTION_RATIO: f64 = 0.7;
const CAUSE_GAP_POINTS: f64 = 20.0;
const COLD_SHARE: f64 = 0.05;
const DIGEST_EDGE_CHARS: u32 = 256;
const SERVER_EMPTY_BELOW: f64 = 128.0;
const HEAD_MARGIN: f64 = 1.15;
const BULK_MIN_EVALUATED: u32 = 50;
const UNKNOWN_CLIENT: &str = "unknown";
impl BreakCause {
fn name(self) -> &'static str {
match self {
BreakCause::ToolsChanged => "tools_changed",
BreakCause::SystemChanged => "system_changed",
BreakCause::HistoryRewritten => "history_rewritten",
}
}
}
struct Weights {
messages: Vec<f64>,
tools: Vec<f64>,
tool_block: f64,
total: f64,
source: Source,
}
struct Tip {
row: usize,
conversation: u32,
turns: u32,
}
enum Choice {
Continue(usize),
Fresh(Option<usize>),
}
struct Link {
prefix_len: usize,
prefix_tokens: f64,
break_cause: Option<BreakCause>,
compaction: bool,
already_sent: f64,
}
pub(crate) fn build(exchanges: &[Exchange], config: &Config) -> Table {
let mut order: Vec<&Exchange> = exchanges.iter().collect();
order.sort_by_key(|x| (x.t_start_ms, x.id));
let weights: Vec<Weights> = order.iter().map(|x| weights(x)).collect();
let sessions = sessions(&order, config.session_gap_ms);
let kinds = classify(&order);
let bulk_rates = bulk_rates(&order);
let mut rows: Vec<Row> = order
.iter()
.enumerate()
.map(|(i, x)| Row {
id: x.id,
api: x.api,
model: x.model.clone(),
client: String::new(),
t_start_ms: x.t_start_ms,
kind: kinds[i].0,
side_call_shape: kinds[i].1,
session: sessions[i],
conversation: None,
turn: None,
counts: x.counts,
composition: composition(x, &weights[i], 0),
linked_to: None,
new_conversation: false,
prefix_tokens: None,
break_cause: None,
compaction: false,
predicted_reuse: None,
observed_reuse: None,
cause: None,
lost_client: None,
lost_external: None,
side_calls_between: Vec::new(),
other_conversations_between: Vec::new(),
other_domains_between: Vec::new(),
loss_measured: false,
})
.collect();
let mut tips: Vec<Tip> = Vec::new();
for cur in 0..order.len() {
if rows[cur].kind != RowKind::Turn {
continue;
}
let candidates: Vec<usize> = (0..tips.len())
.filter(|&t| {
let prev = order[tips[t].row];
prev.model == order[cur].model
&& prev.user_agent == order[cur].user_agent
&& rows[tips[t].row].session == rows[cur].session
})
.collect();
let (prev, fresh) = match choose(&order, &tips, &candidates, cur) {
Choice::Continue(t) => {
let tip = &mut tips[t];
let prev = tip.row;
tip.row = cur;
tip.turns += 1;
rows[cur].conversation = Some(tip.conversation);
rows[cur].turn = Some(tip.turns);
(Some(prev), false)
}
Choice::Fresh(from) => {
let conversation = tips.len() as u32 + 1;
tips.push(Tip {
row: cur,
conversation,
turns: 1,
});
rows[cur].conversation = Some(conversation);
rows[cur].turn = Some(1);
(from.map(|t| tips[t].row), true)
}
};
let Some(prev) = prev else {
continue;
};
let mut latest = None;
for between in prev + 1..cur {
let id = order[between].id;
if order[between].model != order[cur].model {
rows[cur].other_domains_between.push(id);
} else if rows[between].kind == RowKind::SideCall {
rows[cur].side_calls_between.push(id);
latest = Some(Cause::SideCall);
} else {
rows[cur].other_conversations_between.push(id);
latest = Some(Cause::OtherConversation);
}
}
rows[cur].linked_to = Some(order[prev].id);
rows[cur].new_conversation = fresh;
let link = link(
order[prev],
&weights[prev],
order[cur],
&weights[cur],
fresh,
);
let bulk_rate = bulk_rates.get(order[cur].model.as_str()).copied();
apply_link(
&mut rows[cur],
order[cur],
&weights[cur],
&link,
latest,
bulk_rate,
);
}
assign_clients(&order, &mut rows);
let tools = tool_usage(&order, &weights, &rows);
let clients = clients(&order, &rows);
Table {
rows,
tools,
clients,
}
}
fn choose(order: &[&Exchange], tips: &[Tip], candidates: &[usize], cur: usize) -> Choice {
let x = order[cur];
let lead = x.messages.iter().take_while(|m| m.role == "system").count();
let matched = |t: usize| {
let prev = order[tips[t].row];
prev.messages
.iter()
.zip(&x.messages)
.take_while(|(a, b)| a.hash == b.hash)
.count()
};
let recent = |t: &usize| tips[*t].row;
if let Some(&t) = candidates
.iter()
.filter(|&&t| continues(order[tips[t].row], x))
.max_by_key(|&&t| (order[tips[t].row].messages.len(), recent(&t)))
{
return Choice::Continue(t);
}
if let Some(&t) = candidates
.iter()
.filter(|&&t| matched(t) > lead)
.max_by_key(|&&t| (matched(t), recent(&t)))
{
return Choice::Continue(t);
}
Choice::Fresh(
candidates
.iter()
.filter(|&&t| matched(t) > 0)
.max_by_key(|&&t| (matched(t), recent(&t)))
.copied(),
)
}
fn weights(x: &Exchange) -> Weights {
let mut exact = true;
let mut weight = |tokens: Option<Measured<u32>>, chars: u32| match tokens {
Some(tokens) => f64::from(tokens.value),
None => {
exact = false;
f64::from(chars.div_ceil(4))
}
};
let mut messages: Vec<f64> = x
.messages
.iter()
.map(|m| weight(m.tokens, m.chars))
.collect();
let mut tools: Vec<f64> = x.tools.iter().map(|t| weight(t.tokens, t.chars)).collect();
let raw_total = messages.iter().chain(&tools).fold(0.0, |sum, w| sum + w);
if let Some(prompt) = x.counts.prompt_tokens
&& raw_total > 0.0
{
let scale = f64::from(prompt.value) / raw_total;
for w in messages.iter_mut().chain(tools.iter_mut()) {
*w *= scale;
}
}
let tool_block = tools.iter().fold(0.0, |sum, w| sum + w);
let total = messages.iter().fold(tool_block, |sum, w| sum + w);
let source = if exact && !(messages.is_empty() && tools.is_empty()) {
Source::Tokenizer
} else {
Source::Estimate
};
Weights {
messages,
tools,
tool_block,
total,
source,
}
}
fn sessions(order: &[&Exchange], gap_ms: u64) -> Vec<u32> {
let mut open: BTreeMap<&str, (u32, i64)> = BTreeMap::new();
let mut next = 1u32;
order
.iter()
.map(|x| {
let session = match open.get(x.model.as_str()) {
Some(&(session, last_end))
if x.t_start_ms.saturating_sub(last_end) <= gap_ms as i64 =>
{
session
}
_ => {
next += 1;
next - 1
}
};
let last_end = open
.get(x.model.as_str())
.filter(|(s, _)| *s == session)
.map_or(x.t_end_ms, |&(_, end)| end.max(x.t_end_ms));
open.insert(x.model.as_str(), (session, last_end));
session
})
.collect()
}
fn continues(a: &Exchange, b: &Exchange) -> bool {
let (am, bm) = (&a.messages, &b.messages);
if am.is_empty() || am.len() > bm.len() {
return false;
}
let last = am.len() - 1;
am.iter().zip(bm).enumerate().all(|(i, (x, y))| {
x.role == y.role
&& (x.hash == y.hash || x.role == "system" || (i == last && x.role == "assistant"))
})
}
fn system_hash(x: &Exchange) -> Option<&str> {
x.messages
.iter()
.find(|m| m.role == "system")
.map(|m| m.hash.as_str())
}
fn classify(order: &[&Exchange]) -> Vec<(RowKind, Option<Shape>)> {
let mut kinds = vec![(RowKind::Turn, None); order.len()];
let mut domains: BTreeMap<&str, Vec<usize>> = BTreeMap::new();
for (i, x) in order.iter().enumerate() {
domains.entry(x.model.as_str()).or_default().push(i);
}
for members in domains.values() {
let n = members.len();
let mut continued = vec![false; n];
let mut continuing = vec![false; n];
for a in 0..n {
for b in a + 1..n {
if continues(order[members[a]], order[members[b]]) {
continued[a] = true;
continuing[b] = true;
}
}
}
let any_pair = continuing.iter().any(|&c| c);
let continuing_systems: BTreeSet<Option<&str>> = (0..n)
.filter(|&k| continued[k] || continuing[k])
.map(|k| system_hash(order[members[k]]))
.collect();
for (k, &i) in members.iter().enumerate() {
let x = order[i];
let shape = if summarise(x) {
Some(Shape::Summarise)
} else if structured(x) {
Some(Shape::Structured)
} else if x.tools.is_empty()
&& any_pair
&& !continued[k]
&& !continuing[k]
&& !continuing_systems.contains(&system_hash(x))
{
Some(Shape::Oneshot)
} else {
None
};
if shape.is_some() {
kinds[i] = (RowKind::SideCall, shape);
}
}
let rest: Vec<&Exchange> = members
.iter()
.filter(|&&i| kinds[i].0 != RowKind::SideCall)
.map(|&i| order[i])
.collect();
let stateless = rest.len() >= STATELESS_MIN_REQUESTS
&& rest.windows(2).all(|pair| {
!continues(pair[0], pair[1]) && pair[1].messages.len() <= pair[0].messages.len()
});
if stateless {
for &i in members {
if kinds[i].0 != RowKind::SideCall {
kinds[i] = (RowKind::Stateless, None);
}
}
}
}
kinds
}
fn summarise(x: &Exchange) -> bool {
if !x.tools.is_empty() || x.messages.len() > SIDE_CALL_MAX_MESSAGES {
return false;
}
let total: u64 = x.messages.iter().map(|m| u64::from(m.chars)).sum();
let user: u64 = x
.messages
.iter()
.filter(|m| m.role == "user")
.map(|m| u64::from(m.chars))
.sum();
user > SUMMARISE_MIN_CHARS && (user as f64) > SUMMARISE_USER_SHARE * total as f64
}
fn structured(x: &Exchange) -> bool {
x.tools.is_empty()
&& x.messages.len() <= SIDE_CALL_MAX_MESSAGES
&& (x.params.grammar
|| x.params.response_format
|| x.params
.max_tokens
.is_some_and(|n| n > 0 && n <= STRUCTURED_MAX_TOKENS))
}
fn bulk_rates<'a>(order: &[&'a Exchange]) -> BTreeMap<&'a str, f64> {
let mut sums: BTreeMap<&str, (f64, f64)> = BTreeMap::new();
for x in order {
if let (Some(prefill), Some(evaluated)) = (x.counts.prefill_ms, x.counts.evaluated_tokens)
&& prefill.source == Source::Server
&& evaluated.value >= BULK_MIN_EVALUATED
{
let entry = sums.entry(x.model.as_str()).or_default();
entry.0 += prefill.value;
entry.1 += f64::from(evaluated.value);
}
}
sums.into_iter()
.filter(|(_, (_, evaluated))| *evaluated > 0.0)
.map(|(model, (ms, evaluated))| (model, ms / evaluated))
.collect()
}
fn link(prev: &Exchange, pw: &Weights, cur: &Exchange, cw: &Weights, fresh: bool) -> Link {
if fresh {
let k = prev
.messages
.iter()
.zip(&cur.messages)
.take_while(|(a, b)| a.hash == b.hash)
.count();
let prefix_tokens = match (k, prev.tools_hash == cur.tools_hash) {
(0, _) => 0.0,
(_, same_tools) => cw.messages[..k]
.iter()
.fold(if same_tools { cw.tool_block } else { 0.0 }, |sum, w| {
sum + w
}),
};
return Link {
prefix_len: k,
prefix_tokens,
break_cause: None,
compaction: false,
already_sent: prefix_tokens,
};
}
if prev.tools_hash != cur.tools_hash {
return Link {
prefix_len: 0,
prefix_tokens: 0.0,
break_cause: Some(BreakCause::ToolsChanged),
compaction: false,
already_sent: already_sent(prev, cur, cw, false),
};
}
let overlap = prev.messages.len().min(cur.messages.len());
let k = (0..overlap)
.take_while(|&i| prev.messages[i].hash == cur.messages[i].hash)
.count();
let prefix_tokens = if k == 0 {
0.0
} else {
cw.messages[..k]
.iter()
.fold(cw.tool_block, |sum, w| sum + w)
};
let break_cause = if k < overlap {
if prev.messages[k].role == "system" || cur.messages[k].role == "system" {
Some(BreakCause::SystemChanged)
} else {
Some(BreakCause::HistoryRewritten)
}
} else if cur.messages.len() < prev.messages.len() {
Some(BreakCause::HistoryRewritten)
} else {
None
};
let compaction =
break_cause == Some(BreakCause::HistoryRewritten) && cw.total < COMPACTION_RATIO * pw.total;
let already_sent = if break_cause.is_some() && !compaction {
already_sent(prev, cur, cw, true)
} else {
prefix_tokens
};
Link {
prefix_len: k,
prefix_tokens,
break_cause,
compaction,
already_sent,
}
}
fn already_sent(prev: &Exchange, cur: &Exchange, cw: &Weights, same_tools: bool) -> f64 {
let prev_hashes: BTreeSet<&str> = prev.messages.iter().map(|m| m.hash.as_str()).collect();
let cur_hashes: BTreeSet<&str> = cur.messages.iter().map(|m| m.hash.as_str()).collect();
let tools = if same_tools { cw.tool_block } else { 0.0 };
let messages: f64 = prev
.messages
.iter()
.zip(&cur.messages)
.zip(&cw.messages)
.map(|((p, c), &weight)| {
if p.hash == c.hash {
weight
} else if p.role == c.role
&& !prev_hashes.contains(c.hash.as_str())
&& !cur_hashes.contains(p.hash.as_str())
{
weight * unchanged_share(p, c)
} else {
0.0
}
})
.fold(0.0, |sum, w| sum + w);
tools + messages
}
fn unchanged_share(prev: &Msg, cur: &Msg) -> f64 {
if cur.chars == 0 {
return 0.0;
}
let unchanged = match (&prev.text, &cur.text) {
(Some(a), Some(b)) => {
let a: Vec<char> = a.chars().collect();
let b: Vec<char> = b.chars().collect();
let shortest = a.len().min(b.len());
let head = a.iter().zip(&b).take_while(|(x, y)| x == y).count();
let tail = a
.iter()
.rev()
.zip(b.iter().rev())
.take(shortest - head)
.take_while(|(x, y)| x == y)
.count();
(head + tail) as u32
}
_ => {
let edge = DIGEST_EDGE_CHARS.min(cur.chars);
let head = if prev.head == cur.head { edge } else { 0 };
let tail = if prev.tail == cur.tail { edge } else { 0 };
head + tail
}
};
f64::from(unchanged.min(cur.chars)) / f64::from(cur.chars)
}
fn apply_link(
row: &mut Row,
x: &Exchange,
w: &Weights,
link: &Link,
latest: Option<Cause>,
bulk_rate: Option<f64>,
) {
row.prefix_tokens = Some(round(link.prefix_tokens));
row.break_cause = link.break_cause;
row.compaction = link.compaction;
row.composition = composition(x, w, link.prefix_len);
let counts = x.counts;
if let (Some(prompt), Some(reused)) = (counts.prompt_tokens, counts.reused_tokens)
&& prompt.value > 0
{
let predicted = link.prefix_tokens / f64::from(prompt.value);
let observed = f64::from(reused.value) / f64::from(prompt.value);
row.predicted_reuse = Some(predicted);
row.observed_reuse = Some(observed);
row.cause = cause(
(predicted - observed) * 100.0,
observed <= COLD_SHARE,
link.break_cause.is_some(),
latest,
!row.other_domains_between.is_empty(),
);
}
let (Some(_), Some(evaluated), Some(prefill), Some(reused)) = (
counts.prompt_tokens,
counts.evaluated_tokens,
counts.prefill_ms,
counts.reused_tokens,
) else {
return;
};
row.loss_measured = true;
let charged_break = link.break_cause.filter(|_| !link.compaction);
let reused = f64::from(reused.value);
let evaluated_tokens = f64::from(evaluated.value);
let reusable = match charged_break {
Some(_) => link.prefix_tokens.max(link.already_sent),
None => link.prefix_tokens,
};
let lost = (reusable - reused).clamp(0.0, evaluated_tokens);
let client = match charged_break {
Some(_) => lost.min((reusable - link.prefix_tokens.max(reused)).max(0.0)),
None => 0.0,
};
let external = match row.cause {
Some(
Cause::SideCall
| Cause::OtherConversation
| Cause::OtherModel
| Cause::ServerReset
| Cause::Undetermined,
) => lost - client,
_ => 0.0,
};
let own_rate = (evaluated.value > 0).then(|| prefill.value / evaluated_tokens);
let per_token = match (own_rate, bulk_rate) {
(Some(own), Some(bulk)) => Some(own.min(bulk)),
(own, bulk) => own.or(bulk),
};
let loss = |cause: &str, tokens: f64| {
let tokens = round(tokens);
(tokens > 0).then(|| Loss {
cause: cause.to_string(),
tokens,
ms: per_token.map(|rate| Measured {
value: f64::from(tokens) * rate,
source: prefill.source,
}),
})
};
if let Some(cause) = charged_break {
row.lost_client = loss(cause.name(), client);
}
let external_cause = match row.cause {
Some(Cause::SideCall) => Some("side_call"),
Some(Cause::OtherConversation) => Some("other_conversation"),
Some(Cause::OtherModel) => Some("other_model"),
Some(Cause::ServerReset | Cause::Undetermined) => {
let head = w.tool_block
+ x.messages
.iter()
.zip(&w.messages)
.take_while(|(m, _)| m.role == "system")
.fold(0.0, |sum, (_, w)| sum + w);
Some(if reused < SERVER_EMPTY_BELOW {
"server_empty"
} else if reused <= head * HEAD_MARGIN {
"server_head_only"
} else {
"server_rollback"
})
}
_ => None,
};
if let Some(cause) = external_cause {
row.lost_external = loss(cause, external);
}
}
fn cause(
gap: f64,
cold: bool,
has_break: bool,
latest: Option<Cause>,
other_domains: bool,
) -> Option<Cause> {
if gap < -CAUSE_GAP_POINTS {
None
} else if gap.abs() <= CAUSE_GAP_POINTS {
has_break.then_some(Cause::Structural)
} else if !cold {
Some(if has_break {
Cause::Structural
} else {
Cause::Undetermined
})
} else if latest.is_some() {
latest
} else if other_domains {
Some(Cause::OtherModel)
} else {
Some(Cause::ServerReset)
}
}
fn composition(x: &Exchange, w: &Weights, prefix_len: usize) -> Composition {
let mut system = 0.0;
let mut history = 0.0;
let mut new_content = 0.0;
for (i, (m, &weight)) in x.messages.iter().zip(&w.messages).enumerate() {
if m.role == "system" {
system += weight;
} else if i < prefix_len {
history += weight;
} else {
new_content += weight;
}
}
let measured = |value: f64| Measured {
value: round(value),
source: w.source,
};
Composition {
system: measured(system),
tools: measured(w.tool_block),
history: measured(history),
new_content: measured(new_content),
}
}
fn assign_clients(order: &[&Exchange], rows: &mut [Row]) {
struct Unit<'a> {
agent: &'a str,
tools: BTreeSet<&'a str>,
rows: Vec<usize>,
}
let mut units: Vec<Unit> = Vec::new();
let mut unit_of_conversation: BTreeMap<u32, usize> = BTreeMap::new();
for (i, x) in order.iter().enumerate() {
let agent = x.user_agent.as_deref().unwrap_or(UNKNOWN_CLIENT);
let unit = match rows[i].conversation {
Some(c) if unit_of_conversation.contains_key(&c) => unit_of_conversation[&c],
conversation => {
units.push(Unit {
agent,
tools: BTreeSet::new(),
rows: Vec::new(),
});
if let Some(c) = conversation {
unit_of_conversation.insert(c, units.len() - 1);
}
units.len() - 1
}
};
units[unit]
.tools
.extend(x.tools.iter().map(|t| t.name.as_str()));
units[unit].rows.push(i);
}
let mut parent: Vec<usize> = (0..units.len()).collect();
fn root(parent: &mut [usize], mut i: usize) -> usize {
while parent[i] != i {
parent[i] = parent[parent[i]];
i = parent[i];
}
i
}
for a in 0..units.len() {
for b in a + 1..units.len() {
let (ua, ub) = (&units[a], &units[b]);
if ua.agent == ub.agent
&& !ua.tools.is_empty()
&& !ub.tools.is_empty()
&& (ua.tools.is_subset(&ub.tools) || ub.tools.is_subset(&ua.tools))
{
let (ra, rb) = (root(&mut parent, a), root(&mut parent, b));
parent[ra.max(rb)] = ra.min(rb);
}
}
}
let mut clusters: BTreeMap<&str, Vec<usize>> = BTreeMap::new();
for (i, unit) in units.iter().enumerate() {
let r = root(&mut parent, i);
if !unit.tools.is_empty() && r == i {
clusters.entry(unit.agent).or_default().push(i);
}
}
for (i, unit) in units.iter().enumerate() {
let roots = clusters.get(unit.agent).map_or(&[][..], Vec::as_slice);
let label = if unit.tools.is_empty() || roots.len() <= 1 {
unit.agent.to_string()
} else {
let r = root(&mut parent, i);
let n = roots.iter().position(|&x| x == r).unwrap_or(0) + 1;
format!("{} #{n}", unit.agent)
};
for &row in &unit.rows {
rows[row].client = label.clone();
}
}
}
fn tool_usage(order: &[&Exchange], weights: &[Weights], rows: &[Row]) -> Vec<ToolUsage> {
let mut usage: BTreeMap<(&str, &str), ToolUsage> = BTreeMap::new();
for ((x, w), row) in order.iter().zip(weights).zip(rows) {
let mut seen = BTreeSet::new();
for (tool, &weight) in x.tools.iter().zip(&w.tools) {
if !seen.insert(tool.name.as_str()) {
continue;
}
let entry = usage
.entry((row.client.as_str(), tool.name.as_str()))
.or_insert_with(|| ToolUsage {
client: row.client.clone(),
name: tool.name.clone(),
schema_tokens: 0,
offered: 0,
called: 0,
});
entry.offered += 1;
entry.schema_tokens = entry.schema_tokens.max(round(weight));
}
}
let mut call_ids = BTreeSet::new();
for (x, row) in order.iter().zip(rows) {
let from_response = x.output.tool_calls.iter().map(|call| (call, true));
let from_history = x
.messages
.iter()
.filter(|m| m.role == "assistant")
.flat_map(|m| &m.tool_calls)
.map(|call| (call, false));
for (call, in_response) in from_response.chain(from_history) {
let new = match &call.id {
Some(id) => call_ids.insert(id.clone()),
None => in_response,
};
if new && let Some(entry) = usage.get_mut(&(row.client.as_str(), call.name.as_str())) {
entry.called += 1;
}
}
}
usage.into_values().collect()
}
fn clients(order: &[&Exchange], rows: &[Row]) -> Vec<Client> {
let mut clients: Vec<Client> = Vec::new();
let mut conversations: BTreeMap<&str, BTreeSet<u32>> = BTreeMap::new();
let mut tools: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new();
for (x, row) in order.iter().zip(rows) {
let position = match clients.iter().position(|c| c.client == row.client) {
Some(position) => position,
None => {
clients.push(Client {
client: row.client.clone(),
user_agent: x.user_agent.clone(),
requests: 0,
conversations: 0,
tools: Vec::new(),
});
clients.len() - 1
}
};
clients[position].requests += 1;
if let Some(c) = row.conversation {
conversations
.entry(row.client.as_str())
.or_default()
.insert(c);
}
tools
.entry(row.client.as_str())
.or_default()
.extend(x.tools.iter().map(|t| t.name.as_str()));
}
for client in &mut clients {
client.conversations = conversations
.get(client.client.as_str())
.map_or(0, |c| c.len() as u32);
client.tools = tools
.get(client.client.as_str())
.map(|t| t.iter().map(|name| name.to_string()).collect())
.unwrap_or_default();
}
clients
}
fn round(value: f64) -> u32 {
if value.is_finite() && value > 0.0 {
value.round().min(f64::from(u32::MAX)) as u32
} else {
0
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::hash;
use crate::normalize::ToolDef;
use crate::record::{Output, Params, ToolCall};
fn msg(role: &str, text: &str) -> Msg {
let chars: Vec<char> = text.chars().collect();
let head: String = chars.iter().take(256).collect();
let tail: String = chars[chars.len().saturating_sub(256)..].iter().collect();
Msg {
role: role.into(),
hash: hash::hex(format!("{role}\0{text}").as_bytes()),
chars: chars.len() as u32,
tokens: None,
head: hash::hex(head.as_bytes()),
tail: hash::hex(tail.as_bytes()),
text: Some(text.into()),
tool_calls: Vec::new(),
tool_call_id: None,
name: None,
}
}
fn sized(role: &str, tag: &str, tokens: u32) -> Msg {
let mut m = msg(role, tag);
m.tokens = Some(Measured {
value: tokens,
source: Source::Tokenizer,
});
m
}
fn ex(id: u64, t: i64, messages: Vec<Msg>) -> Exchange {
Exchange {
id,
t_start_ms: t,
t_end_ms: t + 100,
api: Api::OpenAiChat,
model: "m".into(),
user_agent: Some("agent/1".into()),
messages,
tools: Vec::new(),
tools_hash: None,
params: Params::default(),
counts: Counts::default(),
output: Output::default(),
ok: true,
}
}
fn with_tools(mut x: Exchange, names: &[(&str, u32)]) -> Exchange {
x.tools = names
.iter()
.map(|(name, tokens)| ToolDef {
name: (*name).into(),
hash: hash::hex(name.as_bytes()),
chars: 400,
tokens: Some(Measured {
value: *tokens,
source: Source::Tokenizer,
}),
})
.collect();
x.tools_hash = Some(hash::hex(format!("{names:?}").as_bytes()));
x
}
fn model(mut x: Exchange, name: &str) -> Exchange {
x.model = name.into();
x
}
fn server(value: u32) -> Option<Measured<u32>> {
Some(Measured {
value,
source: Source::Server,
})
}
fn served(
mut x: Exchange,
prompt: u32,
reused: u32,
evaluated: u32,
prefill_ms: f64,
) -> Exchange {
x.counts.prompt_tokens = server(prompt);
x.counts.reused_tokens = server(reused);
x.counts.evaluated_tokens = server(evaluated);
x.counts.prefill_ms = Some(Measured {
value: prefill_ms,
source: Source::Server,
});
x
}
fn table(exchanges: &[Exchange]) -> Table {
build(exchanges, &Config::default())
}
fn row(t: &Table, id: u64) -> &Row {
t.rows.iter().find(|r| r.id == id).unwrap()
}
fn conversation(turns: usize) -> Vec<Vec<Msg>> {
let mut history = vec![sized("system", "sys", 1000)];
let mut out = Vec::new();
for i in 0..turns {
history.push(sized("user", &format!("u{i}"), 100));
out.push(history.clone());
history.push(sized("assistant", &format!("a{i}"), 100));
}
out
}
#[test]
fn rows_are_ordered_by_time_then_id() {
let t = table(&[
ex(3, 10, vec![msg("user", "c")]),
ex(2, 10, vec![msg("user", "b")]),
ex(1, 20, vec![msg("user", "a")]),
]);
assert_eq!(
t.rows.iter().map(|r| r.id).collect::<Vec<_>>(),
vec![2, 3, 1]
);
}
#[test]
fn weights_estimate_from_chars_and_scale_to_the_server_total() {
let mut x = ex(
1,
0,
vec![msg("system", &"s".repeat(10)), msg("user", &"u".repeat(30))],
);
let w = weights(&x);
assert_eq!(w.messages, vec![3.0, 8.0]);
assert_eq!(w.source, Source::Estimate);
x.counts.prompt_tokens = server(22);
let w = weights(&x);
assert_eq!(w.messages, vec![6.0, 16.0]);
assert_eq!(w.total, 22.0);
assert_eq!(w.source, Source::Estimate);
}
#[test]
fn tokenizer_counts_are_weights() {
let x = with_tools(ex(1, 0, vec![sized("user", "u", 40)]), &[("ls", 60)]);
let w = weights(&x);
assert_eq!(
(w.messages.clone(), w.tool_block, w.total),
(vec![40.0], 60.0, 100.0)
);
assert_eq!(w.source, Source::Tokenizer);
let c = table(&[x]).rows[0].composition;
assert_eq!(
(c.tools.value, c.new_content.value, c.tools.source),
(60, 40, Source::Tokenizer)
);
}
#[test]
fn sessions_split_on_the_gap_after_the_last_request_ended() {
let mut a = ex(1, 0, vec![msg("user", "a")]);
a.t_end_ms = 200_000;
let b = ex(2, 250_000, vec![msg("user", "b")]);
let c = ex(3, 400_000, vec![msg("user", "c")]);
let t = table(&[a, b, c]);
assert_eq!(row(&t, 1).session, row(&t, 2).session);
assert_ne!(row(&t, 2).session, row(&t, 3).session);
}
#[test]
fn sessions_are_per_domain() {
let t = table(&[
ex(1, 0, vec![msg("user", "a")]),
model(ex(2, 30_000, vec![msg("user", "x")]), "other"),
ex(3, 50_000, vec![msg("user", "b")]),
]);
assert_eq!(row(&t, 1).session, row(&t, 3).session);
assert_ne!(row(&t, 1).session, row(&t, 2).session);
}
#[test]
fn continuation_rules() {
let a = ex(
1,
0,
vec![
msg("system", "s1"),
msg("user", "u"),
msg("assistant", "par"),
],
);
let b = ex(
2,
0,
vec![
msg("system", "s2"),
msg("user", "u"),
msg("assistant", "partial answer"),
msg("user", "more"),
],
);
assert!(continues(&a, &b));
let c = ex(
3,
0,
vec![
msg("system", "s1"),
msg("user", "other"),
msg("assistant", "par"),
],
);
assert!(!continues(&a, &c));
let d = ex(4, 0, vec![msg("system", "s1"), msg("user", "u")]);
assert!(!continues(&a, &d));
let e = ex(
5,
0,
vec![msg("system", "s1"), msg("user", "u"), msg("user", "x")],
);
assert!(!continues(&a, &e));
assert!(!continues(&ex(6, 0, vec![]), &a));
let question = ex(7, 0, vec![msg("system", "s1"), msg("user", "q1")]);
let other = ex(
8,
0,
vec![
msg("system", "s1"),
msg("user", "q2"),
msg("assistant", "x"),
],
);
assert!(!continues(&question, &other));
}
#[test]
fn summarise_side_call() {
let big = "x".repeat(20_001);
let shape = |x: Exchange| table(&[x]).rows[0].clone();
let r = shape(ex(
1,
0,
vec![msg("system", "Summarise."), msg("user", &big)],
));
assert_eq!(
(r.kind, r.side_call_shape),
(RowKind::SideCall, Some(Shape::Summarise))
);
let r = shape(ex(
2,
0,
vec![msg("system", &"s".repeat(3000)), msg("user", &big)],
));
assert_eq!(r.side_call_shape, None);
let r = shape(ex(3, 0, vec![msg("user", &"y".repeat(20_000))]));
assert_eq!(r.side_call_shape, None);
let r = shape(ex(4, 0, vec![msg("user", &"y".repeat(20_001))]));
assert_eq!(r.side_call_shape, Some(Shape::Summarise));
}
#[test]
fn structured_side_call() {
let base = || ex(1, 0, vec![msg("system", "Title?"), msg("user", "text")]);
let shape = |x: Exchange| table(&[x]).rows[0].side_call_shape;
let mut x = base();
x.params.max_tokens = Some(32);
assert_eq!(shape(x), Some(Shape::Structured));
let mut x = base();
x.params.max_tokens = Some(33);
assert_eq!(shape(x), None);
let mut x = base();
x.params.max_tokens = Some(0);
assert_eq!(shape(x), None);
let mut x = base();
x.params.grammar = true;
assert_eq!(shape(x), Some(Shape::Structured));
let mut x = base();
x.params.response_format = true;
assert_eq!(shape(x.clone()), Some(Shape::Structured));
assert_eq!(shape(with_tools(x.clone(), &[("ls", 10)])), None);
x.messages.extend([msg("assistant", "a"), msg("user", "b")]);
assert_eq!(shape(x), None);
}
#[test]
fn oneshot_side_call() {
let turns = conversation(3);
let oneshot = ex(
10,
150,
vec![msg("system", "Pick a filename."), msg("user", "q")],
);
let same_system = ex(
11,
160,
vec![sized("system", "sys", 1000), msg("user", "unrelated")],
);
let with_tools_x = with_tools(
ex(12, 170, vec![msg("system", "Other."), msg("user", "q3")]),
&[("ls", 5)],
);
let t = table(&[
ex(1, 0, turns[0].clone()),
ex(2, 100, turns[1].clone()),
oneshot,
same_system,
with_tools_x,
ex(3, 200, turns[2].clone()),
]);
assert_eq!(row(&t, 10).side_call_shape, Some(Shape::Oneshot));
assert_eq!(row(&t, 11).kind, RowKind::Turn);
assert_eq!(row(&t, 12).kind, RowKind::Turn);
assert_eq!(row(&t, 1).kind, RowKind::Turn);
}
#[test]
fn no_oneshot_without_a_continuing_pair() {
let t = table(&[
ex(1, 0, vec![msg("system", "A"), msg("user", "q1")]),
ex(
2,
10,
vec![
msg("system", "B"),
msg("user", "q2"),
msg("assistant", "x"),
msg("user", "y"),
],
),
]);
assert!(t.rows.iter().all(|r| r.kind == RowKind::Turn));
}
#[test]
fn stateless_domain() {
let classify_rows =
|xs: Vec<Exchange>| table(&xs).rows.iter().map(|r| r.kind).collect::<Vec<_>>();
let q = |id, text: &str| {
ex(
id,
id as i64 * 10,
vec![msg("system", "Classify."), msg("user", text)],
)
};
assert_eq!(
classify_rows(vec![q(1, "a"), q(2, "b"), q(3, "c")]),
vec![RowKind::Stateless; 3]
);
assert_eq!(
classify_rows(vec![q(1, "a"), q(2, "b")]),
vec![RowKind::Turn; 2]
);
let mut grows = q(3, "c");
grows.messages.push(msg("user", "d"));
assert_eq!(
classify_rows(vec![q(1, "a"), q(2, "b"), grows]),
vec![RowKind::Turn; 3]
);
let t = table(&[q(1, "a"), q(2, "b"), q(3, "c")]);
assert!(
t.rows
.iter()
.all(|r| r.prefix_tokens.is_none() && r.turn.is_none())
);
}
#[test]
fn links_record_what_ran_between() {
let turns = conversation(2);
let mut side = ex(5, 50, vec![msg("system", "Title?"), msg("user", "t")]);
side.params.max_tokens = Some(8);
let other = model(ex(6, 60, vec![msg("user", "o")]), "embedder");
let t = table(&[
ex(1, 0, turns[0].clone()),
side,
other,
ex(2, 100, turns[1].clone()),
]);
let r = row(&t, 2);
assert_eq!(r.side_calls_between, vec![5]);
assert_eq!(r.other_domains_between, vec![6]);
assert_eq!((row(&t, 1).turn, r.turn), (Some(1), Some(2)));
assert_eq!(row(&t, 5).turn, None);
assert_eq!(row(&t, 1).prefix_tokens, None);
}
#[test]
fn growing_conversation_has_no_break() {
let turns = conversation(2);
let t = table(&[ex(1, 0, turns[0].clone()), ex(2, 100, turns[1].clone())]);
let r = row(&t, 2);
assert_eq!(r.break_cause, None);
assert_eq!(r.prefix_tokens, Some(1100));
let c = r.composition;
assert_eq!(
(c.system.value, c.history.value, c.new_content.value),
(1000, 100, 200)
);
let first = row(&t, 1).composition;
assert_eq!(
(
first.system.value,
first.history.value,
first.new_content.value
),
(1000, 0, 100)
);
}
#[test]
fn an_empty_prefix_is_positive_zero() {
let turns = conversation(2);
let mut sys = turns[1].clone();
sys[0] = sized("system", "other", 1000);
let t = link_pair(|x| served(x, 1300, 0, 1300, 1300.0), sys);
let predicted = row(&t, 2).predicted_reuse.unwrap();
assert!(predicted == 0.0 && predicted.is_sign_positive());
assert_eq!(serde_json::to_string(&predicted).unwrap(), "0.0");
}
#[test]
fn tool_change_breaks_the_prefix() {
let turns = conversation(2);
let a = with_tools(ex(1, 0, turns[0].clone()), &[("ls", 50)]);
let b = with_tools(ex(2, 100, turns[1].clone()), &[("ls", 50), ("cat", 50)]);
let t = table(&[a, b]);
let r = row(&t, 2);
assert_eq!(r.break_cause, Some(BreakCause::ToolsChanged));
assert_eq!(r.prefix_tokens, Some(0));
assert_eq!(r.composition.history.value, 0);
}
#[test]
fn tools_are_not_reusable_when_the_first_message_changed() {
let turns = conversation(2);
let a = with_tools(ex(1, 0, turns[0].clone()), &[("ls", 50)]);
let mut changed = turns[1].clone();
changed[0] = sized("system", "other", 1000);
let b = with_tools(ex(2, 100, changed.clone()), &[("ls", 50)]);
let t = table(&[
served(a, 1150, 0, 1150, 1150.0),
served(b, 1350, 0, 1350, 1350.0),
]);
let r = row(&t, 2);
assert_eq!(
(r.prefix_tokens, r.break_cause),
(Some(0), Some(BreakCause::SystemChanged))
);
assert_eq!(r.cause, Some(Cause::Structural));
assert_eq!(r.lost_client.clone().map(|l| l.tokens), Some(150));
assert_eq!(r.lost_external, None);
let mut edited = turns[1].clone();
edited[1] = sized("user", "u0 edited", 100);
let c = with_tools(ex(3, 200, edited), &[("ls", 50)]);
let t = table(&[with_tools(ex(1, 0, turns[0].clone()), &[("ls", 50)]), c]);
assert_eq!(row(&t, 3).prefix_tokens, Some(1050));
}
#[test]
fn system_and_history_breaks() {
let turns = conversation(2);
let mut sys = turns[1].clone();
sys[0] = sized("system", "sys v2", 1000);
let t = table(&[ex(1, 0, turns[0].clone()), ex(2, 100, sys)]);
assert_eq!(row(&t, 2).break_cause, Some(BreakCause::SystemChanged));
assert_eq!(row(&t, 2).prefix_tokens, Some(0));
let turns = conversation(3);
let mut edit = turns[2].clone();
edit[3] = sized("user", "u1 edited", 100);
let t = table(&[ex(1, 0, turns[1].clone()), ex(2, 100, edit)]);
assert_eq!(row(&t, 2).break_cause, Some(BreakCause::HistoryRewritten));
assert_eq!(row(&t, 2).prefix_tokens, Some(1200));
assert_eq!(
(row(&t, 2).linked_to, row(&t, 2).new_conversation),
(Some(1), false)
);
}
#[test]
fn shorter_request_matched_to_its_end_is_rewritten_history() {
let turns = conversation(3);
let t = table(&[ex(1, 0, turns[2].clone()), ex(2, 100, turns[1].clone())]);
let r = row(&t, 2);
assert_eq!(r.break_cause, Some(BreakCause::HistoryRewritten));
assert!(!r.compaction);
}
#[test]
fn compaction_is_a_big_shrink() {
let long = conversation(20).pop().unwrap();
let short = vec![
long[0].clone(),
long[1].clone(),
sized("user", "summary so far", 200),
sized("user", "next", 100),
];
let t = table(&[ex(1, 0, long.clone()), ex(2, 100, short)]);
let r = row(&t, 2);
assert_eq!(r.break_cause, Some(BreakCause::HistoryRewritten));
assert!(r.compaction);
let system_only = vec![long[0].clone(), sized("user", "summary so far", 200)];
let t = table(&[ex(1, 0, long), ex(2, 100, system_only)]);
let r = row(&t, 2);
assert!(r.new_conversation && !r.compaction);
assert_eq!(r.break_cause, None);
}
#[test]
fn already_sent_counts_equal_and_edited_messages_but_not_moved_ones() {
let prev = ex(
1,
0,
vec![
sized("system", "s", 100),
sized("user", "a", 100),
sized("assistant", "b", 100),
],
);
let mut cur = prev.clone();
cur.id = 2;
cur.messages[0] = sized("system", "s2", 100);
let w = weights(&cur);
assert_eq!(already_sent(&prev, &cur, &w, true), 250.0);
cur.messages[0] = sized("system", "t", 100);
assert_eq!(already_sent(&prev, &cur, &w, true), 200.0);
let mut moved = prev.clone();
moved.messages.insert(1, sized("user", "inserted", 100));
let w = weights(&moved);
assert_eq!(already_sent(&prev, &moved, &w, true), 100.0);
}
#[test]
fn unchanged_share_from_text() {
let a = msg("system", "You are helpful. Today is Monday. Be brief.");
let b = msg("system", "You are helpful. Today is Tuesday. Be brief.");
let share = unchanged_share(&a, &b);
let expected = (f64::from(b.chars) - 4.0) / f64::from(b.chars);
assert!((share - expected).abs() < 1e-9, "{share} vs {expected}");
assert_eq!(unchanged_share(&a, &a), 1.0);
}
#[test]
fn unchanged_share_from_head_and_tail_in_digest_mode() {
let body = "x".repeat(1000);
let mut a = msg("system", &format!("{body}A{body}"));
let mut b = msg("system", &format!("{body}B{body}"));
a.text = None;
b.text = None;
assert_eq!(unchanged_share(&a, &b), 512.0 / 2001.0);
let mut c = msg("system", &format!("C{body}{body}"));
c.text = None;
assert_eq!(unchanged_share(&a, &c), 256.0 / 2001.0);
}
fn link_pair(cur_counts: impl Fn(Exchange) -> Exchange, cur_messages: Vec<Msg>) -> Table {
let turns = conversation(2);
table(&[
served(ex(1, 0, turns[0].clone()), 1100, 0, 1100, 1100.0),
cur_counts(ex(2, 100, cur_messages)),
])
}
#[test]
fn full_reuse_loses_nothing() {
let turns = conversation(2);
let t = link_pair(|x| served(x, 1300, 1100, 200, 200.0), turns[1].clone());
let r = row(&t, 2);
assert_eq!(
(r.predicted_reuse, r.observed_reuse),
(Some(1100.0 / 1300.0), Some(1100.0 / 1300.0))
);
assert_eq!(r.cause, None);
assert!(r.loss_measured);
assert_eq!(
(r.lost_client.clone(), r.lost_external.clone()),
(None, None)
);
}
#[test]
fn cold_server_after_a_clean_prefix_is_a_server_reset() {
let turns = conversation(2);
let t = link_pair(|x| served(x, 1300, 0, 1300, 1300.0), turns[1].clone());
let r = row(&t, 2);
assert_eq!(r.cause, Some(Cause::ServerReset));
let loss = r.lost_external.clone().unwrap();
assert_eq!((loss.cause.as_str(), loss.tokens), ("server_empty", 1100));
assert_eq!(
loss.ms,
Some(Measured {
value: 1100.0,
source: Source::Server
})
);
assert_eq!(r.lost_client, None);
}
#[test]
fn server_loss_kinds_follow_where_reuse_stopped() {
let turns = conversation(6);
let kind = |reused: u32| {
let t = table(&[
served(ex(1, 0, turns[4].clone()), 1900, 0, 1900, 1900.0),
served(
ex(2, 100, turns[5].clone()),
2100,
reused,
2100 - reused,
2100.0,
),
]);
let r = row(&t, 2).clone();
(r.cause, r.lost_external.map(|l| (l.cause, l.tokens)))
};
assert_eq!(
kind(64),
(
Some(Cause::ServerReset),
Some(("server_empty".into(), 1836))
)
);
assert_eq!(
kind(1100),
(
Some(Cause::Undetermined),
Some(("server_head_only".into(), 800))
)
);
assert_eq!(
kind(1400),
(
Some(Cause::Undetermined),
Some(("server_rollback".into(), 500))
)
);
assert_eq!(kind(1500), (None, None));
}
#[test]
fn cause_rules_in_order() {
let side = Some(Cause::SideCall);
let other = Some(Cause::OtherConversation);
assert_eq!(cause(-21.0, true, true, None, false), None);
assert_eq!(
cause(-20.0, true, true, None, false),
Some(Cause::Structural)
);
assert_eq!(cause(20.0, true, false, None, false), None);
assert_eq!(
cause(21.0, false, true, None, false),
Some(Cause::Structural)
);
assert_eq!(
cause(21.0, false, false, side, true),
Some(Cause::Undetermined)
);
assert_eq!(cause(21.0, true, true, side, true), Some(Cause::SideCall));
assert_eq!(
cause(21.0, true, false, other, true),
Some(Cause::OtherConversation)
);
assert_eq!(
cause(21.0, true, false, None, true),
Some(Cause::OtherModel)
);
assert_eq!(
cause(21.0, true, false, None, false),
Some(Cause::ServerReset)
);
}
#[test]
fn the_latest_request_in_between_takes_the_blame() {
let a = thread("main agent", 2);
let b = thread("subagent", 2);
let side = || {
let mut x = ex(0, 0, vec![msg("system", "Title?"), msg("user", "t")]);
x.params.max_tokens = Some(8);
x
};
let at = |mut x: Exchange, id: u64, t: i64| {
x.id = id;
x.t_start_ms = t;
x
};
let run = |middle: Vec<Exchange>| {
let mut xs = vec![served(ex(1, 0, a[0].clone()), 1100, 0, 1100, 1100.0)];
xs.extend(middle);
xs.push(served(ex(9, 100, a[1].clone()), 1300, 0, 1300, 1300.0));
table(&xs).rows.iter().find(|r| r.id == 9).unwrap().cause
};
assert_eq!(
run(vec![
at(side(), 2, 10),
ex(3, 20, b[0].clone()),
ex(4, 30, b[1].clone())
]),
Some(Cause::OtherConversation)
);
assert_eq!(
run(vec![
ex(3, 20, b[0].clone()),
ex(4, 30, b[1].clone()),
at(side(), 5, 40)
]),
Some(Cause::SideCall)
);
}
#[test]
fn side_call_between_turns_takes_the_blame() {
let turns = conversation(2);
let mut side = ex(5, 50, vec![msg("system", "Title?"), msg("user", "t")]);
side.params.max_tokens = Some(8);
let t = table(&[
served(ex(1, 0, turns[0].clone()), 1100, 0, 1100, 1100.0),
side,
served(ex(2, 100, turns[1].clone()), 1300, 0, 1300, 1300.0),
]);
let r = row(&t, 2);
assert_eq!(r.cause, Some(Cause::SideCall));
assert_eq!(
r.lost_external.clone().map(|l| (l.cause, l.tokens)),
Some(("side_call".into(), 1100))
);
}
#[test]
fn other_model_between_turns_takes_the_blame() {
let turns = conversation(2);
let t = table(&[
served(ex(1, 0, turns[0].clone()), 1100, 0, 1100, 1100.0),
model(ex(6, 50, vec![msg("user", "x")]), "other"),
served(ex(2, 100, turns[1].clone()), 1300, 0, 1300, 1300.0),
]);
assert_eq!(row(&t, 2).cause, Some(Cause::OtherModel));
assert_eq!(
row(&t, 2).lost_external.clone().map(|l| l.cause),
Some("other_model".into())
);
}
#[test]
fn client_loss_is_what_was_sent_before_but_not_reused() {
let turns = conversation(2);
let mut edited = turns[1].clone();
edited[0] = sized("system", "other", 1000);
let t = link_pair(|x| served(x, 1300, 0, 1300, 1300.0), edited);
let r = row(&t, 2);
assert_eq!(r.break_cause, Some(BreakCause::SystemChanged));
assert_eq!(r.cause, Some(Cause::Structural));
let client = r.lost_client.clone().unwrap();
assert_eq!(
(client.cause.as_str(), client.tokens),
("system_changed", 100)
);
assert_eq!(r.lost_external, None);
}
#[test]
fn client_loss_counts_the_unchanged_share_of_an_edited_system_prompt() {
let base = "x".repeat(4000);
let prev = vec![msg("system", &format!("{base} Monday")), msg("user", "q")];
let cur = vec![
msg("system", &format!("{base} Tuesday")),
msg("user", "q"),
msg("assistant", "a"),
msg("user", "r"),
];
let t = table(&[
served(ex(1, 0, prev), 1000, 0, 1000, 1000.0),
served(ex(2, 100, cur), 1010, 0, 1010, 1010.0),
]);
let r = row(&t, 2);
assert_eq!(r.break_cause, Some(BreakCause::SystemChanged));
let client = r.lost_client.clone().unwrap().tokens;
assert!((1000..=1010).contains(&client), "{client}");
}
#[test]
fn compaction_charges_nothing_to_the_client() {
let long = conversation(20).pop().unwrap();
let short = vec![
long[0].clone(),
long[1].clone(),
sized("user", "summary", 200),
sized("user", "next", 100),
];
let t = table(&[
served(ex(1, 0, long), 4900, 0, 4900, 4900.0),
served(ex(2, 100, short), 1400, 1100, 300, 300.0),
]);
let r = row(&t, 2);
assert!(r.compaction);
assert_eq!(r.lost_client, None);
assert_eq!(r.cause, Some(Cause::Structural));
assert_eq!(r.lost_external, None);
}
#[test]
fn loss_never_exceeds_what_was_evaluated() {
let turns = conversation(2);
let t = link_pair(|x| served(x, 1300, 0, 300, 300.0), turns[1].clone());
assert_eq!(row(&t, 2).lost_external.clone().unwrap().tokens, 300);
}
#[test]
fn loss_needs_every_counter() {
let turns = conversation(2);
let without_prefill = |x: Exchange| {
let mut x = served(x, 1300, 0, 1300, 0.0);
x.counts.prefill_ms = None;
x
};
let t = link_pair(without_prefill, turns[1].clone());
let r = row(&t, 2);
assert!(!r.loss_measured);
assert_eq!(r.cause, Some(Cause::ServerReset));
assert_eq!(r.lost_external, None);
let without_reuse = |x: Exchange| {
let mut x = served(x, 1300, 0, 1300, 1300.0);
x.counts.reused_tokens = None;
x
};
let t = link_pair(without_reuse, turns[1].clone());
let r = row(&t, 2);
assert!(!r.loss_measured);
assert_eq!((r.cause, r.observed_reuse), (None, None));
}
#[test]
fn time_uses_the_cheaper_of_own_and_bulk_rate() {
let turns = conversation(3);
let t = table(&[
served(ex(1, 0, turns[0].clone()), 1100, 0, 1100, 550.0),
served(ex(2, 100, turns[1].clone()), 1300, 1100, 200, 100.0),
served(ex(3, 200, turns[2].clone()), 1500, 0, 20, 400.0),
]);
let loss = row(&t, 3).lost_external.clone().unwrap();
assert_eq!(loss.tokens, 20);
let bulk = (550.0 + 100.0) / (1100.0 + 200.0);
assert!((loss.ms.unwrap().value - 20.0 * bulk).abs() < 1e-9);
}
#[test]
fn proxy_timed_loss_keeps_the_proxy_source() {
let turns = conversation(2);
let proxy_timed = |x: Exchange| {
let mut x = served(x, 1300, 0, 1300, 0.0);
x.counts.prefill_ms = Some(Measured {
value: 2600.0,
source: Source::Proxy,
});
x
};
let t = link_pair(proxy_timed, turns[1].clone());
let loss = row(&t, 2).lost_external.clone().unwrap();
assert_eq!(
loss.ms,
Some(Measured {
value: 1100.0,
source: Source::Proxy
})
);
}
fn agent(mut x: Exchange, user_agent: Option<&str>) -> Exchange {
x.user_agent = user_agent.map(str::to_string);
x
}
fn thread(system: &str, turns: usize) -> Vec<Vec<Msg>> {
let mut history = vec![sized("system", system, 1000)];
let mut out = Vec::new();
for i in 0..turns {
history.push(sized("user", &format!("{system} u{i}"), 100));
out.push(history.clone());
history.push(sized("assistant", &format!("{system} a{i}"), 100));
}
out
}
#[test]
fn interleaved_conversations_link_to_their_own_turns() {
let a = thread("main agent", 3);
let b = thread("subagent", 2);
let t = table(&[
ex(1, 0, a[0].clone()),
ex(2, 10, b[0].clone()),
ex(3, 20, a[1].clone()),
ex(4, 30, b[1].clone()),
ex(5, 40, a[2].clone()),
]);
type Link = (u64, Option<u64>, Option<u32>, Option<u32>);
let links: Vec<Link> = t
.rows
.iter()
.map(|r| (r.id, r.linked_to, r.conversation, r.turn))
.collect();
assert_eq!(
links,
vec![
(1, None, Some(1), Some(1)),
(2, None, Some(2), Some(1)),
(3, Some(1), Some(1), Some(2)),
(4, Some(2), Some(2), Some(2)),
(5, Some(3), Some(1), Some(3)),
]
);
assert!(t.rows.iter().all(|r| r.break_cause.is_none()));
assert_eq!(row(&t, 3).other_conversations_between, vec![2]);
assert_eq!(row(&t, 5).other_conversations_between, vec![4]);
assert!(row(&t, 3).side_calls_between.is_empty());
}
#[test]
fn another_conversation_in_between_takes_the_blame() {
let a = thread("main agent", 2);
let b = thread("subagent", 2);
let t = table(&[
served(ex(1, 0, a[0].clone()), 1100, 0, 1100, 1100.0),
served(ex(2, 10, b[0].clone()), 1100, 0, 1100, 1100.0),
served(ex(3, 15, b[1].clone()), 1300, 1100, 200, 200.0),
served(ex(4, 20, a[1].clone()), 1300, 0, 1300, 1300.0),
]);
let r = row(&t, 4);
assert_eq!(r.other_conversations_between, vec![2, 3]);
assert_eq!(r.cause, Some(Cause::OtherConversation));
let loss = r.lost_external.clone().unwrap();
assert_eq!(
(loss.cause.as_str(), loss.tokens),
("other_conversation", 1100)
);
}
#[test]
fn a_new_conversation_shares_only_the_system_prompt() {
let first = thread("agent", 2);
let second = vec![
sized("system", "agent", 1000),
sized("user", "a new task", 100),
];
let t = table(&[
served(ex(1, 0, first[1].clone()), 1300, 0, 1300, 1300.0),
served(ex(2, 100, second), 1100, 0, 1100, 1100.0),
]);
let r = row(&t, 2);
assert!(r.new_conversation);
assert_eq!(
(r.linked_to, r.conversation, r.turn),
(Some(1), Some(2), Some(1))
);
assert_eq!((r.break_cause, r.prefix_tokens), (None, Some(1000)));
assert_eq!(r.lost_client, None);
assert_eq!(r.cause, Some(Cause::ServerReset));
assert_eq!(r.lost_external.clone().map(|l| l.tokens), Some(1000));
}
#[test]
fn a_continuation_wins_over_a_more_recent_similar_request() {
let main = thread("agent", 2);
let other = vec![
sized("system", "agent", 1000),
sized("user", "a new task", 100),
];
let t = table(&[
ex(1, 0, main[0].clone()),
ex(2, 10, other),
ex(3, 20, main[1].clone()),
]);
assert_eq!(row(&t, 3).linked_to, Some(1));
assert_eq!(row(&t, 3).conversation, Some(1));
assert_eq!(row(&t, 3).other_conversations_between, vec![2]);
}
#[test]
fn different_user_agents_are_never_linked() {
let a = thread("agent", 2);
let t = table(&[
agent(ex(1, 0, a[0].clone()), Some("one/1")),
agent(ex(2, 10, a[1].clone()), Some("two/1")),
]);
assert_eq!(row(&t, 2).linked_to, None);
assert_ne!(row(&t, 1).conversation, row(&t, 2).conversation);
assert_eq!(
(row(&t, 1).client.as_str(), row(&t, 2).client.as_str()),
("one/1", "two/1")
);
}
#[test]
fn clients_split_a_shared_user_agent_by_tool_set() {
let main = thread("main", 2);
let sub = thread("sub", 1);
let other = thread("other harness", 1);
let sdk = Some("OpenAI/Python 1.54.0");
let t = table(&[
agent(
with_tools(ex(1, 0, main[0].clone()), &[("read", 10), ("write", 10)]),
sdk,
),
agent(with_tools(ex(2, 10, sub[0].clone()), &[("read", 10)]), sdk),
agent(
with_tools(ex(3, 20, other[0].clone()), &[("grep", 10), ("edit", 10)]),
sdk,
),
agent(
with_tools(ex(4, 30, main[1].clone()), &[("read", 10), ("write", 10)]),
sdk,
),
agent(ex(5, 40, vec![msg("user", "hello")]), None),
]);
let clients: Vec<&str> = t.rows.iter().map(|r| r.client.as_str()).collect();
assert_eq!(
clients,
vec![
"OpenAI/Python 1.54.0 #1",
"OpenAI/Python 1.54.0 #1",
"OpenAI/Python 1.54.0 #2",
"OpenAI/Python 1.54.0 #1",
"unknown"
]
);
let summary: Vec<(&str, u32, u32, Vec<&str>)> = t
.clients
.iter()
.map(|c| {
(
c.client.as_str(),
c.requests,
c.conversations,
c.tools.iter().map(String::as_str).collect(),
)
})
.collect();
assert_eq!(
summary,
vec![
("OpenAI/Python 1.54.0 #1", 3, 2, vec!["read", "write"]),
("OpenAI/Python 1.54.0 #2", 1, 1, vec!["edit", "grep"]),
("unknown", 1, 0, vec![]),
]
);
assert_eq!(t.clients[2].user_agent, None);
let tools: Vec<(&str, &str, u32)> = t
.tools
.iter()
.map(|u| (u.client.as_str(), u.name.as_str(), u.offered))
.collect();
assert_eq!(
tools,
vec![
("OpenAI/Python 1.54.0 #1", "read", 3),
("OpenAI/Python 1.54.0 #1", "write", 2),
("OpenAI/Python 1.54.0 #2", "edit", 1),
("OpenAI/Python 1.54.0 #2", "grep", 1),
]
);
}
#[test]
fn one_tool_set_keeps_the_plain_user_agent() {
let a = thread("agent", 1);
let t = table(&[
agent(
with_tools(ex(1, 0, a[0].clone()), &[("read", 10)]),
Some("aider/0.86"),
),
agent(
ex(2, 10, vec![msg("system", "Title?"), msg("user", "t")]),
Some("aider/0.86"),
),
]);
assert!(t.rows.iter().all(|r| r.client == "aider/0.86"));
assert_eq!(t.clients.len(), 1);
}
#[test]
fn tool_usage_counts_offers_and_distinct_calls() {
let mut a = with_tools(
ex(1, 0, vec![sized("user", "u", 100)]),
&[("ls", 50), ("cat", 70)],
);
a.output.tool_calls = vec![ToolCall {
id: Some("c1".into()),
name: "ls".into(),
}];
let mut b = with_tools(
ex(2, 100, vec![sized("user", "u", 100)]),
&[("ls", 60), ("cat", 70)],
);
let mut assistant = sized("assistant", "call", 10);
assistant.tool_calls = vec![
ToolCall {
id: Some("c1".into()),
name: "ls".into(),
},
ToolCall {
id: Some("c0".into()),
name: "ls".into(),
},
];
b.messages.push(assistant);
b.output.tool_calls = vec![
ToolCall {
id: None,
name: "cat".into(),
},
ToolCall {
id: None,
name: "unknown".into(),
},
];
let t = table(&[a, b]);
assert_eq!(
t.tools,
vec![
ToolUsage {
client: "agent/1".into(),
name: "cat".into(),
schema_tokens: 70,
offered: 2,
called: 1
},
ToolUsage {
client: "agent/1".into(),
name: "ls".into(),
schema_tokens: 60,
offered: 2,
called: 2
},
]
);
}
}