use std::io::Cursor;
use llmp::{
BreakCause, Cause, Confidence, Config, FindingKind, Report, RowKind, Shape, Source, analyze,
read_jsonl,
};
use serde_json::{Value, json};
const T0: i64 = 1_789_900_000_000;
fn system(day: &str) -> Value {
json!({"role": "system", "content": format!("{}Today is {day}.", "You are a careful coding agent. ".repeat(125))})
}
fn user(n: u32) -> Value {
json!({"role": "user", "content": format!("{}(request {n})", "Please change the code as described. ".repeat(10))})
}
fn call(n: u32) -> Value {
json!({"role": "assistant", "content": null, "tool_calls": [{"id": format!("call_{n}"), "type": "function",
"function": {"name": "read_file", "arguments": "{\"path\":\"src/main.rs\"}"}}]})
}
fn result(n: u32) -> Value {
json!({"role": "tool", "tool_call_id": format!("call_{n}"), "content": "fn main() {}\n".repeat(20)})
}
fn answer(n: u32) -> Value {
json!({"role": "assistant", "content": format!("{}(answer {n})", "Done, the change is in place. ".repeat(8))})
}
fn tools() -> Value {
json!([
{"type": "function", "function": {"name": "read_file", "description": "Read a file from the workspace.",
"parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}},
{"type": "function", "function": {"name": "write_file", "description": "Write a file in the workspace.",
"parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}}}}
])
}
fn estimate(messages: &Value, tools: &Value) -> u32 {
let chars = |s: &str| s.chars().count() as u32;
let mut total = 0;
for m in messages.as_array().unwrap() {
let mut n = m["content"].as_str().map_or(0, chars);
if let Some(calls) = m.get("tool_calls") {
n += chars(&calls.to_string());
}
total += n.div_ceil(4);
}
for t in tools.as_array().unwrap() {
total += chars(&t.to_string()).div_ceil(4);
}
total
}
struct Chat {
id: u64,
t: i64,
model: &'static str,
messages: Value,
tools: Value,
extra: Value,
prompt: u32,
cached: u32,
reply: Value,
}
fn chat(c: Chat) -> Value {
let mut body = json!({"model": c.model, "messages": c.messages, "tools": c.tools});
for (k, v) in c.extra.as_object().unwrap() {
body[k] = v.clone();
}
let evaluated = c.prompt - c.cached;
json!({
"v": 1, "id": c.id, "conn": 1,
"t_start_ms": c.t, "t_first_byte_ms": c.t + 100, "t_end_ms": c.t + 2000,
"request": {"method": "POST", "path": "/v1/chat/completions", "headers": {"content-type": "application/json"},
"body_bytes": 1, "body": body, "body_text": null, "digest": null},
"response": {"status": 200, "headers": {"content-type": "application/json"}, "stream": null,
"chunks": null, "bad_chunks": null, "body_bytes": 1, "truncated": false,
"body": {
"choices": [{"index": 0, "message": c.reply, "finish_reason": "stop"}],
"usage": {"prompt_tokens": c.prompt, "completion_tokens": 5},
"timings": {"cache_n": c.cached, "prompt_n": evaluated, "prompt_ms": f64::from(evaluated), "predicted_ms": 50.0}
},
"body_text": null, "digest": null},
"error": null, "tokens": null
})
}
fn recording() -> String {
let t1 = json!([system("Monday"), user(1)]);
let t2 = json!([system("Monday"), user(1), call(1), result(1), user(2)]);
let t3 = json!([
system("Monday"),
user(1),
call(1),
result(1),
user(2),
answer(2),
user(3)
]);
let t4 = json!([
system("Tuesday"),
user(1),
call(1),
result(1),
user(2),
answer(2),
user(3),
answer(3),
user(4)
]);
let t5 = json!([
system("Tuesday"),
user(1),
call(1),
result(1),
user(2),
answer(2),
user(3),
answer(3),
user(4),
answer(4),
user(5)
]);
let side =
json!([{"role": "system", "content": "Write a five-word title for this chat."}, user(1)]);
let tools = tools();
let (e1, e2, e3, e4, e5) = (
estimate(&t1, &tools),
estimate(&t2, &tools),
estimate(&t3, &tools),
estimate(&t4, &tools),
estimate(&t5, &tools),
);
let tools_only = estimate(&json!([]), &tools);
let q = |id, t, messages: &Value, prompt, cached, reply: Value| Chat {
id,
t: T0 + t,
model: "qwen",
messages: messages.clone(),
tools: tools.clone(),
extra: json!({}),
prompt,
cached,
reply,
};
let lines = [
chat(q(1, 0, &t1, e1, 0, call(1))),
chat(q(2, 10_000, &t2, e2, e1, answer(2))),
chat(Chat {
tools: json!([]),
extra: json!({"max_tokens": 16}),
..q(
3,
12_000,
&side,
estimate(&side, &json!([])),
0,
json!({"role": "assistant", "content": "Fixing the main file"}),
)
}),
chat(q(4, 20_000, &t3, e3, 0, answer(3))),
chat(q(5, 30_000, &t4, e4, tools_only + 990, answer(4))),
chat(Chat {
model: "small",
tools: json!([]),
..q(
6,
35_000,
&json!([{"role": "user", "content": "Classify: bug or feature?"}]),
12,
0,
json!({"role": "assistant", "content": "bug"}),
)
}),
chat(q(7, 40_000, &t5, e5, 0, answer(5))),
json!({"v": 1, "id": 8, "conn": 2, "t_start_ms": T0 + 41_000, "t_first_byte_ms": T0 + 41_001, "t_end_ms": T0 + 41_002,
"request": {"method": "GET", "path": "/v1/models", "headers": {}, "body_bytes": 0, "body": null, "body_text": null, "digest": null},
"response": {"status": 200, "headers": {}, "stream": null, "chunks": null, "bad_chunks": null, "body_bytes": 2,
"truncated": false, "body": {"data": []}, "body_text": null, "digest": null},
"error": null, "tokens": null}),
json!({"v": 1, "id": 9, "conn": 3, "t_start_ms": T0 + 42_000, "t_first_byte_ms": null, "t_end_ms": T0 + 42_001,
"request": {"method": "POST", "path": "/v1/chat/completions", "headers": {}, "body_bytes": 2, "body": {}, "body_text": null, "digest": null},
"response": null, "error": "upstream_unreachable", "tokens": null}),
];
let mut out: String = lines.iter().map(|l| format!("{l}\n")).collect();
out.push_str("{\"v\":1,\"id\":10, truncated line\n");
out
}
fn report() -> Report {
let read = read_jsonl(Cursor::new(recording()));
assert_eq!(read.bad_lines.len(), 1);
analyze(&read.records, &Config::default())
}
fn row(report: &Report, id: u64) -> &llmp::Row {
report.rows.iter().find(|r| r.id == id).unwrap()
}
#[test]
fn coverage_counts_every_record() {
let r = report();
let c = &r.coverage;
assert_eq!((c.records, c.analysed), (9, 7));
assert_eq!(
(
c.skipped.not_chat,
c.skipped.error,
c.skipped.non_2xx,
c.skipped.unparsed
),
(1, 1, 0, 0)
);
assert_eq!(c.paths.get("/v1/models"), Some(&1));
assert_eq!(
(c.turns, c.turns_with_reuse, c.turns_with_timing),
(6, 6, 6)
);
let window = r.window.unwrap();
assert_eq!((window.start_ms, window.end_ms), (T0, T0 + 42_001));
assert_eq!(r.servers.len(), 1);
assert_eq!(
serde_json::to_value(&r.servers[0]).unwrap()["fields"],
json!([
"prompt_tokens",
"completion_tokens",
"reused_tokens",
"evaluated_tokens",
"prefill_ms",
"decode_ms"
])
);
}
#[test]
fn rows_explain_each_request() {
let r = report();
assert_eq!(
r.rows.iter().map(|r| r.id).collect::<Vec<_>>(),
vec![1, 2, 3, 4, 5, 6, 7]
);
assert_eq!(
(row(&r, 1).kind, row(&r, 1).turn, row(&r, 1).prefix_tokens),
(RowKind::Turn, Some(1), None)
);
let second = row(&r, 2);
assert_eq!((second.break_cause, second.cause), (None, None));
assert!(second.loss_measured && second.lost_client.is_none() && second.lost_external.is_none());
let side = row(&r, 3);
assert_eq!(
(side.kind, side.side_call_shape),
(RowKind::SideCall, Some(Shape::Structured))
);
let after_side = row(&r, 4);
assert_eq!(after_side.side_calls_between, vec![3]);
assert_eq!(after_side.cause, Some(Cause::SideCall));
let loss = after_side.lost_external.clone().unwrap();
assert_eq!(loss.cause, "side_call");
assert!(loss.tokens > 1000);
assert_eq!(loss.ms.unwrap().source, Source::Server);
let new_date = row(&r, 5);
assert_eq!(new_date.break_cause, Some(BreakCause::SystemChanged));
assert_eq!(new_date.cause, None);
let client = new_date.lost_client.clone().unwrap();
assert_eq!(client.cause, "system_changed");
assert!(client.tokens > 300, "{client:?}");
assert_eq!(
(row(&r, 6).model.as_str(), row(&r, 6).session),
("small", 2)
);
let after_other = row(&r, 7);
assert_eq!(after_other.other_domains_between, vec![6]);
assert_eq!(after_other.cause, Some(Cause::OtherModel));
}
#[test]
fn every_number_reads_from_the_rows() {
let r = report();
let mut tokens = 0u64;
for row in &r.rows {
let evaluated = row.counts.evaluated_tokens.map_or(0, |e| e.value);
let lost: u32 = [&row.lost_client, &row.lost_external]
.into_iter()
.flatten()
.map(|l| l.tokens)
.sum();
assert!(
lost <= evaluated,
"row {} lost {lost} of {evaluated}",
row.id
);
tokens += u64::from(lost);
}
assert_eq!(r.time.lost_tokens, tokens);
let by_cause: u64 = r.time.by_cause.values().map(|c| c.tokens).sum();
assert_eq!(by_cause, tokens);
assert_eq!(
r.time
.by_cause
.keys()
.map(String::as_str)
.collect::<Vec<_>>(),
vec!["other_model", "side_call", "system_changed"]
);
let lost_ms = r.time.lost_ms.unwrap();
let cause_ms: f64 = r
.time
.by_cause
.values()
.filter_map(|c| c.ms)
.map(|m| m.value)
.sum();
assert!((lost_ms.value - cause_ms).abs() < 1e-6);
let prefill: f64 = r
.rows
.iter()
.filter_map(|row| row.counts.prefill_ms)
.map(|m| m.value)
.sum();
assert!((r.time.prefill_ms.unwrap().value - prefill).abs() < 1e-6);
let system: u64 = r
.rows
.iter()
.map(|row| u64::from(row.composition.system.value))
.sum();
assert_eq!(r.composition.system, system);
}
#[test]
fn findings_are_typed_and_ranked() {
let r = report();
let kinds: Vec<FindingKind> = r.findings.iter().map(|f| f.kind).collect();
for expected in [
FindingKind::SideCallOnConversationModel,
FindingKind::PrefixBreak,
FindingKind::UnusedTool,
FindingKind::LargeSystemPrompt,
] {
assert!(
kinds.contains(&expected),
"{expected:?} missing from {kinds:?}"
);
}
assert!(!kinds.contains(&FindingKind::OtherModelEviction));
let side = r
.findings
.iter()
.find(|f| f.kind == FindingKind::SideCallOnConversationModel)
.unwrap();
assert_eq!(side.model.as_deref(), Some("qwen"));
assert_eq!(side.rows, vec![3]);
assert!(side.lost_tokens.unwrap() > 1000);
assert_eq!(side.evidence["shapes"], json!({"structured": 1}));
let prefix = r
.findings
.iter()
.find(|f| f.kind == FindingKind::PrefixBreak)
.unwrap();
assert_eq!(prefix.rows, vec![5]);
assert_eq!(prefix.evidence["break_cause"], "system_changed");
let unused = r
.findings
.iter()
.find(|f| f.kind == FindingKind::UnusedTool)
.unwrap();
assert_eq!(unused.evidence["tool"], "write_file");
assert_eq!(unused.confidence, Confidence::Low);
assert!(
!r.findings
.iter()
.any(|f| f.kind == FindingKind::UnusedTool && f.evidence["tool"] == "read_file")
);
let tools: Vec<(&str, u32, u32)> = r
.tools
.iter()
.map(|t| (t.name.as_str(), t.offered, t.called))
.collect();
assert_eq!(tools, vec![("read_file", 5, 1), ("write_file", 5, 0)]);
let ranks: Vec<Confidence> = r.findings.iter().map(|f| f.confidence).collect();
let mut sorted = ranks.clone();
sorted.sort();
assert_eq!(ranks, sorted);
}
#[test]
fn json_has_no_nan_and_unmeasured_is_null() {
let r = report();
let json = serde_json::to_string(&r).unwrap();
assert!(!json.contains("NaN"));
let value: Value = serde_json::from_str(&json).unwrap();
assert_eq!(value["v"], 1);
assert_eq!(value["rows"][0]["counts"]["load_ms"], Value::Null);
let back: Report = serde_json::from_str(&json).unwrap();
assert_eq!(back, r);
}
#[test]
fn summary_is_short_and_names_the_losses() {
let text = report().summary();
assert!(text.lines().count() <= 15, "{text}");
assert!(text.contains("time lost re-processing prompts:"), "{text}");
assert!(text.contains("side-calls"), "{text}");
assert!(text.contains("coverage: 7 of 9 records analysed"), "{text}");
assert!(text.contains("findings:"), "{text}");
}
#[test]
fn nothing_to_analyse() {
let r = analyze(&[], &Config::default());
assert_eq!(r.window, None);
assert_eq!(r.time.lost_ms, None);
assert!(r.rows.is_empty() && r.findings.is_empty());
assert!(r.summary().contains("not measured"));
}
#[cfg(feature = "cli")]
mod cli {
use super::*;
use std::path::PathBuf;
use std::process::Command;
fn file(name: &str, content: &str) -> PathBuf {
let path = PathBuf::from(env!("CARGO_TARGET_TMPDIR")).join(name);
std::fs::write(&path, content).unwrap();
path
}
fn llmp(args: &[&str]) -> std::process::Output {
Command::new(env!("CARGO_BIN_EXE_llmp"))
.args(args)
.output()
.unwrap()
}
#[test]
fn analyze_prints_the_summary() {
let path = file("cli_summary.jsonl", &recording());
let out = llmp(&["analyze", path.to_str().unwrap()]);
assert_eq!(out.status.code(), Some(0));
let text = String::from_utf8(out.stdout).unwrap();
assert_eq!(text, {
let mut r = report();
r.coverage.bad_lines = 1;
r.summary()
});
assert!(text.contains("1 bad line;"));
assert!(String::from_utf8_lossy(&out.stderr).contains(":10:"));
}
#[test]
fn analyze_json_is_the_library_report() {
let path = file("cli_json.jsonl", &recording());
let out = llmp(&[
"analyze",
path.to_str().unwrap(),
"--json",
"--gap-seconds",
"60",
]);
assert_eq!(out.status.code(), Some(0));
let parsed: Report = serde_json::from_slice(&out.stdout).unwrap();
let mut expected = report();
expected.coverage.bad_lines = 1;
assert_eq!(parsed, expected);
}
#[test]
fn several_files_are_merged_by_time_and_renumbered() {
let text = recording();
let lines: Vec<&str> = text.lines().collect();
let a = file("cli_part_a.jsonl", &format!("{}\n", lines[4..].join("\n")));
let b = file("cli_part_b.jsonl", &format!("{}\n", lines[..4].join("\n")));
let out = llmp(&[
"analyze",
a.to_str().unwrap(),
b.to_str().unwrap(),
"--json",
]);
assert_eq!(out.status.code(), Some(0));
let parsed: Report = serde_json::from_slice(&out.stdout).unwrap();
assert_eq!(
parsed.rows.iter().map(|r| r.id).collect::<Vec<_>>(),
vec![1, 2, 3, 4, 5, 6, 7]
);
assert_eq!(parsed.coverage.records, 9);
assert_eq!(parsed.rows[3].cause, Some(Cause::SideCall));
}
#[test]
fn exit_codes() {
assert_eq!(llmp(&["analyze"]).status.code(), Some(2));
assert_eq!(
llmp(&["analyze", "x.jsonl", "--bogus"]).status.code(),
Some(2)
);
assert_eq!(
llmp(&["analyze", "x.jsonl", "--gap-seconds", "soon"])
.status
.code(),
Some(2)
);
assert_eq!(
llmp(&["analyze", "/nonexistent/llmp.jsonl"]).status.code(),
Some(1)
);
}
}