Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions crates/tinyjevclient/src/response/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,14 +166,31 @@ fn validate_distribution(
validate_probability(*probability, name)?;
}
let sum: f64 = probabilities.values().sum();
if (sum - 1.0).abs() > PROBABILITY_TOLERANCE {
if (sum - 1.0).abs() > distribution_tolerance(probabilities.len()) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Update the specified probability-sum contract

This changes accepted response behavior to an option-count-dependent tolerance, but docs/specs/system-one-client.md still requires every distribution to differ from 1.0 by at most 0.000001. Since callers can no longer rely on the implemented spec—and the repository explicitly treats specs as accepted behavior—update the specification alongside this validation change.

AGENTS.md reference: AGENTS.md:L64-L68

Useful? React with 👍 / 👎.

return Err(Error::invalid_response(format!(
"{name} probabilities must sum to one"
"{name} probabilities must sum to one (sum {sum:.6} over {} options)",
probabilities.len()
)));
}
Ok(())
}

/// How far a distribution's sum may stray from one.
///
/// Providers serialise each probability with a fixed number of decimals, so
/// the rounding error grows with the number of options. Measured on the
/// `OpenRouter` System One endpoint (2026-09): two decimals per option, and a
/// 21-option Choice answered with probabilities summing to 0.99. The
/// tolerance is therefore half a unit in the second decimal per option,
/// floored at `PROBABILITY_TOLERANCE` so a two-option answer is held as
/// tightly as before. It is a bound on rounding, not on the model: a
/// distribution that is off by more than that is still rejected.
fn distribution_tolerance(options: usize) -> f64 {
// A Choice holds at most 255 options, so the cast is exact.
let options = f64::from(u32::try_from(options).unwrap_or(u32::MAX));
PROBABILITY_TOLERANCE.max(options * 0.005)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Bound the scaled tolerance before validation becomes vacuous

For valid 200–255-option Choice requests, options * 0.005 is at least 1.0, so a response containing 0.0 for every matching option passes both this sum check and the highest-probability-choice check. The relaxation also starts immediately—a two-option distribution summing to 0.995 is now accepted despite the stated intent to preserve the old strict behavior. Rework or cap the allowance so rounding tolerance cannot admit completely non-normalized distributions.

AGENTS.md reference: AGENTS.md:L66-L68

Useful? React with 👍 / 👎.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

priority high critique confident

Keep the distribution tolerance from accepting arbitrary sums

For a 255-option Choice, this tolerance is 1.275, so a response with every probability equal to 0.0 has sum difference 1.0 and is accepted. That contradicts the comment that distributions farther from one are rejected and allows a completely uninformative response through validation. Bound the tolerance so it cannot make materially invalid sums pass, or validate the rounded distribution with a stricter invariant.


Additional security observation

priority medium confident

Bound the size-based distribution tolerance

[RULE] unbounded-validation-tolerance

With the maximum 255 options, this returns a tolerance of 1.275. Because each probability is only checked to be within [0, 1], a distribution containing all zeroes has a sum difference of 1.0 and is therefore accepted. That can make an invalid response pass validation, and downstream choice validation may accept any label as tying for the highest probability. Keep the rounding allowance bounded so it cannot accept materially non-normalized distributions; the tolerance should account for serialization rounding without exceeding a meaningful distribution error bound.

[RULE] invalid-distribution-acceptance ·

}

fn validate_probability(value: f64, name: &str) -> Result<()> {
if !value.is_finite() || !(0.0..=1.0).contains(&value) {
return Err(Error::invalid_response(format!(
Expand Down
17 changes: 17 additions & 0 deletions crates/tinyjevclient/src/response/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,3 +230,20 @@ fn accepts_the_exact_score_rounding_boundary() {
score.score = 0.05;
boundary.validate_for(&request()).unwrap();
}

#[test]
fn many_option_distributions_tolerate_per_option_rounding() {
// 21 options rounded to two decimals summing to 0.99: what the OpenRouter
// endpoint answers. Rejected before the size-aware tolerance.
let mut probabilities = std::collections::BTreeMap::new();
for i in 0..20 {
probabilities.insert(format!("o{i}"), 0.04);
}
probabilities.insert("o20".to_owned(), 0.19);
assert!((probabilities.values().sum::<f64>() - 0.99).abs() < 1e-9);
assert!(validate_distribution(&probabilities, "choice").is_ok());

// Two options are still held tightly.
let two = std::collections::BTreeMap::from([("a".to_owned(), 0.6), ("b".to_owned(), 0.39)]);
assert!(validate_distribution(&two, "choice").is_err());
}
Loading