use std::collections::HashMap;
use std::path::PathBuf;
use newt_core::TokenEstimation;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolConformance {
Native,
TextMode,
NoTools,
}
impl ToolConformance {
pub fn symbol(&self) -> &'static str {
match self {
Self::Native => "✓ native",
Self::TextMode => "~ text ",
Self::NoTools => "✗ none ",
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TuneConfidence {
#[default]
None,
Low,
Medium,
High,
}
impl TuneConfidence {
pub fn promote(&self) -> Self {
match self {
Self::None => Self::Low,
Self::Low => Self::Medium,
Self::Medium | Self::High => Self::High,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CapabilityEntry {
pub conformance: ToolConformance,
pub tested_date: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_window: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub safe_context: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub overflow_at: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_ok_input: Option<u32>,
#[serde(default)]
pub consecutive_ok: u32,
#[serde(default)]
pub tune_confidence: TuneConfidence,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tune_date: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub estimate_ratio: Option<f32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub emits_thinking: Option<bool>,
#[serde(default)]
pub accounting_version: u32,
}
pub const ACCOUNTING_VERSION: u32 = 1;
impl Default for CapabilityEntry {
fn default() -> Self {
Self {
conformance: ToolConformance::NoTools,
tested_date: String::new(),
context_window: None,
safe_context: None,
overflow_at: None,
max_ok_input: None,
consecutive_ok: 0,
tune_confidence: TuneConfidence::None,
tune_date: None,
estimate_ratio: None,
emits_thinking: None,
accounting_version: ACCOUNTING_VERSION,
}
}
}
impl CapabilityEntry {
pub fn record_success(&mut self, input_tokens: u32, today: &str) -> bool {
let mut changed = false;
if self.max_ok_input.map(|m| input_tokens > m).unwrap_or(true) {
self.max_ok_input = Some(input_tokens);
changed = true;
}
self.consecutive_ok = self.consecutive_ok.saturating_add(1);
if self.consecutive_ok >= 5 && self.tune_confidence != TuneConfidence::High {
self.tune_confidence = self.tune_confidence.promote();
self.tune_date = Some(today.to_string());
self.consecutive_ok = 0;
changed = true;
}
changed
}
pub fn record_context_window_400(&mut self, hard_limit: u32, today: &str) -> bool {
let new_cap = (hard_limit as u64 * 80 / 100) as u32;
self.max_ok_input = Some(new_cap);
self.consecutive_ok = 0;
self.tune_confidence = TuneConfidence::Low;
self.tune_date = Some(today.to_string());
if self.safe_context.map(|s| new_cap < s).unwrap_or(true) {
self.safe_context = Some(new_cap);
}
true }
pub fn record_overflow(&mut self, input_tokens: u32, today: &str) -> bool {
let new_safe = input_tokens * 75 / 100;
self.overflow_at = Some(input_tokens);
self.consecutive_ok = 0;
self.tune_confidence = TuneConfidence::Low;
self.tune_date = Some(today.to_string());
let changed = self.safe_context.map(|s| new_safe < s).unwrap_or(true);
if changed {
self.safe_context = Some(new_safe);
}
if self.max_ok_input.map(|m| new_safe < m).unwrap_or(false) {
self.max_ok_input = Some(new_safe);
}
true }
pub fn record_accepted_prompt(&mut self, prompt_tokens: u32, today: &str) -> bool {
if self.max_ok_input.map(|m| prompt_tokens > m).unwrap_or(true) {
self.max_ok_input = Some(prompt_tokens);
self.tune_date = Some(today.to_string());
return true;
}
false
}
pub fn record_estimate_sample(&mut self, observed: u32, estimated: usize) -> bool {
if estimated == 0 {
return false;
}
let raw = observed as f32 / estimated as f32;
if raw < 0.5 {
return false;
}
let sample = raw.clamp(0.5, 3.0);
let new = match self.estimate_ratio {
None => sample,
Some(old) => (0.75 * old + 0.25 * sample).clamp(0.5, 3.0),
};
let dirty = match self.estimate_ratio {
None => true,
Some(old) => (new - old).abs() > 0.01,
};
self.estimate_ratio = Some(new);
dirty
}
pub fn record_thinking_only(&mut self) -> bool {
if self.emits_thinking == Some(true) {
return false;
}
self.emits_thinking = Some(true);
true
}
}
pub fn apply_observation(
entry: &mut CapabilityEntry,
obs: &newt_core::RoundObservation,
today: &str,
) -> bool {
match *obs {
newt_core::RoundObservation::Accepted {
prompt_tokens,
estimated_tokens,
} => {
entry.record_accepted_prompt(prompt_tokens, today)
| entry.record_estimate_sample(prompt_tokens, estimated_tokens)
}
newt_core::RoundObservation::SuspectedOverflow { prompt_tokens } => {
entry.record_overflow(prompt_tokens, today)
}
newt_core::RoundObservation::ThinkingOnly => entry.record_thinking_only(),
}
}
pub type CapabilityCache = HashMap<String, CapabilityEntry>;
#[derive(Debug, Clone)]
pub struct ModelInfo {
pub name: String,
pub param_size: String,
}
#[cfg(test)]
thread_local! {
static CACHE_DIR_OVERRIDE: std::cell::RefCell<Option<PathBuf>> =
const { std::cell::RefCell::new(None) };
}
#[cfg(test)]
pub(crate) fn set_cache_dir_override(dir: Option<PathBuf>) {
CACHE_DIR_OVERRIDE.with(|c| *c.borrow_mut() = dir);
}
fn cache_path() -> Option<PathBuf> {
#[cfg(test)]
if let Some(dir) = CACHE_DIR_OVERRIDE.with(|c| c.borrow().clone()) {
return Some(dir.join("model-capabilities.json"));
}
newt_core::Config::user_config_path().map(|p| p.with_file_name("model-capabilities.json"))
}
pub fn load_cache() -> CapabilityCache {
let Some(path) = cache_path() else {
return Default::default();
};
let Ok(data) = std::fs::read_to_string(&path) else {
return Default::default();
};
let mut cache: CapabilityCache = serde_json::from_str(&data).unwrap_or_default();
if migrate_accounting(&mut cache) {
save_cache(&cache);
}
cache
}
pub fn migrate_accounting(cache: &mut CapabilityCache) -> bool {
let mut dirty = false;
for (model, entry) in cache.iter_mut() {
if entry.accounting_version >= ACCOUNTING_VERSION {
continue;
}
if entry.max_ok_input.is_some() {
tracing::info!(
model,
max_ok_input = entry.max_ok_input,
"invalidating max_ok_input recorded under the double-counting \
regime (Step 18.1); the ratchet will re-learn"
);
entry.max_ok_input = None;
entry.consecutive_ok = 0;
entry.tune_confidence = TuneConfidence::None;
}
entry.accounting_version = ACCOUNTING_VERSION;
dirty = true;
}
dirty
}
pub fn save_cache(cache: &CapabilityCache) {
let Some(path) = cache_path() else { return };
if let Ok(data) = serde_json::to_string_pretty(cache) {
let _ = std::fs::write(path, data);
}
}
pub fn resolve_memory_budget(explicit: Option<u32>, cache: &CapabilityCache, model: &str) -> u32 {
explicit
.or_else(|| {
cache
.get(model)
.and_then(|e| match (e.max_ok_input, e.safe_context) {
(Some(m), Some(s)) => Some(m.max(s)),
(m, s) => m.or(s),
})
})
.unwrap_or(newt_core::DEFAULT_CONTEXT_TOKENS)
}
pub fn fetch_ollama_models(endpoint: &str) -> anyhow::Result<Vec<ModelInfo>> {
let url = format!("{}/api/tags", endpoint.trim_end_matches('/'));
let json: serde_json::Value = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
let resp = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(5))
.build()?
.get(&url)
.send()
.await?;
if !resp.status().is_success() {
anyhow::bail!("HTTP {}", resp.status());
}
resp.json::<serde_json::Value>()
.await
.map_err(anyhow::Error::from)
})
})?;
Ok(json["models"]
.as_array()
.map(|arr| {
arr.iter()
.filter_map(|m| {
let name = m["name"].as_str()?.to_string();
let param_size = m["details"]["parameter_size"]
.as_str()
.unwrap_or("")
.to_string();
Some(ModelInfo { name, param_size })
})
.collect()
})
.unwrap_or_default())
}
pub fn fetch_context_window(endpoint: &str, model: &str) -> Option<u32> {
let url = format!("{}/api/show", endpoint.trim_end_matches('/'));
let body = serde_json::json!({"name": model});
let json: serde_json::Value = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
.ok()?
.post(&url)
.json(&body)
.send()
.await
.ok()?
.json::<serde_json::Value>()
.await
.ok()
})
})?;
parse_show_response(&json)
}
pub(crate) fn parse_show_response(json: &serde_json::Value) -> Option<u32> {
let arch_limit: Option<u32> = json["model_info"].as_object().and_then(|info| {
if let Some(v) = info.get("context_length").and_then(|v| v.as_u64()) {
return Some(v as u32);
}
info.iter()
.filter(|(k, _)| k.ends_with(".context_length"))
.filter_map(|(_, v)| v.as_u64())
.map(|v| v as u32)
.min() });
let modelfile_ctx: Option<u32> = json["parameters"].as_str().and_then(|params| {
params.lines().find_map(|line| {
let mut parts = line.split_whitespace();
if parts.next()? == "num_ctx" {
parts.next()?.parse::<u32>().ok()
} else {
None
}
})
});
match (arch_limit, modelfile_ctx) {
(Some(a), Some(b)) => Some(a.min(b)),
(Some(a), None) => Some(a),
(None, Some(b)) => Some(b),
(None, None) => None,
}
}
pub fn ensure_context_window(
entry: &mut CapabilityEntry,
endpoint: &str,
model: &str,
trust_declared: bool,
) -> bool {
if !trust_declared && entry.context_window.is_some() {
return false;
}
let mut changed = false;
if entry.context_window.is_none() {
let Some(window) = fetch_context_window(endpoint, model) else {
return false;
};
entry.context_window = Some(window);
changed = true;
}
let Some(window) = entry.context_window else {
return changed;
};
let declared_safe = window * 80 / 100;
if trust_declared {
if entry.safe_context != Some(declared_safe) {
entry.safe_context = Some(declared_safe);
changed = true;
}
} else if entry.safe_context.is_none() {
entry.safe_context = Some(declared_safe);
changed = true;
}
changed
}
pub fn today_local_date() -> String {
chrono::Local::now().format("%Y-%m-%d").to_string()
}
pub fn refresh_context_window(
entry: &mut CapabilityEntry,
endpoint: &str,
model: &str,
trust_declared: bool,
) -> bool {
let Some(window) = fetch_context_window(endpoint, model) else {
return false;
};
let mut changed = entry.context_window != Some(window);
entry.context_window = Some(window);
let declared_safe = window * 80 / 100;
if trust_declared {
if entry.safe_context != Some(declared_safe) {
entry.safe_context = Some(declared_safe);
changed = true;
}
} else if entry.safe_context.is_none() {
entry.safe_context = Some(declared_safe);
}
changed
}
pub fn message_thinking_fields(message: &serde_json::Value) -> bool {
let content_empty = message["content"]
.as_str()
.map(|c| c.trim().is_empty())
.unwrap_or(true);
if !content_empty {
return false;
}
["thinking", "reasoning", "reasoning_content"]
.iter()
.any(|field| {
message[*field]
.as_str()
.map(|v| !v.trim().is_empty())
.unwrap_or(false)
})
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProbeThinking {
pub emits_thinking: bool,
pub calibration: Option<(u32, usize)>,
}
pub fn probe_thinking(
endpoint: &str,
model: &str,
est: TokenEstimation,
) -> anyhow::Result<ProbeThinking> {
let url = format!("{}/api/chat", endpoint.trim_end_matches('/'));
let prompt = "Reply with the single word: ok";
let body = serde_json::json!({
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": false,
});
let estimated = serde_json::to_string(&body)
.map(|s| est.tokens_for_chars(s.chars().count()))
.unwrap_or(est.tokens_for_chars(prompt.chars().count()));
let json: serde_json::Value = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
let resp = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()?
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
if !resp.status().is_success() {
anyhow::bail!("Ollama returned {}", resp.status());
}
resp.json::<serde_json::Value>()
.await
.map_err(anyhow::Error::from)
})
})?;
let emits_thinking = message_thinking_fields(&json["message"]);
let calibration = json["prompt_eval_count"]
.as_u64()
.map(|observed| (observed as u32, estimated));
Ok(ProbeThinking {
emits_thinking,
calibration,
})
}
fn sanitize_ratio(ratio: f32) -> f32 {
if ratio.is_finite() && (0.5..=3.0).contains(&ratio) {
ratio
} else {
1.0
}
}
pub fn build_padded_prompt(target_real_tokens: u32, ratio: f32) -> String {
let ratio = sanitize_ratio(ratio);
let target_estimate = (target_real_tokens as f32 / ratio).round().max(1.0);
let target_chars = (target_estimate * 4.0).round() as usize;
const CLAUSE: &str = "lorem ipsum dolor sit amet ";
let mut s = String::with_capacity(target_chars + CLAUSE.len());
while s.chars().count() < target_chars {
s.push_str(CLAUSE);
}
s
}
#[derive(Debug, Clone, PartialEq)]
pub enum BoundaryClass {
Accepted { prompt_tokens: u32 },
Truncated,
CtxWindow400 { limit: u32 },
Inconclusive,
}
pub fn classify_boundary_probe(
http: Result<&serde_json::Value, &anyhow::Error>,
sent_real_estimate: u32,
) -> BoundaryClass {
match http {
Ok(json) => {
let message = &json["message"];
let content_nonempty = message["content"]
.as_str()
.map(|c| !c.trim().is_empty())
.unwrap_or(false);
let has_tool_call = message["tool_calls"]
.as_array()
.map(|a| !a.is_empty())
.unwrap_or(false);
let eval_count = json["eval_count"].as_u64().unwrap_or(0);
let usable = content_nonempty || has_tool_call || eval_count > 0;
let prompt_eval = json["prompt_eval_count"].as_u64().unwrap_or(0) as u32;
let threshold = (sent_real_estimate as u64 * 90 / 100) as u32;
if usable && prompt_eval >= threshold {
BoundaryClass::Accepted {
prompt_tokens: prompt_eval,
}
} else {
BoundaryClass::Truncated
}
}
Err(e) => match parse_context_window_error(&e.to_string()) {
Some((_, limit)) => BoundaryClass::CtxWindow400 {
limit: limit as u32,
},
None => BoundaryClass::Inconclusive,
},
}
}
fn boundary_probe_request(
endpoint: &str,
model: &str,
prompt: &str,
num_ctx: u32,
est: TokenEstimation,
) -> (anyhow::Result<serde_json::Value>, usize) {
let url = format!("{}/api/chat", endpoint.trim_end_matches('/'));
let body = serde_json::json!({
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": false,
"options": {"num_ctx": num_ctx, "num_predict": 8},
});
let sent_chars4 = serde_json::to_string(&body)
.map(|s| est.tokens_for_chars(s.chars().count()))
.unwrap_or(est.tokens_for_chars(prompt.chars().count()));
let result = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
let resp = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()?
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
if !status.is_success() {
anyhow::bail!("inference endpoint {status}: {text}");
}
serde_json::from_str::<serde_json::Value>(&text)
.map_err(|e| anyhow::anyhow!("bad JSON from /api/chat: {e}"))
})
});
(result, sent_chars4)
}
#[derive(Debug, Clone, PartialEq)]
pub struct BoundarySearchOutcome {
pub highest_accepted: Option<u32>,
pub steps: u32,
pub final_bounds: (u32, u32),
pub error: Option<String>,
}
pub fn probe_input_boundary(
endpoint: &str,
model: &str,
entry: &mut CapabilityEntry,
mut progress: impl FnMut(&str),
today: &str,
est: TokenEstimation,
) -> anyhow::Result<BoundarySearchOutcome> {
const STEP_CAP: u32 = 12;
const REPLY_MARGIN: u32 = 256;
const LOW_FLOOR: u32 = 2_048;
const DOUBLE_CAP: u32 = 1_000_000;
let ratio = entry.estimate_ratio.map(sanitize_ratio).unwrap_or(1.0);
let mut low = entry.safe_context.unwrap_or(LOW_FLOOR).max(LOW_FLOOR);
let declared = entry.context_window;
let mut highest_accepted: Option<u32> = None;
let mut steps = 0u32;
let mut error: Option<String> = None;
let run_probe =
|n: u32, entry: &mut CapabilityEntry, progress: &mut dyn FnMut(&str)| -> BoundaryClass {
let num_ctx = match declared {
Some(w) => (n + REPLY_MARGIN).min(w),
None => n + REPLY_MARGIN,
};
let prompt = build_padded_prompt(n, ratio);
let (http, sent_chars4) =
boundary_probe_request(endpoint, model, &prompt, num_ctx, est);
let class = classify_boundary_probe(http.as_ref(), n);
match &class {
BoundaryClass::Accepted { prompt_tokens } => {
entry.record_accepted_prompt(*prompt_tokens, today);
entry.record_estimate_sample(*prompt_tokens, sent_chars4);
progress(&format!(
" num_ctx={num_ctx}: accepted (prompt_eval={prompt_tokens})"
));
}
BoundaryClass::Truncated => {
progress(&format!(" num_ctx={num_ctx}: truncated/rejected"));
}
BoundaryClass::CtxWindow400 { limit } => {
progress(&format!(
" num_ctx={num_ctx}: context-window 400 (limit {limit})"
));
}
BoundaryClass::Inconclusive => {
progress(&format!(
" num_ctx={num_ctx}: inconclusive — {}",
http.err()
.map(|e| e.to_string())
.unwrap_or_else(|| "transport error".into())
));
}
}
class
};
let mut high = match declared {
Some(w) => w.max(low + 1),
None => {
let mut candidate = low.saturating_mul(2).max(low + 1024).min(DOUBLE_CAP);
let mut bracket_high = DOUBLE_CAP;
loop {
if steps >= STEP_CAP {
break;
}
steps += 1;
match run_probe(candidate, entry, &mut progress) {
BoundaryClass::Accepted { prompt_tokens } => {
highest_accepted =
Some(highest_accepted.map_or(prompt_tokens, |h| h.max(prompt_tokens)));
low = low.max(candidate);
if candidate >= DOUBLE_CAP {
bracket_high = DOUBLE_CAP;
break;
}
candidate = candidate.saturating_mul(2).min(DOUBLE_CAP);
}
BoundaryClass::CtxWindow400 { limit } => {
entry.record_context_window_400(limit, today);
bracket_high = limit;
break;
}
BoundaryClass::Truncated => {
bracket_high = candidate;
break;
}
BoundaryClass::Inconclusive => {
error = Some("boundary search stopped: inconclusive probe".into());
bracket_high = candidate;
break;
}
}
}
bracket_high
}
};
while error.is_none() && steps < STEP_CAP {
let tolerance = 1_024u32.max((high as u64 * 5 / 100) as u32);
if high.saturating_sub(low) <= tolerance {
break;
}
let mid = low + (high - low) / 2;
steps += 1;
match run_probe(mid, entry, &mut progress) {
BoundaryClass::Accepted { prompt_tokens } => {
highest_accepted =
Some(highest_accepted.map_or(prompt_tokens, |h| h.max(prompt_tokens)));
low = mid;
}
BoundaryClass::Truncated => {
high = mid;
}
BoundaryClass::CtxWindow400 { limit } => {
entry.record_context_window_400(limit, today);
high = limit.min(high);
if limit < low {
low = limit;
}
}
BoundaryClass::Inconclusive => {
error = Some("boundary search stopped: inconclusive probe".into());
break;
}
}
}
if highest_accepted.is_some() {
entry.tune_confidence = TuneConfidence::High;
entry.tune_date = Some(today.to_string());
}
Ok(BoundarySearchOutcome {
highest_accepted,
steps,
final_bounds: (low, high),
error,
})
}
pub fn is_tuning_stale(tune_date: Option<&str>, today: &str, max_age_days: i64) -> bool {
let Some(stamp) = tune_date else {
return true;
};
let (Ok(then), Ok(now)) = (
chrono::NaiveDate::parse_from_str(stamp, "%Y-%m-%d"),
chrono::NaiveDate::parse_from_str(today, "%Y-%m-%d"),
) else {
return true;
};
(now - then).num_days() > max_age_days
}
pub fn tuning_age_days(tune_date: Option<&str>, today: &str) -> Option<i64> {
let then = chrono::NaiveDate::parse_from_str(tune_date?, "%Y-%m-%d").ok()?;
let now = chrono::NaiveDate::parse_from_str(today, "%Y-%m-%d").ok()?;
Some((now - then).num_days())
}
#[derive(Debug, Clone, PartialEq)]
pub struct FullProbeReport {
pub conformance: ToolConformance,
pub context_window: Option<u32>,
pub emits_thinking: bool,
pub estimate_ratio: Option<f32>,
pub boundary: Option<BoundarySearchOutcome>,
pub notes: Vec<String>,
}
pub fn full_probe(
endpoint: &str,
model: &str,
entry: &mut CapabilityEntry,
do_window: bool,
today: &str,
mut progress: impl FnMut(&str),
est: TokenEstimation,
) -> FullProbeReport {
let mut notes = Vec::new();
let conformance = match tokio::task::block_in_place(|| {
tokio::runtime::Handle::current()
.block_on(probe_tool_conformance_calibrated(endpoint, model, est))
}) {
Ok(pc) => {
if let Some((observed, estimated)) = pc.calibration {
entry.record_estimate_sample(observed, estimated);
}
pc.conformance
}
Err(e) => {
notes.push(format!("conformance probe failed: {e}"));
entry.conformance.clone()
}
};
entry.conformance = conformance.clone();
entry.tested_date = today.to_string();
refresh_context_window(entry, endpoint, model, false);
let mut emits_thinking = entry.emits_thinking.unwrap_or(false);
match probe_thinking(endpoint, model, est) {
Ok(pt) => {
if pt.emits_thinking {
entry.record_thinking_only();
emits_thinking = true;
}
if let Some((observed, estimated)) = pt.calibration {
entry.record_estimate_sample(observed, estimated);
}
}
Err(e) => notes.push(format!("thinking probe failed: {e}")),
}
let boundary = if do_window {
match probe_input_boundary(endpoint, model, entry, &mut progress, today, est) {
Ok(outcome) => Some(outcome),
Err(e) => {
notes.push(format!("boundary search failed: {e}"));
None
}
}
} else {
None
};
FullProbeReport {
conformance,
context_window: entry.context_window,
emits_thinking,
estimate_ratio: entry.estimate_ratio,
boundary,
notes,
}
}
fn probe_tool_schema() -> serde_json::Value {
serde_json::json!([{
"type": "function",
"function": {
"name": "list_dir",
"description": "List files in a directory",
"parameters": {
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Directory path (use '.' for current directory)"
}
},
"required": ["path"]
}
}
}])
}
pub fn parse_context_window_error(msg: &str) -> Option<(u64, u64)> {
let after = msg.split("prompt is too long:").nth(1)?;
let prompt = first_number(after)?;
let after_gt = after.split('>').nth(1)?;
let max = first_number(after_gt)?;
Some((prompt, max))
}
fn first_number(s: &str) -> Option<u64> {
s.split(|c: char| !c.is_ascii_digit())
.find(|t| !t.is_empty())
.and_then(|t| t.parse().ok())
}
pub fn looks_like_tool_call_json(content: &str) -> bool {
let trimmed = content.trim();
if !trimmed.contains("\"name\"") || !trimmed.contains("\"arguments\"") {
return false;
}
if let Ok(val) = serde_json::from_str::<serde_json::Value>(trimmed) {
let is_call =
|v: &serde_json::Value| v.get("name").is_some() && v.get("arguments").is_some();
if is_call(&val) {
return true;
}
if val
.as_array()
.map(|a| a.iter().any(is_call))
.unwrap_or(false)
{
return true;
}
}
false
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProbeConformance {
pub conformance: ToolConformance,
pub calibration: Option<(u32, usize)>,
}
fn classify_conformance(message: &serde_json::Value) -> ToolConformance {
if let Some(tcs) = message["tool_calls"].as_array() {
if !tcs.is_empty() {
return ToolConformance::Native;
}
}
let content = message["content"].as_str().unwrap_or("");
if looks_like_tool_call_json(content) {
return ToolConformance::TextMode;
}
ToolConformance::NoTools
}
pub async fn probe_tool_conformance_calibrated(
endpoint: &str,
model: &str,
est: TokenEstimation,
) -> anyhow::Result<ProbeConformance> {
let url = format!("{}/api/chat", endpoint.trim_end_matches('/'));
let body = serde_json::json!({
"model": model,
"messages": [{
"role": "user",
"content": "Call the list_dir tool on path '.'. \
Do not explain — just call the tool."
}],
"tools": probe_tool_schema(),
"stream": false,
});
let estimated = serde_json::to_string(&body)
.map(|s| est.tokens_for_chars(s.chars().count()))
.unwrap_or(0);
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()?;
let resp = client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| anyhow::anyhow!("request failed: {e}"))?;
if !resp.status().is_success() {
anyhow::bail!("Ollama returned {}", resp.status());
}
let json: serde_json::Value = resp.json().await?;
let conformance = classify_conformance(&json["message"]);
let calibration = json["prompt_eval_count"]
.as_u64()
.map(|observed| (observed as u32, estimated));
Ok(ProbeConformance {
conformance,
calibration,
})
}
pub async fn probe_tool_conformance(
endpoint: &str,
model: &str,
est: TokenEstimation,
) -> anyhow::Result<ToolConformance> {
probe_tool_conformance_calibrated(endpoint, model, est)
.await
.map(|p| p.conformance)
}
pub fn print_capabilities_table(
models: &[ModelInfo],
cache: &CapabilityCache,
active: &str,
endpoint: &str,
color: bool,
) {
let tested = models
.iter()
.filter(|m| cache.contains_key(&m.name))
.count();
println!(
"Models on {} ({} total, {} tested)\n",
endpoint,
models.len(),
tested,
);
let name_w = models
.iter()
.map(|m| m.name.len())
.max()
.unwrap_or(20)
.max(20);
let sep = "─".repeat(name_w);
println!(
" {:<name_w$} {:>6} {:<8} {:<5} {:>8} {:>8} Conf Tested",
"Model", "Size", "Tool Use", "Think", "Ctx Win", "Safe Ctx"
);
println!(" {sep} ────── ──────── ───── ──────── ──────── ──── ──────────");
for m in models {
let is_active = m.name == active;
let active_tag = if is_active { " ◀" } else { " " };
let size = if m.param_size.is_empty() {
" — ".to_string()
} else {
format!("{:>6}", m.param_size)
};
let (conformance_str, think_str, ctx_win_str, safe_ctx_str, conf_str, date_str) =
match cache.get(&m.name) {
Some(e) => {
let ctx = e
.context_window
.map(|c| format!("{:>8}", fmt_k(c)))
.unwrap_or_else(|| " —".to_string());
let safe = e
.safe_context
.map(|c| format!("{:>8}", fmt_k(c)))
.unwrap_or_else(|| " —".to_string());
let conf = match e.tune_confidence {
TuneConfidence::None => " — ".to_string(),
TuneConfidence::Low => " Low".to_string(),
TuneConfidence::Medium => " Med".to_string(),
TuneConfidence::High => "High".to_string(),
};
let think = if e.emits_thinking == Some(true) {
"✓".to_string()
} else {
"—".to_string()
};
(
e.conformance.symbol().to_string(),
think,
ctx,
safe,
conf,
e.tested_date.clone(),
)
}
None => (
"— ".to_string(),
"—".to_string(),
" —".to_string(),
" —".to_string(),
" — ".to_string(),
"(untested)".to_string(),
),
};
let name = &m.name;
let row = format!(
" {name:<name_w$}{active_tag} {size} {conformance_str} {think_str:<5} {ctx_win_str} {safe_ctx_str} {conf_str} {date_str}"
);
if color && is_active {
use crossterm::style::Color as CtColor;
use crossterm::{
execute,
style::{Print, ResetColor, SetForegroundColor},
};
execute!(
std::io::stdout(),
SetForegroundColor(CtColor::Rgb {
r: 220,
g: 60,
b: 20
}),
Print(format!("{row}\n")),
ResetColor,
)
.ok();
} else {
println!("{row}");
}
}
println!();
println!("Legend:");
println!(" ✓ native tool_calls field — works with this harness");
println!(" ~ text JSON embedded in content — NOT dispatched by newt");
println!(" ✗ none ignores tools, answers directly");
println!(" — untested → /probe <model> to classify");
println!();
println!(" Think ✓ reasoning model — emits chain-of-thought tokens (auto-detected;");
println!(" newt streams it dimmed and keeps it out of the saved answer)");
println!(" Ctx Win declared context window from Ollama /api/show");
println!(" Safe Ctx num_ctx sent to Ollama (auto-tuned; human-overridable in config)");
println!(" Conf tuning confidence: None | Low | Med | High");
println!();
println!("Run /probe <model> to test a model (warm-up included).");
println!("Run /probe all to test every untested model in sequence.");
println!("Run /probe window <model> for an empirical input-boundary search (High confidence).");
}
fn fmt_k(n: u32) -> String {
if n >= 1024 {
format!("{}k", n / 1024)
} else {
n.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn looks_like_tool_call_json_native_object() {
assert!(looks_like_tool_call_json(
r#"{"name":"list_dir","arguments":{"path":"."}}"#
));
}
#[test]
fn looks_like_tool_call_json_array() {
assert!(looks_like_tool_call_json(
r#"[{"name":"list_dir","arguments":{"path":"."}}]"#
));
}
#[test]
fn looks_like_tool_call_json_plain_text() {
assert!(!looks_like_tool_call_json(
"Here are the files: README.md, src/"
));
}
#[test]
fn looks_like_tool_call_json_incomplete_object() {
assert!(!looks_like_tool_call_json(r#"{"name":"list_dir"}"#));
}
#[test]
fn load_cache_returns_empty_on_missing_file() {
let _ = load_cache();
}
#[test]
fn conformance_symbol_coverage() {
assert!(ToolConformance::Native.symbol().contains('✓'));
assert!(ToolConformance::TextMode.symbol().contains('~'));
assert!(ToolConformance::NoTools.symbol().contains('✗'));
}
#[test]
fn tune_confidence_promotes_correctly() {
assert_eq!(TuneConfidence::None.promote(), TuneConfidence::Low);
assert_eq!(TuneConfidence::Low.promote(), TuneConfidence::Medium);
assert_eq!(TuneConfidence::Medium.promote(), TuneConfidence::High);
assert_eq!(TuneConfidence::High.promote(), TuneConfidence::High);
}
fn make_entry() -> CapabilityEntry {
CapabilityEntry {
conformance: ToolConformance::Native,
tested_date: "2026-06-06".to_string(),
context_window: Some(32768),
safe_context: Some(26214),
..Default::default()
}
}
#[test]
fn record_success_updates_max_ok_input() {
let mut e = make_entry();
e.record_success(10_000, "2026-06-06");
assert_eq!(e.max_ok_input, Some(10_000));
e.record_success(8_000, "2026-06-06");
assert_eq!(e.max_ok_input, Some(10_000));
}
#[test]
fn record_success_promotes_confidence_after_five() {
let mut e = make_entry();
for i in 0..4 {
e.record_success(5_000, "2026-06-06");
assert_eq!(e.tune_confidence, TuneConfidence::None, "early iter {i}");
}
e.record_success(5_000, "2026-06-06");
assert_eq!(e.tune_confidence, TuneConfidence::Low);
assert_eq!(e.consecutive_ok, 0); }
#[test]
fn record_overflow_reduces_safe_context() {
let mut e = make_entry();
e.record_overflow(30_000, "2026-06-06");
assert_eq!(e.safe_context, Some(22_500));
assert_eq!(e.tune_confidence, TuneConfidence::Low);
assert_eq!(e.overflow_at, Some(30_000));
}
#[test]
fn record_overflow_reins_max_ok_input_down() {
let mut e = make_entry();
e.max_ok_input = Some(28_000);
e.record_overflow(20_000, "2026-06-12");
assert_eq!(e.safe_context, Some(15_000));
assert_eq!(e.max_ok_input, Some(15_000));
let mut e2 = make_entry();
e2.max_ok_input = Some(10_000);
e2.record_overflow(20_000, "2026-06-12");
assert_eq!(e2.max_ok_input, Some(10_000));
let mut e3 = make_entry();
e3.record_overflow(20_000, "2026-06-12");
assert_eq!(e3.max_ok_input, None);
}
#[test]
fn record_accepted_prompt_is_a_pure_high_water_ratchet() {
let mut e = make_entry();
e.consecutive_ok = 3;
e.tune_confidence = TuneConfidence::Medium;
assert!(
e.record_accepted_prompt(8_734, "2026-06-12"),
"first: dirty"
);
assert_eq!(e.max_ok_input, Some(8_734));
assert_eq!(e.tune_date.as_deref(), Some("2026-06-12"), "date stamped");
assert_eq!(e.consecutive_ok, 3, "untouched");
assert_eq!(e.tune_confidence, TuneConfidence::Medium, "untouched");
assert!(
!e.record_accepted_prompt(8_734, "2026-06-13"),
"equal: clean"
);
assert!(
!e.record_accepted_prompt(4_000, "2026-06-13"),
"lower: clean"
);
assert_eq!(e.max_ok_input, Some(8_734), "HWM only raises");
assert_eq!(e.tune_date.as_deref(), Some("2026-06-12"), "no re-stamp");
assert!(e.record_accepted_prompt(9_000, "2026-06-13"));
assert_eq!(e.max_ok_input, Some(9_000));
assert_eq!(e.tune_date.as_deref(), Some("2026-06-13"));
}
#[test]
fn record_estimate_sample_initializes_then_emas() {
let mut e = make_entry();
assert!(e.record_estimate_sample(8_734, 6_600));
let first = e.estimate_ratio.unwrap();
assert!((first - 8_734.0 / 6_600.0).abs() < 1e-6, "got {first}");
assert!(e.record_estimate_sample(2_000, 1_000));
let second = e.estimate_ratio.unwrap();
assert!(
(second - (0.75 * first + 0.25 * 2.0)).abs() < 1e-6,
"got {second}"
);
}
#[test]
fn record_estimate_sample_clamps_both_ends() {
let mut e = make_entry();
assert!(e.record_estimate_sample(10_000, 1_000)); assert_eq!(e.estimate_ratio, Some(3.0), "init clamped to 3.0");
let mut e2 = make_entry();
assert!(e2.record_estimate_sample(500, 1_000));
assert_eq!(e2.estimate_ratio, Some(0.5));
assert!(!e.record_estimate_sample(10_000, 1_000), "3.0 → 3.0: clean");
assert_eq!(e.estimate_ratio, Some(3.0));
}
#[test]
fn record_estimate_sample_skips_cache_hits_and_zero_estimates() {
let mut e = make_entry();
assert!(!e.record_estimate_sample(400, 1_000));
assert_eq!(e.estimate_ratio, None);
assert!(!e.record_estimate_sample(400, 0));
assert_eq!(e.estimate_ratio, None);
e.estimate_ratio = Some(1.3);
assert!(!e.record_estimate_sample(100, 1_000));
assert_eq!(e.estimate_ratio, Some(1.3));
}
#[test]
fn record_estimate_sample_dirty_only_above_threshold() {
let mut e = make_entry();
assert!(e.record_estimate_sample(1_300, 1_000)); let stored = e.estimate_ratio.unwrap();
assert!(!e.record_estimate_sample(1_301, 1_000));
let drifted = e.estimate_ratio.unwrap();
assert!((drifted - stored).abs() < 0.01, "stored as-is, tiny drift");
assert!(e.record_estimate_sample(2_000, 1_000));
}
#[test]
fn record_thinking_only_is_sticky_and_dirty_once() {
let mut e = make_entry();
assert_eq!(e.emits_thinking, None);
assert!(e.record_thinking_only(), "first observation: dirty");
assert_eq!(e.emits_thinking, Some(true));
assert!(!e.record_thinking_only(), "repeat: clean");
assert_eq!(e.emits_thinking, Some(true));
}
#[test]
fn apply_observation_dispatches_each_variant() {
let today = "2026-06-12";
let mut e = make_entry();
let obs = newt_core::RoundObservation::Accepted {
prompt_tokens: 8_734,
estimated_tokens: 6_600,
};
assert!(apply_observation(&mut e, &obs, today));
assert_eq!(e.max_ok_input, Some(8_734));
assert!(e.estimate_ratio.is_some());
assert!(!apply_observation(&mut e, &obs, today));
let recal = newt_core::RoundObservation::Accepted {
prompt_tokens: 8_000,
estimated_tokens: 3_000,
};
assert!(apply_observation(&mut e, &recal, today));
assert_eq!(e.max_ok_input, Some(8_734), "lower prompt: no ratchet");
let mut e = make_entry();
e.max_ok_input = Some(28_000);
let obs = newt_core::RoundObservation::SuspectedOverflow {
prompt_tokens: 20_000,
};
assert!(apply_observation(&mut e, &obs, today));
assert_eq!(e.safe_context, Some(15_000));
assert_eq!(e.max_ok_input, Some(15_000));
assert_eq!(e.overflow_at, Some(20_000));
let mut e = make_entry();
assert!(apply_observation(
&mut e,
&newt_core::RoundObservation::ThinkingOnly,
today
));
assert_eq!(e.emits_thinking, Some(true));
assert!(!apply_observation(
&mut e,
&newt_core::RoundObservation::ThinkingOnly,
today
));
}
#[test]
fn estimate_ratio_and_emits_thinking_roundtrip_json() {
let mut e = make_entry();
e.estimate_ratio = Some(1.29);
e.emits_thinking = Some(true);
let json = serde_json::to_string(&e).unwrap();
let back: CapabilityEntry = serde_json::from_str(&json).unwrap();
assert_eq!(back.estimate_ratio, Some(1.29));
assert_eq!(back.emits_thinking, Some(true));
let bare = serde_json::to_string(&make_entry()).unwrap();
assert!(!bare.contains("estimate_ratio"), "{bare}");
assert!(!bare.contains("emits_thinking"), "{bare}");
}
#[test]
fn record_overflow_does_not_increase_safe_context() {
let mut e = make_entry();
e.safe_context = Some(10_000);
e.record_overflow(5_000, "2026-06-06");
assert_eq!(e.safe_context, Some(3_750));
e.record_overflow(40_000, "2026-06-06");
assert_eq!(e.safe_context, Some(3_750));
}
#[test]
fn parse_context_window_error_none_for_unrelated_400() {
let msg = "inference endpoint 400: invalid api key";
assert_eq!(super::parse_context_window_error(msg), None);
}
#[test]
fn parse_context_window_error_extracts_prompt_and_max() {
let msg = "inference endpoint 400: litellm.ContextWindowExceededError: prompt is too long: 5960028 tokens > 1000000 maximum";
assert_eq!(
super::parse_context_window_error(msg),
Some((5_960_028, 1_000_000))
);
}
#[test]
fn parse_context_window_error_none_without_max_clause() {
let msg = "prompt is too long: 5960028 tokens";
assert_eq!(super::parse_context_window_error(msg), None);
}
#[test]
fn record_context_window_400_tightens_max_ok_input_to_80pct() {
let mut e = make_entry();
e.max_ok_input = Some(251_640);
let dirty = e.record_context_window_400(1_000_000, "2026-06-08");
assert!(dirty);
assert_eq!(e.max_ok_input, Some(800_000));
assert_eq!(e.tune_confidence, TuneConfidence::Low);
assert_eq!(e.consecutive_ok, 0);
}
#[test]
fn record_context_window_400_lowers_an_overshot_cap() {
let mut e = make_entry();
e.max_ok_input = Some(2_000_000);
e.record_context_window_400(1_000_000, "2026-06-08");
assert_eq!(e.max_ok_input, Some(800_000));
}
#[test]
fn record_context_window_400_caps_safe_context_without_raising_it() {
let mut e = make_entry();
e.safe_context = Some(64_000); e.record_context_window_400(1_000_000, "2026-06-08");
assert_eq!(e.safe_context, Some(64_000));
}
#[test]
fn fmt_k_formats_correctly() {
assert_eq!(fmt_k(1024), "1k");
assert_eq!(fmt_k(32768), "32k");
assert_eq!(fmt_k(131072), "128k");
assert_eq!(fmt_k(512), "512");
}
#[test]
fn capability_entry_roundtrips_json_with_new_fields() {
let mut e = make_entry();
e.overflow_at = Some(28_000);
e.max_ok_input = Some(25_000);
e.tune_confidence = TuneConfidence::Medium;
e.tune_date = Some("2026-06-06".to_string());
let json = serde_json::to_string(&e).unwrap();
let back: CapabilityEntry = serde_json::from_str(&json).unwrap();
assert_eq!(back.context_window, Some(32768));
assert_eq!(back.overflow_at, Some(28_000));
assert_eq!(back.tune_confidence, TuneConfidence::Medium);
}
#[test]
fn capability_entry_deserializes_legacy_json_without_new_fields() {
let legacy = r#"{"conformance":"native","tested_date":"2026-06-04"}"#;
let e: CapabilityEntry = serde_json::from_str(legacy).unwrap();
assert_eq!(e.conformance, ToolConformance::Native);
assert_eq!(e.context_window, None);
assert_eq!(e.tune_confidence, TuneConfidence::None);
assert_eq!(e.accounting_version, 0);
}
#[test]
fn migrate_accounting_invalidates_poisoned_entry() {
let mut cache = CapabilityCache::default();
cache.insert(
"llama3.1:8b".into(),
CapabilityEntry {
conformance: ToolConformance::Native,
tested_date: "2026-06-08".into(),
context_window: Some(8_192),
safe_context: Some(6_553),
overflow_at: None,
max_ok_input: Some(25_602),
consecutive_ok: 3,
tune_confidence: TuneConfidence::High,
tune_date: Some("2026-06-08".into()),
estimate_ratio: None,
emits_thinking: None,
accounting_version: 0, },
);
assert!(
migrate_accounting(&mut cache),
"migration must report dirty"
);
let e = &cache["llama3.1:8b"];
assert_eq!(e.max_ok_input, None, "poisoned ratchet value dropped");
assert_eq!(e.consecutive_ok, 0);
assert_eq!(e.tune_confidence, TuneConfidence::None);
assert_eq!(e.accounting_version, ACCOUNTING_VERSION);
assert_eq!(e.context_window, Some(8_192));
assert_eq!(e.safe_context, Some(6_553));
assert_eq!(e.conformance, ToolConformance::Native);
}
#[test]
fn migrate_accounting_leaves_current_version_entry_untouched() {
let mut cache = CapabilityCache::default();
let entry = CapabilityEntry {
conformance: ToolConformance::Native,
tested_date: "2026-06-09".into(),
safe_context: Some(64_000),
max_ok_input: Some(800_000), consecutive_ok: 2,
tune_confidence: TuneConfidence::Medium,
..Default::default() };
cache.insert("hosted-model".into(), entry.clone());
assert!(!migrate_accounting(&mut cache), "nothing to migrate");
let e = &cache["hosted-model"];
assert_eq!(e.max_ok_input, Some(800_000));
assert_eq!(e.consecutive_ok, 2);
assert_eq!(e.tune_confidence, TuneConfidence::Medium);
}
#[test]
fn migrate_accounting_stamps_untuned_legacy_entry() {
let mut cache = CapabilityCache::default();
cache.insert(
"old-model".into(),
CapabilityEntry {
conformance: ToolConformance::TextMode,
tested_date: "2026-06-04".into(),
accounting_version: 0,
..Default::default()
},
);
assert!(migrate_accounting(&mut cache));
assert_eq!(cache["old-model"].accounting_version, ACCOUNTING_VERSION);
assert_eq!(cache["old-model"].conformance, ToolConformance::TextMode);
}
#[test]
fn migrate_accounting_is_idempotent() {
let mut cache = CapabilityCache::default();
let mut e = make_entry();
e.max_ok_input = Some(25_602);
e.accounting_version = 0;
cache.insert("m".into(), e);
assert!(migrate_accounting(&mut cache), "first pass migrates");
let snapshot = serde_json::to_string(&cache).unwrap();
assert!(!migrate_accounting(&mut cache), "second pass is a no-op");
assert_eq!(serde_json::to_string(&cache).unwrap(), snapshot);
}
fn fixture_cache(max_ok_input: Option<u32>, safe_context: Option<u32>) -> CapabilityCache {
let mut cache = CapabilityCache::default();
cache.insert(
"tuned-model".into(),
CapabilityEntry {
conformance: ToolConformance::Native,
tested_date: "2026-06-10".into(),
context_window: Some(32_768),
safe_context,
max_ok_input,
..Default::default()
},
);
cache
}
#[test]
fn resolve_memory_budget_explicit_config_wins() {
let cache = fixture_cache(Some(24_000), Some(26_214));
assert_eq!(
resolve_memory_budget(Some(16_000), &cache, "tuned-model"),
16_000
);
}
#[test]
fn resolve_memory_budget_capability_max_ok_input_second() {
let cache = fixture_cache(Some(24_000), Some(6_553));
assert_eq!(resolve_memory_budget(None, &cache, "tuned-model"), 24_000);
}
#[test]
fn resolve_memory_budget_max_keeps_safe_context_over_low_hwm() {
let cache = fixture_cache(Some(6_068), Some(26_214));
assert_eq!(resolve_memory_budget(None, &cache, "tuned-model"), 26_214);
}
#[test]
fn resolve_memory_budget_falls_back_to_safe_context() {
let cache = fixture_cache(None, Some(6_553));
assert_eq!(resolve_memory_budget(None, &cache, "tuned-model"), 6_553);
let cache = fixture_cache(Some(24_000), None);
assert_eq!(resolve_memory_budget(None, &cache, "tuned-model"), 24_000);
}
#[test]
fn resolve_memory_budget_static_default_last() {
let empty = CapabilityCache::default();
assert_eq!(
resolve_memory_budget(None, &empty, "fresh-model"),
newt_core::DEFAULT_CONTEXT_TOKENS
);
let untuned = fixture_cache(None, None);
assert_eq!(
resolve_memory_budget(None, &untuned, "tuned-model"),
newt_core::DEFAULT_CONTEXT_TOKENS
);
let cache = fixture_cache(Some(24_000), Some(26_214));
assert_eq!(
resolve_memory_budget(None, &cache, "different-model"),
newt_core::DEFAULT_CONTEXT_TOKENS
);
}
#[test]
fn resolve_memory_budget_never_ignores_probe_data() {
let cache = fixture_cache(Some(24_000), Some(26_214));
let budget = resolve_memory_budget(None, &cache, "tuned-model");
assert_ne!(
budget,
newt_core::DEFAULT_CONTEXT_TOKENS,
"capability data present — the static default must not win"
);
assert_eq!(budget, 26_214);
}
#[test]
fn message_thinking_fields_truth_table() {
assert!(message_thinking_fields(&serde_json::json!({
"content": "", "thinking": "reasoning here"
})));
assert!(message_thinking_fields(&serde_json::json!({
"content": " \n", "reasoning": "x"
})));
assert!(message_thinking_fields(&serde_json::json!({
"content": "", "reasoning_content": "y"
})));
assert!(message_thinking_fields(&serde_json::json!({
"thinking": "z"
})));
assert!(!message_thinking_fields(&serde_json::json!({
"content": "ok", "thinking": "z"
})));
assert!(!message_thinking_fields(&serde_json::json!({
"content": "", "thinking": " "
})));
assert!(!message_thinking_fields(
&serde_json::json!({"content": ""})
));
}
#[test]
fn build_padded_prompt_sizes_near_target_across_ratios() {
for &(target, ratio) in &[
(512u32, 1.0f32),
(2_048, 1.0),
(8_000, 1.3),
(4_096, 0.8),
(16_000, 2.5),
] {
let s = build_padded_prompt(target, ratio);
let est = s.chars().count() as f32 / 4.0;
let predicted_real = est * sanitize_ratio(ratio);
let rel = (predicted_real - target as f32).abs() / target as f32;
assert!(
rel <= 0.10,
"target={target} ratio={ratio}: chars/4*ratio={predicted_real} ({:.1}% off)",
rel * 100.0
);
}
}
#[test]
fn build_padded_prompt_sanitizes_bad_ratio_to_one() {
let s = build_padded_prompt(4_000, f32::NAN);
let est = s.chars().count() as f32 / 4.0;
assert!((est - 4_000.0).abs() / 4_000.0 <= 0.10);
let s2 = build_padded_prompt(4_000, 99.0);
let est2 = s2.chars().count() as f32 / 4.0;
assert!((est2 - 4_000.0).abs() / 4_000.0 <= 0.10);
}
#[test]
fn classify_boundary_probe_accepted_when_eval_meets_threshold() {
let json = serde_json::json!({
"message": {"content": "ok"},
"prompt_eval_count": 9_500,
"eval_count": 3
});
assert_eq!(
classify_boundary_probe(Ok(&json), 10_000),
BoundaryClass::Accepted {
prompt_tokens: 9_500
}
);
}
#[test]
fn classify_boundary_probe_accepted_via_tool_call_or_eval_count() {
let tc = serde_json::json!({
"message": {"content": "", "tool_calls": [{"function": {"name": "x"}}]},
"prompt_eval_count": 9_900
});
assert!(matches!(
classify_boundary_probe(Ok(&tc), 10_000),
BoundaryClass::Accepted { .. }
));
let ec = serde_json::json!({
"message": {"content": ""},
"prompt_eval_count": 9_900,
"eval_count": 5
});
assert!(matches!(
classify_boundary_probe(Ok(&ec), 10_000),
BoundaryClass::Accepted { .. }
));
}
#[test]
fn classify_boundary_probe_truncated_when_eval_below_threshold() {
let json = serde_json::json!({
"message": {"content": "ok"},
"prompt_eval_count": 4_000,
"eval_count": 2
});
assert_eq!(
classify_boundary_probe(Ok(&json), 10_000),
BoundaryClass::Truncated
);
let no_count = serde_json::json!({"message": {"content": "ok"}, "eval_count": 1});
assert_eq!(
classify_boundary_probe(Ok(&no_count), 10_000),
BoundaryClass::Truncated
);
}
#[test]
fn classify_boundary_probe_ctx_window_400() {
let err = anyhow::anyhow!(
"inference endpoint 400: litellm.ContextWindowExceededError: \
prompt is too long: 12000 tokens > 8192 maximum"
);
assert_eq!(
classify_boundary_probe(Err(&err), 12_000),
BoundaryClass::CtxWindow400 { limit: 8_192 }
);
}
#[test]
fn classify_boundary_probe_inconclusive_for_other_errors() {
let err = anyhow::anyhow!("request failed: connection reset");
assert_eq!(
classify_boundary_probe(Err(&err), 10_000),
BoundaryClass::Inconclusive
);
}
#[test]
fn is_tuning_stale_boundaries() {
assert!(is_tuning_stale(None, "2026-06-13", 30));
assert!(!is_tuning_stale(Some("2026-05-14"), "2026-06-13", 30));
assert!(is_tuning_stale(Some("2026-05-13"), "2026-06-13", 30));
assert!(!is_tuning_stale(Some("2026-06-13"), "2026-06-13", 30));
assert!(is_tuning_stale(Some("not-a-date"), "2026-06-13", 30));
assert!(is_tuning_stale(Some("2026-06-13"), "garbage", 30));
}
#[test]
fn sanitize_ratio_clamps_and_defaults() {
assert_eq!(sanitize_ratio(1.3), 1.3);
assert_eq!(sanitize_ratio(0.5), 0.5);
assert_eq!(sanitize_ratio(3.0), 3.0);
assert_eq!(sanitize_ratio(0.4), 1.0); assert_eq!(sanitize_ratio(3.1), 1.0); assert_eq!(sanitize_ratio(f32::NAN), 1.0);
assert_eq!(sanitize_ratio(f32::INFINITY), 1.0);
}
#[test]
fn parse_show_response_reads_llama_key() {
let json = serde_json::json!({"model_info": {"llama.context_length": 32768}});
assert_eq!(super::parse_show_response(&json), Some(32768));
}
#[test]
fn parse_show_response_reads_nemotron_key() {
let json = serde_json::json!({"model_info": {"nemotron_h_omni.context_length": 131072}});
assert_eq!(super::parse_show_response(&json), Some(131072));
}
#[test]
fn parse_show_response_bare_context_length_key() {
let json = serde_json::json!({"model_info": {"context_length": 8192}});
assert_eq!(super::parse_show_response(&json), Some(8192));
}
#[test]
fn parse_show_response_modelfile_num_ctx_wins_when_smaller() {
let json = serde_json::json!({
"model_info": {"llama.context_length": 131072},
"parameters": "num_ctx 32768\ntemperature 0.7"
});
assert_eq!(super::parse_show_response(&json), Some(32768));
}
#[test]
fn parse_show_response_arch_wins_when_num_ctx_larger() {
let json = serde_json::json!({
"model_info": {"llama.context_length": 4096},
"parameters": "num_ctx 32768"
});
assert_eq!(super::parse_show_response(&json), Some(4096));
}
#[test]
fn parse_show_response_returns_none_when_no_keys() {
let json = serde_json::json!({"model_info": {"general.architecture": "llama"}});
assert_eq!(super::parse_show_response(&json), None);
}
#[test]
fn parse_show_response_uses_minimum_when_multiple_arch_keys() {
let json = serde_json::json!({
"model_info": {
"llama.context_length": 131072,
"gemma.context_length": 8192
}
});
assert_eq!(super::parse_show_response(&json), Some(8192));
}
#[test]
fn parse_show_response_modelfile_only_no_model_info() {
let json = serde_json::json!({
"parameters": "stop \"<|end|>\"\nnum_ctx 16384\ntemperature 0.2"
});
assert_eq!(super::parse_show_response(&json), Some(16384));
}
#[test]
fn parse_show_response_ignores_unparsable_num_ctx() {
let json = serde_json::json!({"parameters": "num_ctx lots"});
assert_eq!(super::parse_show_response(&json), None);
}
#[test]
fn parse_show_response_parameters_without_num_ctx() {
let json = serde_json::json!({"parameters": "temperature 0.7\ntop_p 0.9"});
assert_eq!(super::parse_show_response(&json), None);
}
#[test]
fn parse_show_response_non_numeric_context_length_ignored() {
let json = serde_json::json!({"model_info": {"llama.context_length": "32768"}});
assert_eq!(super::parse_show_response(&json), None);
}
#[test]
fn parse_show_response_empty_json() {
assert_eq!(super::parse_show_response(&serde_json::json!({})), None);
}
#[test]
fn probe_tool_schema_is_single_list_dir_function() {
let schema = super::probe_tool_schema();
let arr = schema.as_array().expect("schema is a JSON array");
assert_eq!(arr.len(), 1, "probe uses exactly one tool");
let f = &arr[0];
assert_eq!(f["type"], "function");
assert_eq!(f["function"]["name"], "list_dir");
let params = &f["function"]["parameters"];
assert_eq!(params["properties"]["path"]["type"], "string");
assert_eq!(params["required"][0], "path");
}
#[test]
fn capability_entry_default_is_untested_no_tools() {
let e = CapabilityEntry::default();
assert_eq!(e.conformance, ToolConformance::NoTools);
assert!(e.tested_date.is_empty());
assert_eq!(e.context_window, None);
assert_eq!(e.safe_context, None);
assert_eq!(e.overflow_at, None);
assert_eq!(e.max_ok_input, None);
assert_eq!(e.consecutive_ok, 0);
assert_eq!(e.tune_confidence, TuneConfidence::None);
assert_eq!(e.tune_date, None);
}
#[test]
fn tune_confidence_default_is_none() {
assert_eq!(TuneConfidence::default(), TuneConfidence::None);
}
#[test]
fn print_capabilities_table_handles_empty_model_list() {
let cache = CapabilityCache::default();
print_capabilities_table(&[], &cache, "none", "http://localhost:11434", false);
}
#[test]
fn print_capabilities_table_renders_all_branches() {
let mut cache = CapabilityCache::default();
for (name, conf) in [
("m-none", TuneConfidence::None),
("m-low", TuneConfidence::Low),
("m-med", TuneConfidence::Medium),
("m-high", TuneConfidence::High),
] {
let mut e = make_entry();
e.tune_confidence = conf;
cache.insert(name.to_string(), e);
}
cache.insert(
"m-noctx".to_string(),
CapabilityEntry {
conformance: ToolConformance::TextMode,
tested_date: "2026-06-06".to_string(),
..Default::default()
},
);
cache.insert(
"m-think".to_string(),
CapabilityEntry {
conformance: ToolConformance::Native,
tested_date: "2026-06-10".to_string(),
emits_thinking: Some(true),
..Default::default()
},
);
let models: Vec<ModelInfo> = [
("m-none", "7B"),
("m-low", "13B"),
("m-med", ""),
("m-high", "32.8B"),
("m-noctx", "3B"),
("m-think", "30B"),
("m-untested", "1B"),
]
.into_iter()
.map(|(n, s)| ModelInfo {
name: n.to_string(),
param_size: s.to_string(),
})
.collect();
print_capabilities_table(&models, &cache, "m-low", "http://localhost:11434", false);
print_capabilities_table(&models, &cache, "m-high", "http://localhost:11434", true);
print_capabilities_table(&models, &cache, "absent", "http://localhost:11434", true);
}
}