mod code_syntax;
mod context_faithfulness;
mod context_ladder;
mod edit_format;
mod json_output;
mod max_tokens_compliance;
mod multi_turn_memory;
mod multi_turn_task_sequencing;
mod one_shot_tool_plan;
mod parallel_tool_scale;
mod streaming_tool_calls;
mod system_message_adherence;
mod token_efficiency;
mod tool_calling;
mod vision;
mod xml_fallback;
pub use code_syntax::probe_code_syntax;
pub use context_faithfulness::probe_context_faithfulness;
pub use context_ladder::{ContextLadder, probe_effective_context_tokens};
pub use edit_format::{probe_search_replace, probe_unified_diff};
pub use json_output::{probe_instruction_following, probe_json_output};
pub use max_tokens_compliance::probe_max_tokens_compliance;
pub use multi_turn_memory::probe_multi_turn_memory;
pub use multi_turn_task_sequencing::probe_multi_turn_task_sequencing;
pub use one_shot_tool_plan::probe_one_shot_tool_plan;
pub use parallel_tool_scale::probe_parallel_tool_scale;
pub use streaming_tool_calls::probe_streaming_tool_calls;
pub use system_message_adherence::probe_system_message_adherence;
pub use token_efficiency::probe_token_efficiency;
pub use tool_calling::{
probe_complex_tool_calling, probe_nested_arguments, probe_tool_calling, probe_tool_selection,
};
pub use vision::probe_vision;
pub use xml_fallback::probe_xml_tool_calling;
use crate::client::{
ProbeContent, ProbeFinish, ProbeMessage, ProbeResponse, ProbeRole, ProbeTool, ProbeToolCall,
};
use crate::error::ProbeError;
pub(crate) fn user_text(text: impl Into<String>) -> ProbeMessage {
ProbeMessage {
role: ProbeRole::User,
content: ProbeContent::Text(text.into()),
tool_calls: None,
tool_call_id: None,
}
}
pub(crate) fn system_text(text: impl Into<String>) -> ProbeMessage {
ProbeMessage {
role: ProbeRole::System,
content: ProbeContent::Text(text.into()),
tool_calls: None,
tool_call_id: None,
}
}
pub(crate) fn assistant_text(text: impl Into<String>) -> ProbeMessage {
ProbeMessage {
role: ProbeRole::Assistant,
content: ProbeContent::Text(text.into()),
tool_calls: None,
tool_call_id: None,
}
}
pub(crate) fn assistant_tool_calls(
text: impl Into<String>,
tool_calls: Vec<ProbeToolCall>,
) -> ProbeMessage {
ProbeMessage {
role: ProbeRole::Assistant,
content: ProbeContent::Text(text.into()),
tool_calls: Some(tool_calls),
tool_call_id: None,
}
}
pub(crate) fn tool_result(
tool_call_id: impl Into<String>,
text: impl Into<String>,
) -> ProbeMessage {
ProbeMessage {
role: ProbeRole::Tool,
content: ProbeContent::Text(text.into()),
tool_calls: None,
tool_call_id: Some(tool_call_id.into()),
}
}
pub(crate) fn extract_json_from_text(text: &str) -> &str {
let trimmed = text.trim();
if let Some(body) = fenced_json_body(trimmed) {
return body;
}
let object = match (trimmed.find('{'), trimmed.rfind('}')) {
(Some(start), Some(end)) if end > start => Some((start, end)),
_ => None,
};
let array = match (trimmed.find('['), trimmed.rfind(']')) {
(Some(start), Some(end)) if end > start => Some((start, end)),
_ => None,
};
match (object, array) {
(Some((os, oe)), Some((as_, ae))) if as_ < os && ae >= oe => {
return &trimmed[as_..=ae];
}
(Some((start, end)), _) => return &trimmed[start..=end],
(None, Some((start, end))) => return &trimmed[start..=end],
_ => {}
}
trimmed
}
fn fenced_json_body(trimmed: &str) -> Option<&str> {
let start = trimmed.find("```")?;
let after = &trimmed[start + 3..];
let (inner, _) = after.split_once("```")?;
let inner = inner.trim_start_matches('\r');
let body = if let Some((first, rest)) = inner.split_once('\n') {
let tag = first.trim().trim_end_matches('\r');
if is_fence_language_tag(tag) {
rest
} else {
inner
}
} else {
inner
};
let body = body.trim();
if body.starts_with('{') || body.starts_with('[') {
Some(body)
} else {
None
}
}
fn is_fence_language_tag(tag: &str) -> bool {
!tag.is_empty()
&& tag
.chars()
.all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '_'))
}
pub(crate) fn utf8_prefix(s: &str, max: usize) -> &str {
if s.len() <= max {
return s;
}
let mut end = max;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
&s[..end]
}
pub(crate) fn has_visible_arg_text(s: &str) -> bool {
s.chars().any(|c| {
!c.is_whitespace()
&& !matches!(
c,
'\u{200B}' | '\u{200C}' | '\u{200D}' | '\u{2060}' | '\u{FEFF}'
)
})
}
pub(crate) fn nonempty_string_arg(
args: &serde_json::Map<String, serde_json::Value>,
key: &str,
) -> bool {
args.get(key)
.and_then(|v| v.as_str())
.is_some_and(has_visible_arg_text)
}
pub(crate) fn nonempty_string_arg_any(
args: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> bool {
keys.iter().any(|key| nonempty_string_arg(args, key))
}
pub(crate) fn tool(name: &str, description: &str, parameters: serde_json::Value) -> ProbeTool {
ProbeTool {
name: name.to_string(),
description: description.to_string(),
parameters,
}
}
pub fn refuse_truncated_tool_call(resp: &ProbeResponse) -> Result<(), ProbeError> {
if resp.finish == ProbeFinish::Length && resp.tool_calls.is_empty() {
Err(ProbeError::Transient(
"response truncated before a tool call".into(),
))
} else {
Ok(())
}
}
pub fn refuse_truncated_incomplete(finish: ProbeFinish, score: f32) -> Result<(), ProbeError> {
if finish == ProbeFinish::Length && score < 1.0 {
Err(ProbeError::Transient(
"response truncated before a complete tool call".into(),
))
} else {
Ok(())
}
}
#[cfg(test)]
pub(crate) mod test_support {
use std::collections::VecDeque;
use std::future::Future;
use std::sync::Mutex;
use crate::client::{
ProbeClient, ProbeContent, ProbeContentPart, ProbeFinish, ProbeRequest, ProbeResponse,
ProbeRole, ProbeStreamChunk, ProbeToolCall,
};
use crate::error::ProbeError;
use futures::Stream;
pub(crate) struct MockLlm {
pub(crate) response: ProbeResponse,
}
impl ProbeClient for MockLlm {
fn chat(
&self,
_req: ProbeRequest,
) -> impl Future<Output = Result<ProbeResponse, ProbeError>> + Send {
let resp = self.response.clone();
async move { Ok(resp) }
}
fn stream_chat(
&self,
_req: ProbeRequest,
) -> impl Stream<Item = Result<ProbeStreamChunk, ProbeError>> + Send {
futures::stream::empty()
}
fn model_id(&self) -> &str {
"test-model"
}
fn provider(&self) -> &str {
"test-provider"
}
}
pub(crate) struct SequentialMock {
responses: Mutex<VecDeque<ProbeResponse>>,
}
impl SequentialMock {
pub(crate) fn new(responses: Vec<ProbeResponse>) -> Self {
Self {
responses: Mutex::new(VecDeque::from(responses)),
}
}
}
impl ProbeClient for SequentialMock {
fn chat(
&self,
_req: ProbeRequest,
) -> impl Future<Output = Result<ProbeResponse, ProbeError>> + Send {
let next = self
.responses
.lock()
.expect("sequential mock lock")
.pop_front()
.unwrap_or_else(|| text_response("done"));
async move { Ok(next) }
}
fn stream_chat(
&self,
_req: ProbeRequest,
) -> impl Stream<Item = Result<ProbeStreamChunk, ProbeError>> + Send {
futures::stream::empty()
}
fn model_id(&self) -> &str {
"test-model"
}
fn provider(&self) -> &str {
"test-provider"
}
}
pub(crate) fn tool_call_response() -> ProbeResponse {
ProbeResponse {
text: String::new(),
tool_calls: vec![ProbeToolCall {
id: "call_1".into(),
name: "read_file".into(),
arguments: serde_json::json!({"path": "/tmp/test.txt"})
.as_object()
.unwrap()
.clone(),
}],
finish: ProbeFinish::ToolCalls,
usage: None,
}
}
pub(crate) fn text_response(text: &str) -> ProbeResponse {
ProbeResponse {
text: text.to_string(),
tool_calls: Vec::new(),
finish: ProbeFinish::Stop,
usage: None,
}
}
pub(crate) fn length_text_response(text: &str) -> ProbeResponse {
ProbeResponse {
text: text.to_string(),
tool_calls: Vec::new(),
finish: ProbeFinish::Length,
usage: None,
}
}
pub(crate) fn multi_tool_call_response(calls: Vec<ProbeToolCall>) -> ProbeResponse {
ProbeResponse {
text: String::new(),
tool_calls: calls,
finish: ProbeFinish::ToolCalls,
usage: None,
}
}
pub(crate) struct RecordingMock {
inner: MockLlm,
pub(crate) requests: Mutex<Vec<ProbeRequest>>,
}
impl RecordingMock {
pub(crate) fn new(response: ProbeResponse) -> Self {
Self {
inner: MockLlm { response },
requests: Mutex::new(Vec::new()),
}
}
}
impl ProbeClient for RecordingMock {
fn chat(
&self,
req: ProbeRequest,
) -> impl Future<Output = Result<ProbeResponse, ProbeError>> + Send {
self.requests.lock().expect("lock").push(req.clone());
self.inner.chat(req)
}
fn stream_chat(
&self,
req: ProbeRequest,
) -> impl Stream<Item = Result<ProbeStreamChunk, ProbeError>> + Send {
self.requests.lock().expect("lock").push(req.clone());
self.inner.stream_chat(req)
}
fn model_id(&self) -> &str {
self.inner.model_id()
}
fn provider(&self) -> &str {
self.inner.provider()
}
}
pub(crate) fn request_user_text(req: &ProbeRequest) -> String {
let mut out = String::new();
for message in &req.messages {
if message.role != ProbeRole::User {
continue;
}
match &message.content {
ProbeContent::Text(text) => out.push_str(text),
ProbeContent::Parts(parts) => {
for part in parts {
if let ProbeContentPart::Text { text } = part {
out.push_str(text);
}
}
}
}
}
out
}
}
#[cfg(test)]
mod nonempty_string_arg_tests {
use super::nonempty_string_arg;
fn args(value: serde_json::Value) -> serde_json::Map<String, serde_json::Value> {
value.as_object().unwrap().clone()
}
#[test]
fn nonempty_string_arg_rejects_empty_and_whitespace() {
assert!(!nonempty_string_arg(
&args(serde_json::json!({"path": ""})),
"path"
));
assert!(!nonempty_string_arg(
&args(serde_json::json!({"path": " "})),
"path"
));
assert!(!nonempty_string_arg(
&args(serde_json::json!({"path": "\n"})),
"path"
));
assert!(!nonempty_string_arg(
&args(serde_json::json!({"path": 1})),
"path"
));
assert!(!nonempty_string_arg(&args(serde_json::json!({})), "path"));
assert!(!nonempty_string_arg(
&args(serde_json::json!({"path": "\u{200b}"})),
"path"
));
assert!(nonempty_string_arg(
&args(serde_json::json!({"path": "/tmp/a"})),
"path"
));
}
}
#[cfg(test)]
mod extract_json_tests {
use super::extract_json_from_text;
#[test]
fn extract_json_from_bare_object() {
assert_eq!(extract_json_from_text(r#" {"a": 1} "#), r#"{"a": 1}"#);
}
#[test]
fn extract_json_from_fenced_block() {
let input = "```json\n{\"a\": 1}\n```";
assert_eq!(extract_json_from_text(input), "{\"a\": 1}");
}
#[test]
fn extract_json_from_uppercase_json_fence() {
let input = "```JSON\n{\"a\": 1}\n```";
assert_eq!(extract_json_from_text(input), "{\"a\": 1}");
}
#[test]
fn extract_json_from_jsonc_fence_falls_through_to_object() {
let input = "```jsonc\n{\"a\": 1}\n```";
assert_eq!(extract_json_from_text(input), "{\"a\": 1}");
}
#[test]
fn extract_json_does_not_peel_array_wrapper() {
let input = r#"[{"word": "hello", "length": 5, "reversed": "olleh"}]"#;
assert_eq!(extract_json_from_text(input), input);
}
#[test]
fn extract_json_keeps_array_when_prose_wraps_it() {
let input = r#"Here: [{"word": "hello", "length": 5, "reversed": "olleh"}]"#;
assert_eq!(
extract_json_from_text(input),
r#"[{"word": "hello", "length": 5, "reversed": "olleh"}]"#
);
}
#[test]
fn extract_json_keeps_object_after_citation_brackets() {
let input = r#"See [1] {"word": "hello", "length": 5, "reversed": "olleh"}"#;
assert_eq!(
extract_json_from_text(input),
r#"{"word": "hello", "length": 5, "reversed": "olleh"}"#
);
}
#[test]
fn extract_json_unclosed_bracket_still_takes_object() {
let input = r#"[unclosed {"word": "hello", "length": 5, "reversed": "olleh"}"#;
assert_eq!(
extract_json_from_text(input),
r#"{"word": "hello", "length": 5, "reversed": "olleh"}"#
);
}
}
#[cfg(test)]
mod refuse_truncated_tests {
use super::test_support::{text_response, tool_call_response};
use super::{refuse_truncated_incomplete, refuse_truncated_tool_call};
use crate::client::{ProbeFinish, ProbeResponse};
use crate::error::ProbeError;
#[test]
fn refuse_truncated_tool_call_errors_on_length_without_tools() {
let resp = ProbeResponse {
text: "leftover reasoning".into(),
tool_calls: Vec::new(),
finish: ProbeFinish::Length,
usage: None,
};
let err = refuse_truncated_tool_call(&resp).expect_err("must refuse");
assert!(
matches!(&err, ProbeError::Transient(msg) if msg.contains("truncated")),
"{err:?}"
);
}
#[test]
fn refuse_truncated_tool_call_allows_stop_without_tools() {
let resp = text_response("I would read the file");
assert!(refuse_truncated_tool_call(&resp).is_ok());
}
#[test]
fn refuse_truncated_tool_call_allows_length_with_tools() {
let mut resp = tool_call_response();
resp.finish = ProbeFinish::Length;
assert!(refuse_truncated_tool_call(&resp).is_ok());
}
#[test]
fn refuse_truncated_incomplete_errors_on_length_below_strong() {
let err = refuse_truncated_incomplete(ProbeFinish::Length, 0.5).expect_err("must refuse");
assert!(
matches!(&err, ProbeError::Transient(msg) if msg.contains("truncated")),
"{err:?}"
);
}
#[test]
fn refuse_truncated_incomplete_allows_length_at_strong() {
assert!(refuse_truncated_incomplete(ProbeFinish::Length, 1.0).is_ok());
}
#[test]
fn refuse_truncated_incomplete_allows_stop_below_strong() {
assert!(refuse_truncated_incomplete(ProbeFinish::Stop, 0.5).is_ok());
}
}
#[cfg(test)]
mod length_policy_tests {
use super::test_support::{MockLlm, length_text_response};
use super::*;
use crate::error::ProbeError;
use crate::types::ProbeResult;
async fn run(name: &str) -> Result<ProbeResult, ProbeError> {
let llm = MockLlm {
response: length_text_response("partial"),
};
match name {
"code_syntax" => probe_code_syntax(&llm).await,
"context_faithfulness" => probe_context_faithfulness(&llm).await,
"json_output" => probe_json_output(&llm).await,
"instruction_following" => probe_instruction_following(&llm).await,
"search_replace" => probe_search_replace(&llm).await,
"unified_diff" => probe_unified_diff(&llm).await,
"max_tokens_compliance" => probe_max_tokens_compliance(&llm).await,
"multi_turn_memory" => probe_multi_turn_memory(&llm).await,
"multi_turn_task_sequencing" => probe_multi_turn_task_sequencing(&llm).await,
"one_shot_tool_plan" => probe_one_shot_tool_plan(&llm).await,
"parallel_tool_scale" => probe_parallel_tool_scale(&llm).await,
"system_message_adherence" => probe_system_message_adherence(&llm).await,
"token_efficiency" => probe_token_efficiency(&llm).await,
"complex_tool_calling" => probe_complex_tool_calling(&llm).await,
"nested_arguments" => probe_nested_arguments(&llm).await,
"tool_calling" => probe_tool_calling(&llm).await,
"tool_selection" => probe_tool_selection(&llm).await,
"vision" => probe_vision(&llm).await,
"xml_tool_calling" => probe_xml_tool_calling(&llm).await,
other => panic!("unknown probe {other}"),
}
}
#[tokio::test]
async fn length_partial_is_refused_or_allowlisted() {
const ALLOW: &[&str] = &["max_tokens_compliance", "token_efficiency"];
const PROBES: &[&str] = &[
"code_syntax",
"context_faithfulness",
"json_output",
"instruction_following",
"search_replace",
"unified_diff",
"max_tokens_compliance",
"multi_turn_memory",
"multi_turn_task_sequencing",
"one_shot_tool_plan",
"parallel_tool_scale",
"system_message_adherence",
"token_efficiency",
"complex_tool_calling",
"nested_arguments",
"tool_calling",
"tool_selection",
"vision",
"xml_tool_calling",
];
for name in PROBES {
match run(name).await {
Ok(pr) => {
assert!(
ALLOW.contains(name),
"{name} scored {} on Length partial; must refuse or be allow-listed",
pr.score
);
}
Err(ProbeError::Transient(msg)) => {
assert!(msg.contains("truncated"), "{name}: {msg}");
assert!(
!ALLOW.contains(name),
"{name} is allow-listed but refused: {msg}"
);
}
Err(other) => panic!("{name}: unexpected error {other:?}"),
}
}
}
}