ai.rs
Rust · 560 lines · 25.8 KB
//! Remote model access: OpenAI-compatible HTTP APIs (LM Studio, Azure OpenAI v1, others),
//! for embeddings and chat. Azure signs in without keys: managed identity on an Azure VM,
//! else the user's `az login`.
use anyhow::{anyhow, bail, Context, Result};
use serde_json::{json, Value};
use std::path::PathBuf;
use std::sync::Mutex;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
const AZURE_SCOPE: &str = "https://cognitiveservices.azure.com";
#[derive(Clone, Debug, PartialEq)]
pub enum Kind {
/// Plain OpenAI-compatible server (LM Studio, vLLM, Ollama's /v1, …).
OpenAi,
/// Azure OpenAI through its v1 API (`<resource>/openai/v1`): deployment names as `model`.
Azure,
}
#[derive(Clone, Debug)]
pub enum Auth {
None,
/// API key read from an environment variable (never stored in the brain).
KeyEnv(String),
/// Azure keyless: managed identity if on an Azure VM, else `az login`. Holds the token
/// resource (default Cognitive Services; Microsoft Foundry resources may need https://ai.azure.com).
Azure(String),
/// A command that prints a bearer token (for company-specific sign-in tools).
Command(String),
}
impl Auth {
/// For messages: the kind of sign-in, never a command line (it may hold a credential).
pub fn label(&self) -> String {
match self {
Auth::None => "none".into(),
Auth::KeyEnv(v) => format!("key from ${v}"),
Auth::Azure(_) => "Azure sign-in".into(),
Auth::Command(_) => "token command".into(),
}
}
pub fn parse(s: &str) -> Auth {
let s = s.trim();
if let Some(v) = s.strip_prefix("env:") {
Auth::KeyEnv(v.to_string())
} else if let Some(c) = s.strip_prefix("command:") {
Auth::Command(c.to_string())
} else if s == "azure" || s == "managed-identity" || s == "az" {
Auth::Azure(AZURE_SCOPE.to_string())
} else if let Some(scope) = s.strip_prefix("azure:") {
Auth::Azure(scope.trim_end_matches("/.default").trim_end_matches('/').to_string())
} else {
Auth::None
}
}
}
pub struct Client {
pub kind: Kind,
pub base: String,
auth: Auth,
agent: ureq::Agent,
token: Mutex<Option<(String, Instant)>>,
/// Where an Azure token is cached between runs (the user's private cache folder, outside the brain).
token_cache: Option<PathBuf>,
/// Hosts besides Azure's own that may receive an Azure sign-in token (e.g. a company API gateway).
token_hosts: Vec<String>,
/// Masks personal details in text sent to a model that isn't on this machine.
masker: crate::privacy::Masker,
/// Set once an OpenAI-style service rejected `max_tokens`/`temperature` (OpenAI's GPT-5 and
/// reasoning models do): chat then sends `max_completion_tokens` and the default temperature.
newer_chat_params: std::sync::atomic::AtomicBool,
}
/// Azure OpenAI / AI Services endpoint domains: an Azure sign-in token is only ever sent to these
/// (or to hosts listed in `token_hosts`). API Management gateways (`*.azure-api.net`) are not
/// trusted automatically: anyone can create one, so a company gateway has to be listed.
pub const AZURE_HOST_SUFFIXES: &[&str] = &[
".openai.azure.com",
".cognitiveservices.azure.com",
".services.ai.azure.com",
".openai.azure.us",
".cognitiveservices.azure.us",
".openai.azure.cn",
".cognitiveservices.azure.cn",
];
/// The host of an http(s) URL, lowercased, without port; None if it isn't a plain http(s) URL
/// (credentials in the URL are refused too).
pub fn url_host(url: &str) -> Option<(bool, String)> {
let lower = url.trim().to_ascii_lowercase();
let (https, rest) = match lower.strip_prefix("https://") {
Some(r) => (true, r),
None => (false, lower.strip_prefix("http://")?),
};
let authority = rest.split(['/', '?', '#']).next().unwrap_or("");
if authority.contains('@') || authority.is_empty() {
return None;
}
let host = if let Some(v6) = authority.strip_prefix('[') {
v6.split(']').next().unwrap_or("").to_string()
} else {
authority.split(':').next().unwrap_or("").to_string()
};
(!host.is_empty()).then_some((https, host))
}
/// This machine: exactly `localhost`, or an IP address in 127.0.0.0/8 or ::1 (a name like
/// `127.evil.example` is not).
pub fn is_loopback(host: &str) -> bool {
host == "localhost" || host.parse::<std::net::IpAddr>().is_ok_and(|ip| ip.is_loopback())
}
/// HTTP agent that trusts the operating system's certificate store (Windows, macOS, Linux), so
/// company TLS-inspection certificates installed there work.
/// `direct` (for this machine's own servers) ignores HTTP_PROXY/HTTPS_PROXY/ALL_PROXY: plain-http
/// traffic to localhost is unmasked and may carry a key, so it must never go through a proxy.
fn http_agent(timeout: Duration, direct: bool) -> ureq::Agent {
use ureq::tls::{RootCerts, TlsConfig};
let b = ureq::Agent::config_builder();
let b = if direct { b.proxy(None) } else { b };
b.timeout_global(Some(timeout))
.http_status_as_error(false)
// Never follow redirects: a redirect to another host would carry the API key with it.
.max_redirects(0)
.tls_config(TlsConfig::builder().root_certs(RootCerts::PlatformVerifier).build())
.build()
.into()
}
fn unix_now() -> u64 {
SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)
}
fn read_cache(path: &PathBuf) -> Value {
std::fs::read_to_string(path).ok().and_then(|s| serde_json::from_str(&s).ok()).unwrap_or(json!({}))
}
/// Write the token cache privately: the temp file is created owner-only (Unix), then renamed over.
fn write_cache(path: &PathBuf, v: &Value) {
use std::io::Write;
if let Some(dir) = path.parent() {
crate::store::create_private_dir(dir);
}
let tmp = path.with_extension(format!("tmp{}", std::process::id()));
let mut o = std::fs::OpenOptions::new();
o.write(true).create_new(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
o.mode(0o600);
}
let _ = std::fs::remove_file(&tmp);
if let Ok(mut f) = o.open(&tmp)
&& f.write_all(v.to_string().as_bytes()).is_ok()
{
drop(f);
if std::fs::rename(&tmp, path).is_ok() {
return;
}
}
let _ = std::fs::remove_file(&tmp);
}
/// Cache key for an Azure token: the scope it was issued for and the identity that asked.
fn cache_key(scope: &str, host: &str) -> String {
format!("{scope}|{}|{host}", std::env::var("AZURE_CLIENT_ID").unwrap_or_default())
}
impl Client {
pub fn new(kind: Kind, url: &str, auth: Auth, timeout_secs: u64) -> Client {
let url = url.trim().trim_end_matches('/');
let base = match kind {
// Accept the resource endpoint as shown in the Azure portal.
// Accept the resource endpoint, …/openai, or …/openai/v1.
Kind::Azure if url.ends_with("/openai/v1") => url.to_string(),
Kind::Azure if url.ends_with("/openai") => format!("{url}/v1"),
Kind::Azure => format!("{url}/openai/v1"),
_ => url.to_string(),
};
let agent = http_agent(Duration::from_secs(timeout_secs), url_host(&base).is_some_and(|(_, h)| is_loopback(&h)));
Client {
kind,
base,
auth,
agent,
token: Mutex::new(None),
token_cache: None,
token_hosts: vec![],
masker: crate::privacy::Masker::off(),
newer_chat_params: std::sync::atomic::AtomicBool::new(false),
}
}
fn host(&self) -> String {
url_host(&self.base).map(|h| h.1).unwrap_or_default()
}
/// Mask personal details before text goes to this endpoint, unless it is on this machine.
pub fn with_privacy(mut self, cfg: &crate::store::Privacy) -> Self {
let local = url_host(&self.base).is_some_and(|(_, h)| is_loopback(&h));
self.masker = if local { crate::privacy::Masker::off() } else { crate::privacy::Masker::new(cfg) };
self
}
/// Text about to leave for this endpoint: never anything that looks like a secret (whatever its
/// origin: an entry, a search query, a diff), and masked if the endpoint isn't on this machine.
fn outbound(&self, text: &str) -> Result<String> {
if let Some(what) = crate::store::secret_kind(text) {
bail!("not sent: the text looks like it contains {what}");
}
self.masker.mask(text)
}
/// Extra hosts allowed to receive an Azure sign-in token (from the config's `token_hosts`).
pub fn with_token_hosts(mut self, hosts: &[String]) -> Self {
self.token_hosts = hosts.iter().map(|h| h.trim().to_ascii_lowercase()).filter(|h| !h.is_empty()).collect();
self
}
/// Refuse to send credentials where they could leak: over plain http (except to this machine),
/// or an Azure sign-in token to a host that isn't Azure's (or listed in `token_hosts`).
pub fn check_destination(&self) -> Result<()> {
if self.base.contains(['?', '#']) {
bail!("the endpoint URL must not have a ?query or #fragment (keys belong in `auth`, not the URL)");
}
// Checked for every auth mode: a `user:password@` URL would send Basic credentials by itself.
let (https, host) = url_host(&self.base)
.ok_or_else(|| anyhow!("the endpoint isn't a plain http(s) URL (and must not hold user:password@)"))?;
if matches!(self.auth, Auth::None) {
return Ok(());
}
if !https && !is_loopback(&host) {
bail!("refusing to send credentials over plain http to {host}: use https");
}
if let Auth::Azure(_) = self.auth
&& !is_loopback(&host)
&& !AZURE_HOST_SUFFIXES.iter().any(|s| host.ends_with(s))
&& !self.token_hosts.contains(&host)
{
bail!(
"refusing to send your Azure sign-in token to {host}: it isn't an Azure domain. If it is your \
company's gateway, add it to `token_hosts` in that section of config.toml"
);
}
Ok(())
}
/// Cache Azure tokens in this file between runs, so each command doesn't sign in again.
pub fn with_token_cache(mut self, path: PathBuf) -> Self {
self.token_cache = Some(path);
self
}
/// (header name, value) for this client's auth, if any.
fn auth_header(&self) -> Result<Option<(&'static str, String)>> {
self.check_destination()?;
Ok(match &self.auth {
Auth::None => None,
Auth::KeyEnv(var) => {
let key = std::env::var(var).map_err(|_| anyhow!("environment variable {var} (API key) isn't set"))?;
Some(if self.kind == Kind::Azure { ("api-key", key) } else { ("Authorization", format!("Bearer {key}")) })
}
Auth::Azure(_) | Auth::Command(_) => {
let mut cached = self.token.lock().unwrap();
if let Some((t, until)) = cached.as_ref()
&& Instant::now() < *until {
return Ok(Some(("Authorization", format!("Bearer {t}"))));
}
// A still-valid token from an earlier run (Azure only; command tokens are re-fetched).
if let (Auth::Azure(scope), Some(path)) = (&self.auth, &self.token_cache) {
let c = read_cache(path);
let c = &c["tokens"][cache_key(scope, &self.host())];
if let (Some(t), Some(exp)) = (c["token"].as_str(), c["expires"].as_u64())
&& exp > unix_now() + 120 {
*cached = Some((t.to_string(), Instant::now() + Duration::from_secs(exp - unix_now() - 120)));
return Ok(Some(("Authorization", format!("Bearer {t}"))));
}
}
let (t, ttl) = match &self.auth {
Auth::Command(cmd) => (run_token_command(cmd)?, Duration::from_secs(600)),
_ => {
let skip_imds = self.token_cache.as_ref().is_some_and(|p| read_cache(p)["no_imds_until"].as_u64().unwrap_or(0) > unix_now());
let Auth::Azure(scope) = &self.auth else { unreachable!() };
let (t, ttl, from_imds) = azure_token(scope, skip_imds)?;
if let Some(path) = &self.token_cache {
let mut c = read_cache(path);
if !c.is_object() || c.get("token").is_some() {
c = json!({}); // older single-token format
}
let now = unix_now();
if let Some(m) = c.get_mut("tokens").and_then(Value::as_object_mut) {
m.retain(|_, v| v["expires"].as_u64().unwrap_or(0) > now);
}
if !c["tokens"].is_object() {
c["tokens"] = json!({});
}
c["tokens"][cache_key(scope, &self.host())] = json!({"token": t, "expires": now + ttl.as_secs()});
if !from_imds {
// Not an Azure VM: don't wait on the metadata endpoint for a day.
c["no_imds_until"] = json!(unix_now() + 86400);
}
write_cache(path, &c);
}
(t, ttl)
}
};
*cached = Some((t.clone(), Instant::now() + ttl));
Some(("Authorization", format!("Bearer {t}")))
}
})
}
fn post(&self, path: &str, body: &Value) -> Result<Value> {
let url = format!("{}/{path}", self.base);
let mut req = self.agent.post(&url);
if let Some((k, v)) = self.auth_header()? {
req = req.header(k, v);
}
let mut resp = req.send_json(body).map_err(|e| anyhow!("request to {url} failed: {}", transport_error(&e)))?;
let status = resp.status().as_u16();
let text = resp.body_mut().read_to_string().unwrap_or_default();
if !(200..300).contains(&status) {
// The reply body is the service's text, not ours: it could carry instructions or echo a
// key, and errors reach agents. Report only the status and a short error code.
if older_params_refused(status, &text) {
bail!("{url} returned HTTP {status}: {OLDER_PARAMS_REFUSED}");
}
let code = serde_json::from_str::<Value>(&text).ok().and_then(|v| {
let e = &v["error"];
e["code"].as_str().or(e["type"].as_str()).filter(|c| is_plain_code(c)).map(str::to_string)
});
bail!("{url} returned HTTP {status}{}", code.map(|c| format!(" ({c})")).unwrap_or_default());
}
serde_json::from_str(&text).with_context(|| format!("{url} returned non-JSON"))
}
/// Check the endpoint and credentials without spending tokens: GET <base>/models.
/// Returns the HTTP status (not every deployment serves the model list, so callers decide).
pub fn ping(&self) -> Result<u16> {
let url = format!("{}/models", self.base);
let mut req = self.agent.get(&url);
if let Some((k, v)) = self.auth_header()? {
req = req.header(k, v);
}
let resp = req.call().map_err(|e| anyhow!("{url}: {}", transport_error(&e)))?;
Ok(resp.status().as_u16())
}
pub fn reachable(&self) -> bool {
// Unauthenticated probe; still never follow redirects.
if self.check_destination().is_err() {
return false;
}
let direct = url_host(&self.base).is_some_and(|(_, h)| is_loopback(&h));
http_agent(Duration::from_secs(3), direct).get(format!("{}/models", self.base)).call().is_ok_and(|r| r.status().as_u16() < 500)
}
pub fn embeddings(&self, model: &str, texts: &[String]) -> Result<Vec<Vec<f32>>> {
let mut out = Vec::with_capacity(texts.len());
for batch in texts.chunks(64) {
let batch: Vec<String> = batch.iter().map(|t| self.outbound(t)).collect::<Result<_>>()?;
let v = self.post("embeddings", &json!({"model": model, "input": batch}))?;
let data = v["data"].as_array().ok_or_else(|| anyhow!("embeddings reply has no data"))?;
let mut rows: Vec<(usize, Vec<f32>)> = data
.iter()
.map(|d| {
let i = d["index"].as_u64().unwrap_or(0) as usize;
let e = d["embedding"]
.as_array()
.map(|xs| xs.iter().filter_map(|x| x.as_f64()).map(|x| x as f32).collect())
.unwrap_or_default();
(i, e)
})
.collect();
rows.sort_by_key(|r| r.0);
if rows.len() != batch.len() {
bail!("embeddings reply had {} vectors for {} inputs", rows.len(), batch.len());
}
out.extend(rows.into_iter().map(|r| r.1));
}
Ok(out)
}
pub fn chat(&self, model: &str, system: &str, user: &str, max_tokens: u32) -> Result<String> {
let messages = json!([{"role": "system", "content": system}, {"role": "user", "content": self.outbound(user)?}]);
// Azure's GPT-5-family deployments (and OpenAI's own GPT-5/reasoning models) take
// max_completion_tokens and only the default temperature. Other OpenAI-compatible servers
// (LM Studio, Ollama, vLLM, …) may only know max_tokens, so they get that until they refuse it.
use std::sync::atomic::Ordering::Relaxed;
let newer = |m: &Value| json!({"model": model, "messages": m, "max_completion_tokens": max_tokens});
let v = if self.kind == Kind::Azure || self.newer_chat_params.load(Relaxed) {
self.post("chat/completions", &newer(&messages))?
} else {
let older = json!({"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": 0.1});
match self.post("chat/completions", &older) {
Err(e) if e.to_string().contains(OLDER_PARAMS_REFUSED) => {
self.newer_chat_params.store(true, Relaxed);
self.post("chat/completions", &newer(&messages))?
}
other => other?,
}
};
// Thinking models can spend the whole budget before answering.
match v["choices"][0]["message"]["content"].as_str() {
Some(s) if !s.trim().is_empty() => Ok(s.to_string()),
_ => Err(anyhow!(
"model returned no answer (finish_reason: {})",
v["choices"][0]["finish_reason"].as_str().filter(|c| is_plain_code(c)).unwrap_or("?")
)),
}
}
}
/// A network failure as one of a few fixed descriptions. ureq's own message can quote a proxy URL
/// (with any user:password in it) or server text, and errors reach agents.
fn transport_error(e: &ureq::Error) -> &'static str {
match e {
ureq::Error::Timeout(_) => "timed out",
ureq::Error::HostNotFound => "host not found (DNS)",
ureq::Error::ConnectionFailed => "couldn't connect",
ureq::Error::Io(_) => "connection error",
ureq::Error::TooManyRedirects | ureq::Error::RedirectFailed => "redirected (redirects are never followed)",
_ => {
let d = format!("{e:?}").to_ascii_lowercase();
if d.contains("proxy") {
"proxy error (check HTTPS_PROXY)"
} else if d.contains("tls") || d.contains("certificate") || d.contains("rustls") {
"TLS/certificate error"
} else {
"network error"
}
}
}
}
const OLDER_PARAMS_REFUSED: &str = "the model refuses max_tokens/temperature";
/// Whether an HTTP 400 reply says the service wants `max_completion_tokens` / the default temperature.
pub fn older_params_refused(status: u16, body: &str) -> bool {
let b = body.to_ascii_lowercase();
status == 400 && (b.contains("max_tokens") || b.contains("max_completion_tokens") || b.contains("temperature"))
}
/// Error codes and finish reasons OpenAI, Azure OpenAI and LM Studio use. Only these are shown:
/// any other text from the service (however harmless it looks) never reaches agents.
const KNOWN_CODES: [&str; 26] = [
"invalid_request_error", "invalid_api_key", "authentication_error", "permission_error", "unauthorized",
"insufficient_quota", "rate_limit_exceeded", "rate_limit_error", "429", "model_not_found", "not_found_error",
"deploymentnotfound", "unsupported_parameter", "unsupported_value", "invalid_value", "context_length_exceeded",
"content_filter", "content_policy_violation", "server_error", "overloaded_error", "timeout", "stop", "length",
"tool_calls", "function_call", "contentfilter",
];
pub(crate) fn is_plain_code(c: &str) -> bool {
KNOWN_CODES.contains(&c.to_ascii_lowercase().as_str())
}
fn run_token_command(cmd: &str) -> Result<String> {
// Neither the command (it may hold a credential argument) nor its stderr (it may echo a token)
// goes into errors: those reach agents.
let o = shell(cmd).output().context("couldn't run the token command set in `auth`")?;
if !o.status.success() {
bail!("the token command set in `auth` failed ({}); run it yourself to see why", o.status);
}
Ok(String::from_utf8_lossy(&o.stdout).trim().to_string())
}
fn shell(cmd: &str) -> std::process::Command {
if cfg!(windows) {
let mut c = std::process::Command::new("cmd");
c.args(["/C", cmd]);
c
} else {
let mut c = std::process::Command::new("sh");
c.args(["-c", cmd]);
c
}
}
/// Keyless Azure token: managed identity (Azure VM metadata endpoint), else the Azure CLI login.
/// Returns (token, lifetime, came from managed identity).
/// A token scope must be a plain https URL: it ends up in a request to Azure and an `az` argument.
pub fn valid_scope(scope: &str) -> bool {
scope.starts_with("https://")
&& scope.len() > 8
&& scope.chars().all(|c| c.is_ascii_alphanumeric() || ".-_:/".contains(c))
}
fn pct_encode(s: &str) -> String {
s.bytes()
.map(|b| if b.is_ascii_alphanumeric() || b"-._~".contains(&b) { (b as char).to_string() } else { format!("%{b:02X}") })
.collect()
}
fn azure_token(scope: &str, skip_imds: bool) -> Result<(String, Duration, bool)> {
if !valid_scope(scope) {
bail!("the Azure token scope {scope:?} must be an https URL like https://cognitiveservices.azure.com");
}
if !skip_imds {
let imds: ureq::Agent = ureq::Agent::config_builder()
.timeout_global(Some(Duration::from_secs(2)))
.proxy(None)
// The metadata endpoint never redirects; following one would send `Metadata: true` elsewhere.
.max_redirects(0)
.build()
.into();
// With several user-assigned identities, AZURE_CLIENT_ID picks one (as Azure SDKs do).
let client = match std::env::var("AZURE_CLIENT_ID") {
Ok(c) if c.chars().all(|ch| ch.is_ascii_hexdigit() || ch == '-') && !c.is_empty() => format!("&client_id={c}"),
Ok(c) if !c.is_empty() => bail!("AZURE_CLIENT_ID doesn't look like a client id (a GUID)"),
_ => String::new(),
};
let url = format!(
"http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource={}{client}",
pct_encode(scope)
);
for attempt in 0..2 {
match imds.get(&url).header("Metadata", "true").call() {
Ok(mut r) => {
let v: Value = r.body_mut().read_json().unwrap_or(Value::Null);
if let Some(t) = v["access_token"].as_str() {
let secs = v["expires_in"]
.as_u64()
.or_else(|| v["expires_in"].as_str().and_then(|s| s.parse().ok()))
.unwrap_or(3000);
return Ok((t.to_string(), Duration::from_secs(secs.saturating_sub(300).max(60)), true));
}
break; // reachable but no token (e.g. no identity assigned): fall through to az
}
Err(_) if attempt == 0 => std::thread::sleep(Duration::from_millis(300)),
Err(_) => {}
}
}
}
// Run az directly with arguments, never through a shell. On Windows it is az.cmd (Rust escapes
// batch-file arguments safely), and the scope was checked above anyway.
let args = ["account", "get-access-token", "--scope", &format!("{scope}/.default"), "-o", "json"];
let o = if cfg!(windows) {
std::process::Command::new("az.cmd").args(args).output().or_else(|_| std::process::Command::new("az.exe").args(args).output())
} else {
std::process::Command::new("az").args(args).output()
}
.map_err(|_| anyhow!("no managed identity here and the Azure CLI (`az`) isn't installed"))?;
if !o.status.success() {
// az's stderr can quote tokens or account details; it isn't repeated here.
bail!(
"no managed identity here, and `az account get-access-token --scope {scope}/.default` failed \
({}); run `az login`, or run that command yourself to see why",
o.status
);
}
let v: Value = serde_json::from_slice(&o.stdout).context("az returned non-JSON")?;
let t = v["accessToken"].as_str().ok_or_else(|| anyhow!("az returned no accessToken"))?;
// az tokens last about an hour; the CLI reports `expires_on` (unix seconds) in newer versions.
let ttl = v["expires_on"].as_u64().map(|e| e.saturating_sub(unix_now()).saturating_sub(300)).filter(|t| *t > 60).unwrap_or(45 * 60);
Ok((t.to_string(), Duration::from_secs(ttl), false))
}
