pub(crate) mod evidence;
use reqwest::blocking::Client;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::{
collections::BTreeMap,
fs,
path::PathBuf,
sync::{Mutex, OnceLock},
time::{Duration, Instant},
};
pub(crate) const TYPESAFE_API_KEY_ENV: &str = "TYPESAFE_API_KEY";
const SYSTEM_ONE_URL: &str = "https://api.typesafe.ai/v1/systemone";
const CACHE_DIRECTORY: &str = "jev";
pub(crate) const JEV_MODEL: &str = "jev-1.13.0";
const MAX_INSTRUCTIONS_BYTES: usize = 4 * 1024;
static CACHE_LOCK: OnceLock<Mutex<u64>> = OnceLock::new();
#[derive(Debug, Clone)]
pub(crate) struct TypeSafeClient {
http: Client,
api_key: String,
endpoint: String,
cache_dir: Option<PathBuf>,
timeout: Duration,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct JevResult<T> {
pub(crate) value: T,
pub(crate) cached: bool,
pub(crate) model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub(crate) struct ScoreAnswer {
pub(crate) score: f64,
pub(crate) confidence: f64,
pub(crate) probabilities: BTreeMap<String, f64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub(crate) enum JevQuestion {
Choice {
instructions: Value,
criteria: BTreeMap<String, Value>,
},
Score {
instructions: Value,
criteria: Vec<serde_json::Value>,
},
Noul {
instructions: Value,
criteria: Option<NoulCriteriaOwned>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct NoulCriteriaOwned {
#[serde(rename = "true")]
pub(crate) yes: Value,
#[serde(rename = "false")]
pub(crate) no: Value,
}
impl JevQuestion {
pub(crate) fn validate(&self) -> anyhow::Result<()> {
let (instructions, count) = match self {
Self::Choice {
instructions,
criteria,
} => (instructions, criteria.len()),
Self::Score {
instructions,
criteria,
} => (instructions, criteria.len()),
Self::Noul { instructions, .. } => (instructions, 2),
};
validate_description(instructions)?;
anyhow::ensure!(
serde_json::to_vec(instructions)?.len() <= MAX_INSTRUCTIONS_BYTES,
"instructions too large"
);
match self {
Self::Choice { criteria, .. } => {
for description in criteria.values() {
validate_description(description)?;
}
}
Self::Score { criteria, .. } => {
for description in criteria {
validate_description(description)?;
}
}
Self::Noul {
criteria: Some(criteria),
..
} => {
validate_description(&criteria.yes)?;
validate_description(&criteria.no)?;
}
Self::Noul { criteria: None, .. } => {}
}
anyhow::ensure!(
(2..=255).contains(&count),
"criteria must contain between 2 and 255 entries"
);
if matches!(self, Self::Score { .. }) {
anyhow::ensure!(count <= 10, "score supports at most 10 criteria");
}
Ok(())
}
}
fn validate_description(value: &Value) -> anyhow::Result<()> {
anyhow::ensure!(
matches!(
value,
Value::String(_) | Value::Object(_) | Value::Array(_) | Value::Null
),
"descriptions must be strings, objects, arrays, or null"
);
Ok(())
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub(crate) enum JevAnswer {
Choice {
choice: String,
confidence: f64,
probabilities: BTreeMap<String, f64>,
},
Score {
score: f64,
confidence: f64,
probabilities: BTreeMap<String, f64>,
},
Noul {
noul: f64,
},
}
#[derive(Serialize)]
struct SystemOneRequest<'a, S: ?Sized, K, Q> {
state: &'a S,
model: &'static str,
questions: BTreeMap<K, Q>,
}
#[derive(Serialize, Deserialize)]
struct SystemOneResponse<A> {
#[serde(default = "default_model")]
model: String,
answers: BTreeMap<String, A>,
}
fn default_model() -> String {
"unreported".to_string()
}
fn cache_directory() -> anyhow::Result<PathBuf> {
Ok(crate::config::McPaths::resolve()?
.cache
.join(CACHE_DIRECTORY))
}
pub(crate) fn reset_cache(paths: &crate::config::McPaths) -> anyhow::Result<()> {
let mut generation = CACHE_LOCK
.get_or_init(|| Mutex::new(0))
.lock()
.map_err(|_| anyhow::anyhow!("Jev cache lock poisoned"))?;
*generation = generation.wrapping_add(1);
match fs::remove_dir_all(paths.cache.join(CACHE_DIRECTORY)) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
}
}
fn cached_response<A>(
client: &TypeSafeClient,
request: &impl Serialize,
) -> anyhow::Result<JevResult<SystemOneResponse<A>>>
where
A: Serialize + DeserializeOwned,
{
let request_bytes = serde_json::to_vec(request)?;
anyhow::ensure!(
request_bytes.len() <= 128 * 1024,
"Jev request exceeds local byte budget"
);
let cache_path = client.cache_dir.as_ref().map(|directory| {
directory.join(format!(
"{}.json",
crate::hex::lower_hex(Sha256::digest(&request_bytes))
))
});
let generation = if let Some(path) = &cache_path {
let guard = CACHE_LOCK
.get_or_init(|| Mutex::new(0))
.lock()
.map_err(|_| anyhow::anyhow!("Jev cache lock poisoned"))?;
if let Ok(file) = crate::persistence::open_regular_file(path, false)
&& let Ok(bytes) = crate::persistence::read_regular_file_bounded(&file, 2 * 1024 * 1024)
&& let Ok(response) = serde_json::from_slice::<SystemOneResponse<A>>(&bytes)
{
return Ok(JevResult {
model: response.model.clone(),
value: response,
cached: true,
});
}
Some(*guard)
} else {
None
};
let deadline = Instant::now() + client.timeout;
let mut attempt = 0;
let response = loop {
let remaining = deadline.saturating_duration_since(Instant::now());
anyhow::ensure!(!remaining.is_zero(), "Jev assessment timed out");
let response = client
.http
.post(&client.endpoint)
.bearer_auth(&client.api_key)
.body(request_bytes.clone())
.header(reqwest::header::CONTENT_TYPE, "application/json")
.timeout(remaining)
.send()?;
let status = response.status();
if status.is_success() {
break response;
}
let retryable =
status.as_u16() == 429 || status.as_u16() == 408 || status.is_server_error();
if attempt == 0 && retryable && client.cache_dir.is_some() {
let delay = match response.headers().get(reqwest::header::RETRY_AFTER) {
Some(header) => header
.to_str()
.ok()
.and_then(|value| value.parse::<u64>().ok())
.map(Duration::from_secs),
None => Some(Duration::from_millis(100)),
};
if let Some(delay) =
delay.filter(|delay| *delay < deadline.saturating_duration_since(Instant::now()))
{
drop(response);
std::thread::sleep(delay);
attempt += 1;
continue;
}
}
anyhow::bail!("TypeSafe API returned HTTP {status}");
};
let response = serde_json::from_str::<SystemOneResponse<A>>(
&crate::http_body::read_bounded_response_text(response, 2 * 1024 * 1024)?,
)?;
if let Some(path) = cache_path
&& let Ok(guard) = CACHE_LOCK.get_or_init(|| Mutex::new(0)).lock()
&& generation == Some(*guard)
{
let _ = crate::persistence::atomic_write_with_permissions(
&path,
&serde_json::to_vec(&response)?,
Some(0o600),
);
}
Ok(JevResult {
model: response.model.clone(),
value: response,
cached: false,
})
}
impl TypeSafeClient {
pub(crate) fn from_environment() -> anyhow::Result<Self> {
Self::from_environment_with_options(Duration::from_secs(15), Some(cache_directory()?))
}
pub(crate) fn for_background_inference(timeout: Duration) -> anyhow::Result<Self> {
Self::from_environment_with_options(timeout, None)
}
fn from_environment_with_options(
timeout: Duration,
cache_dir: Option<PathBuf>,
) -> anyhow::Result<Self> {
let api_key = std::env::var(TYPESAFE_API_KEY_ENV).map_err(|_| {
anyhow::anyhow!("{TYPESAFE_API_KEY_ENV} is required when Jev tooling is enabled")
})?;
anyhow::ensure!(
!api_key.trim().is_empty(),
"{TYPESAFE_API_KEY_ENV} must not be empty"
);
Ok(Self {
http: Client::builder().timeout(timeout).build()?,
api_key,
endpoint: SYSTEM_ONE_URL.to_string(),
cache_dir,
timeout,
})
}
pub(crate) fn ask_many<S: Serialize + ?Sized>(
&self,
state: &S,
questions: BTreeMap<String, JevQuestion>,
) -> anyhow::Result<JevResult<BTreeMap<String, JevAnswer>>> {
anyhow::ensure!(!questions.is_empty(), "questions must not be empty");
for question in questions.values() {
question.validate()?;
}
let request = SystemOneRequest {
state,
model: JEV_MODEL,
questions,
};
let response = cached_response::<JevAnswer>(self, &request)?;
Ok(JevResult {
value: response.value.answers,
cached: response.cached,
model: response.model,
})
}
}