use std::collections::HashMap;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use crate::error::ProbeError;
use crate::types::{CapabilityLevel, CapabilityProfile, TOOL_PROBE_NAMES};
pub const CACHE_TTL_SECS: u64 = 30 * 24 * 60 * 60;
pub const PROBE_SUITE_VERSION: u32 = 96;
pub const DEFAULT_PROBE_EFFORT: &str = "unset";
pub const DEFAULT_SKIP_EXPENSIVE: bool = false;
pub const DEFAULT_VISION: bool = false;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CacheEntry {
pub profile: CapabilityProfile,
pub cached_at: u64,
#[serde(default = "default_effort_label")]
pub reasoning_effort: String,
#[serde(default = "default_suite_v1")]
pub probe_suite_version: u32,
}
fn default_effort_label() -> String {
DEFAULT_PROBE_EFFORT.to_owned()
}
fn default_suite_v1() -> u32 {
1
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct ProbeCache {
pub profiles: HashMap<String, CacheEntry>,
}
impl ProbeCache {
pub fn load(path: &Path) -> Result<Self, ProbeError> {
if !path.exists() {
return Ok(Self::default());
}
let len = std::fs::metadata(path)?.len();
if len > 8 * 1024 * 1024 {
return Err(ProbeError::Internal(format!(
"probe cache is too large ({} bytes): {}",
len,
path.display()
)));
}
let mut cache = Self::read_disk(path)?;
if cache.migrate_stale_tool_scores() {
if let Err(err) = cache.save(path) {
eprintln!(
"warning: failed to persist migrated probe cache ({}): {err}",
path.display()
);
}
}
Ok(cache)
}
fn read_disk(path: &Path) -> Result<Self, ProbeError> {
let contents = std::fs::read_to_string(path)?;
Ok(serde_json::from_str(&contents)?)
}
pub fn save(&self, path: &Path) -> Result<(), ProbeError> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut profiles = self.profiles.clone();
if path.exists() {
if let Ok(disk) = Self::read_disk(path) {
for (key, theirs) in disk.profiles {
match profiles.get(&key) {
Some(ours) if ours.cached_at >= theirs.cached_at => {}
_ => {
profiles.insert(key, theirs);
}
}
}
}
}
let outgoing = Self { profiles };
let contents = serde_json::to_string_pretty(&outgoing)?;
let tmp = path.with_extension(format!("tmp-{}", std::process::id()));
std::fs::write(&tmp, contents)?;
#[cfg(windows)]
{
let bak = path.with_extension("bak");
if path.exists() {
let _ = std::fs::remove_file(&bak);
std::fs::rename(path, &bak)?;
}
if let Err(err) = std::fs::rename(&tmp, path) {
let _ = std::fs::rename(&bak, path);
return Err(err.into());
}
let _ = std::fs::remove_file(&bak);
}
#[cfg(not(windows))]
std::fs::rename(&tmp, path)?;
Ok(())
}
pub fn find_profile(&self, model_id: &str, provider: &str) -> Option<&CapabilityProfile> {
if let Some(profile) = self.get(model_id, provider) {
return Some(profile);
}
self.find_profile_with_cost(model_id, provider)
.map(|(profile, _)| profile)
}
pub fn find_profile_with_cost(
&self,
model_id: &str,
provider: &str,
) -> Option<(&CapabilityProfile, bool)> {
if let Some(profile) = self.get(model_id, provider) {
return Some((profile, DEFAULT_SKIP_EXPENSIVE));
}
self.profiles
.iter()
.filter(|(_, entry)| {
Self::is_valid(entry)
&& entry.probe_suite_version == PROBE_SUITE_VERSION
&& models_equivalent(
&entry.profile.model_id,
model_id,
provider,
&entry.profile.provider,
)
&& providers_equivalent(&entry.profile.provider, provider)
})
.max_by_key(|(_, entry)| entry.cached_at)
.map(|(key, entry)| (&entry.profile, key.split('|').nth(4) == Some("cheap")))
}
pub fn find_profile_with_cost_and_advertised(
&self,
model_id: &str,
provider: &str,
advertised: Option<u32>,
) -> Option<(&CapabilityProfile, bool)> {
if let Some(profile) = self.get_with_knobs(
model_id,
provider,
DEFAULT_SKIP_EXPENSIVE,
DEFAULT_VISION,
advertised,
) {
return Some((profile, DEFAULT_SKIP_EXPENSIVE));
}
self.profiles
.iter()
.filter(|(key, entry)| {
Self::is_valid(entry)
&& entry.probe_suite_version == PROBE_SUITE_VERSION
&& models_equivalent(
&entry.profile.model_id,
model_id,
provider,
&entry.profile.provider,
)
&& providers_equivalent(&entry.profile.provider, provider)
&& key_advertised(key) == advertised
})
.max_by_key(|(_, entry)| entry.cached_at)
.map(|(key, entry)| (&entry.profile, key.split('|').nth(4) == Some("cheap")))
}
pub fn get(&self, model_id: &str, provider: &str) -> Option<&CapabilityProfile> {
self.get_with_knobs(
model_id,
provider,
DEFAULT_SKIP_EXPENSIVE,
DEFAULT_VISION,
None,
)
}
pub fn get_with_knobs(
&self,
model_id: &str,
provider: &str,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) -> Option<&CapabilityProfile> {
self.get_with_settings(
model_id,
provider,
DEFAULT_PROBE_EFFORT,
PROBE_SUITE_VERSION,
skip_expensive,
vision,
advertised,
)
}
#[allow(clippy::too_many_arguments)]
pub fn get_with_settings(
&self,
model_id: &str,
provider: &str,
reasoning_effort: &str,
suite_version: u32,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) -> Option<&CapabilityProfile> {
let key = Self::cache_key_with_knobs(
model_id,
provider,
reasoning_effort,
suite_version,
skip_expensive,
vision,
advertised,
);
if let Some(profile) = self.profiles.get(&key).and_then(|entry| {
if Self::is_valid(entry) {
Some(&entry.profile)
} else {
None
}
}) {
return Some(profile);
}
self.profiles
.iter()
.filter(|(stored_key, entry)| {
Self::is_valid(entry)
&& entry.probe_suite_version == suite_version
&& entry.reasoning_effort == reasoning_effort
&& models_equivalent(
&entry.profile.model_id,
model_id,
provider,
&entry.profile.provider,
)
&& providers_equivalent(&entry.profile.provider, provider)
&& key_knobs_match(stored_key, skip_expensive, vision, advertised)
})
.max_by_key(|(_, entry)| entry.cached_at)
.map(|(_, entry)| &entry.profile)
}
pub fn get_entry(&self, model_id: &str, provider: &str) -> Option<&CacheEntry> {
self.get_entry_with_knobs(
model_id,
provider,
DEFAULT_SKIP_EXPENSIVE,
DEFAULT_VISION,
None,
)
}
pub fn get_entry_with_knobs(
&self,
model_id: &str,
provider: &str,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) -> Option<&CacheEntry> {
let key = Self::cache_key_with_knobs(
model_id,
provider,
DEFAULT_PROBE_EFFORT,
PROBE_SUITE_VERSION,
skip_expensive,
vision,
advertised,
);
if let Some(entry) = self.profiles.get(&key).filter(|e| Self::is_valid(e)) {
return Some(entry);
}
self.profiles
.iter()
.filter(|(stored_key, entry)| {
Self::is_valid(entry)
&& entry.probe_suite_version == PROBE_SUITE_VERSION
&& entry.reasoning_effort == DEFAULT_PROBE_EFFORT
&& models_equivalent(
&entry.profile.model_id,
model_id,
provider,
&entry.profile.provider,
)
&& providers_equivalent(&entry.profile.provider, provider)
&& key_knobs_match(stored_key, skip_expensive, vision, advertised)
})
.max_by_key(|(_, entry)| entry.cached_at)
.map(|(_, entry)| entry)
}
pub fn put(&mut self, profile: CapabilityProfile) {
self.put_with_knobs(profile, DEFAULT_SKIP_EXPENSIVE, DEFAULT_VISION, None);
}
pub fn put_with_knobs(
&mut self,
profile: CapabilityProfile,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) {
self.put_with_settings(
profile,
DEFAULT_PROBE_EFFORT,
PROBE_SUITE_VERSION,
skip_expensive,
vision,
advertised,
);
}
pub fn put_with_settings(
&mut self,
profile: CapabilityProfile,
reasoning_effort: &str,
suite_version: u32,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) {
let key = Self::cache_key_with_knobs(
&profile.model_id,
&profile.provider,
reasoning_effort,
suite_version,
skip_expensive,
vision,
advertised,
);
let entry = CacheEntry {
cached_at: unix_now(),
profile,
reasoning_effort: reasoning_effort.to_owned(),
probe_suite_version: suite_version,
};
self.profiles.insert(key, entry);
}
fn migrate_stale_tool_scores(&mut self) -> bool {
let mut changed = false;
for entry in self.profiles.values_mut() {
for name in TOOL_PROBE_NAMES {
let Some(probe) = entry.profile.dimension_result_mut(name) else {
continue;
};
if probe.details.contains("does not support tools")
&& probe.level != CapabilityLevel::Weak
{
probe.level = CapabilityLevel::Weak;
probe.score = 0.0;
changed = true;
}
}
}
changed
}
pub fn cache_key(
model_id: &str,
provider: &str,
reasoning_effort: &str,
suite_version: u32,
) -> String {
Self::cache_key_with_knobs(
model_id,
provider,
reasoning_effort,
suite_version,
DEFAULT_SKIP_EXPENSIVE,
DEFAULT_VISION,
None,
)
}
#[allow(clippy::too_many_arguments)]
pub fn cache_key_with_knobs(
model_id: &str,
provider: &str,
reasoning_effort: &str,
suite_version: u32,
skip_expensive: bool,
vision: bool,
advertised: Option<u32>,
) -> String {
let cost = if skip_expensive { "cheap" } else { "full" };
let vis = if vision { "vision" } else { "novision" };
let ctx = match advertised {
Some(n) => format!("ctx{n}"),
None => "ctxnone".to_owned(),
};
format!("{model_id}|{provider}|{reasoning_effort}|v{suite_version}|{cost}|{vis}|{ctx}")
}
fn is_valid(entry: &CacheEntry) -> bool {
let now = unix_now();
now.saturating_sub(entry.cached_at) < CACHE_TTL_SECS
}
}
fn providers_equivalent(stored: &str, requested: &str) -> bool {
let a = stored.to_ascii_lowercase();
let b = requested.to_ascii_lowercase();
if a == b {
return true;
}
provider_family(&a) == provider_family(&b)
}
fn models_equivalent(
stored: &str,
requested: &str,
requested_provider: &str,
stored_provider: &str,
) -> bool {
if stored == requested {
return true;
}
let stored_n = strip_normalized_provider_prefix(stored, requested_provider)
.or_else(|| strip_normalized_provider_prefix(stored, stored_provider))
.unwrap_or(stored);
let requested_n = strip_normalized_provider_prefix(requested, requested_provider)
.or_else(|| strip_normalized_provider_prefix(requested, stored_provider))
.unwrap_or(requested);
stored_n == requested_n
}
fn strip_normalized_provider_prefix<'a>(model_id: &'a str, provider: &str) -> Option<&'a str> {
let lower_model = model_id.to_ascii_lowercase();
let provider_l = provider.to_ascii_lowercase();
let family = provider_family(&provider_l);
for prefix in [provider_l.as_str(), family] {
let with_slash = format!("{prefix}/");
if let Some(rest) = lower_model
.strip_prefix(&with_slash)
.and_then(|_| model_id.get(with_slash.len()..))
{
if !rest.is_empty() {
return Some(rest);
}
}
}
None
}
fn provider_family(provider: &str) -> &str {
match provider {
"openai" | "api.openai.com" => "openai",
"openrouter" | "openrouter.ai" => "openrouter",
"xai" | "grok" | "api.x.ai" | "x.ai" => "xai",
"claude" | "anthropic" | "api.anthropic.com" => "anthropic",
"ollama" | "localhost" | "127.0.0.1" | "::1" | "[::1]" | "0.0.0.0" => "ollama",
other if loopback_default_ollama_port(other) => "ollama",
other => other,
}
}
fn loopback_default_ollama_port(provider: &str) -> bool {
let Some(host) = provider.strip_suffix(":11434") else {
return false;
};
let bare = host
.strip_prefix('[')
.and_then(|h| h.strip_suffix(']'))
.unwrap_or(host);
matches!(bare, "localhost" | "127.0.0.1" | "0.0.0.0" | "::1")
}
fn key_knobs_match(key: &str, skip_expensive: bool, vision: bool, advertised: Option<u32>) -> bool {
let mut parts = key.rsplit('|');
let ctx = parts.next().unwrap_or("");
let vis = parts.next().unwrap_or("");
let cost = parts.next().unwrap_or("");
let want_cost = if skip_expensive { "cheap" } else { "full" };
let want_vis = if vision { "vision" } else { "novision" };
let want_ctx = match advertised {
Some(n) => format!("ctx{n}"),
None => "ctxnone".to_owned(),
};
cost == want_cost && vis == want_vis && ctx == want_ctx
}
fn key_advertised(key: &str) -> Option<u32> {
let ctx = key.rsplit('|').next().unwrap_or("");
ctx.strip_prefix("ctx")
.filter(|s| *s != "none")
.and_then(|s| s.parse().ok())
}
fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}