Add models.dev metadata-enrichment pipeline (providers P0c-2)
Provider /models listings that return bare ids (OpenAI-style) get context windows, thinking levels, image support, and display names from models.dev: kigi-models owns the transform (parse_api_json — ONE field interpretation for the bundled snapshot AND runtime refreshes), enrich_wire_model fills gaps with wire values always winning and model availability strictly wire-truth. Spec rows gained models_dev_id + wire_serves_metadata; all three current platforms are wire-served, so this pipeline is provably inert for them (byte-identical catalogs, zero egress, zero ~/.kigi writes — adversarially verified). Shell side: enrichment_fetch with a 24h disk cache guarded by binary version + keep-set + future-stamp sanity (a registry change or downgrade refetches instead of serving a catalog missing new providers), refresh of https://models.dev/api.json filtered to registry ids, KIGI_MODELS_DEV_URL override with case/whitespace-tolerant kill switch, fallback chain fresh-cache > refresh > stale-cache > bundled (each step logged). The fast path returns an empty catalog without forcing the bundled parse. From the review: blast-radius-confined parsing (one drifted provider on models.dev warn-skips instead of failing the whole refresh), registry- coverage and field-coverage tests guarding script/parser drift, a path- injectable core with 8 state-machine tests (one of which caught a guard patch that had failed to apply), _meta provenance stamp in the snapshot, and models.dev (MIT) attribution in NOTICE. Snapshot: 29 providers, 1124 models, 246KB, regenerated by scripts/gen_enrichment_snapshot.py (pure filter, no transform).
This commit is contained in:
@@ -9,6 +9,7 @@ description = "Kimi platform registry, /models wire contract, capability derivat
|
||||
kigi-env = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,382 @@
|
||||
//! models.dev metadata enrichment (multi-provider expansion).
|
||||
//!
|
||||
//! Most provider `/models` listings return bare ids — no context window, no
|
||||
//! thinking levels. This module carries per-model metadata keyed by
|
||||
//! models.dev provider id, sourced from the bundled snapshot
|
||||
//! (`enrichment_snapshot.json`, regenerated by
|
||||
//! `scripts/gen_enrichment_snapshot.py`) or a runtime refresh fetched by the
|
||||
//! shell. It NEVER invents model availability: the live listing is the only
|
||||
//! source of which models exist — enrichment only fills metadata gaps, and
|
||||
//! wire-served values always win (see [`enrich_wire_model`]).
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
/// One model's enrichment metadata. All fields optional-by-default so the
|
||||
/// snapshot stays minimal.
|
||||
#[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize)]
|
||||
pub struct EnrichmentModel {
|
||||
/// Max context window in tokens (`limit.context`).
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
pub context: u64,
|
||||
/// Max output tokens (`limit.output`). Not yet consumed by
|
||||
/// [`enrich_wire_model`]; feeds `max_completion_tokens` when provider
|
||||
/// cycles start mapping it.
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
pub output: u64,
|
||||
/// Model supports reasoning/thinking.
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub reasoning: bool,
|
||||
/// Selectable effort levels (canonical tokens, e.g. ["low","high","max"]).
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub efforts: Vec<String>,
|
||||
/// Accepts image input.
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub image_in: bool,
|
||||
/// Supports tool calling (used later to filter non-agentic listings).
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub tool_call: bool,
|
||||
/// Human display name.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
}
|
||||
|
||||
fn is_zero(v: &u64) -> bool {
|
||||
*v == 0
|
||||
}
|
||||
|
||||
/// `models.dev provider id -> model id -> metadata`.
|
||||
pub type EnrichmentCatalog = BTreeMap<String, BTreeMap<String, EnrichmentModel>>;
|
||||
|
||||
/// The bundled snapshot, embedded at compile time — RAW models.dev shape,
|
||||
/// filtered to the providers kigi references (pure-filter script:
|
||||
/// `scripts/gen_enrichment_snapshot.py`). OFFLINE FALLBACK for the runtime
|
||||
/// refresh; both parse through [`parse_api_json`] so there is exactly one
|
||||
/// field interpretation.
|
||||
pub const ENRICHMENT_SNAPSHOT_JSON: &str = include_str!("../enrichment_snapshot.json");
|
||||
|
||||
static BUNDLED: LazyLock<EnrichmentCatalog> = LazyLock::new(|| {
|
||||
// Baked-in JSON — a mismatch here is a developer error, not a runtime
|
||||
// condition (same policy as default_models.json).
|
||||
parse_api_json(ENRICHMENT_SNAPSHOT_JSON, None)
|
||||
.expect("enrichment_snapshot.json: invalid JSON (regenerate via script)")
|
||||
});
|
||||
|
||||
// ── Raw models.dev api.json shape (parse-only) ──────────────────────────────
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RawProvider {
|
||||
#[serde(default)]
|
||||
models: BTreeMap<String, RawModel>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawModel {
|
||||
name: Option<String>,
|
||||
reasoning: bool,
|
||||
reasoning_options: Vec<RawReasoningOption>,
|
||||
limit: RawLimit,
|
||||
modalities: RawModalities,
|
||||
tool_call: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawReasoningOption {
|
||||
r#type: String,
|
||||
values: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawLimit {
|
||||
context: u64,
|
||||
output: u64,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
struct RawModalities {
|
||||
input: Vec<String>,
|
||||
}
|
||||
|
||||
/// Parse a models.dev `api.json` document (full download or the bundled
|
||||
/// filtered snapshot) into the in-memory catalog. `keep`: restrict to these
|
||||
/// provider ids (runtime refresh filters the full 3MB download to the
|
||||
/// registry's providers); `None` keeps everything present.
|
||||
///
|
||||
/// Blast-radius confinement: providers are filtered on RAW keys first and
|
||||
/// typed-parsed individually — a schema drift in one of the ~70 providers
|
||||
/// kigi never keeps cannot fail the whole refresh, and a malformed KEPT
|
||||
/// provider is warn-skipped (its models fall to defaults) rather than
|
||||
/// killing the others. Only a document that isn't a JSON object errors
|
||||
/// (the caller falls back to cache/bundled — never a silently empty
|
||||
/// catalog). Keys starting with `_` (provenance stamps) are skipped.
|
||||
pub fn parse_api_json(
|
||||
json: &str,
|
||||
keep: Option<&std::collections::BTreeSet<&str>>,
|
||||
) -> Result<EnrichmentCatalog, serde_json::Error> {
|
||||
let raw: BTreeMap<String, serde_json::Value> = serde_json::from_str(json)?;
|
||||
let mut catalog = EnrichmentCatalog::new();
|
||||
for (pid, value) in raw {
|
||||
if pid.starts_with('_') {
|
||||
continue;
|
||||
}
|
||||
if let Some(keep) = keep
|
||||
&& !keep.contains(pid.as_str())
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let provider: RawProvider = match serde_json::from_value(value) {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
tracing::warn!(provider = %pid, error = %e,
|
||||
"models.dev provider entry malformed; skipping (models fall to defaults)");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let models = provider
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|(mid, m)| {
|
||||
let efforts = m
|
||||
.reasoning_options
|
||||
.into_iter()
|
||||
.find(|o| o.r#type == "effort" && !o.values.is_empty())
|
||||
.map(|o| o.values)
|
||||
.unwrap_or_default();
|
||||
let meta = EnrichmentModel {
|
||||
context: m.limit.context,
|
||||
output: m.limit.output,
|
||||
reasoning: m.reasoning,
|
||||
efforts,
|
||||
image_in: m.modalities.input.iter().any(|s| s == "image"),
|
||||
tool_call: m.tool_call,
|
||||
name: m.name,
|
||||
};
|
||||
(mid, meta)
|
||||
})
|
||||
.collect();
|
||||
catalog.insert(pid, models);
|
||||
}
|
||||
Ok(catalog)
|
||||
}
|
||||
|
||||
/// The compiled-in enrichment catalog.
|
||||
pub fn bundled_enrichment() -> &'static EnrichmentCatalog {
|
||||
&BUNDLED
|
||||
}
|
||||
|
||||
/// Look up one model's metadata. `None` when the provider or model is
|
||||
/// unknown to the catalog (callers fall through to defaults).
|
||||
pub fn lookup<'a>(
|
||||
catalog: &'a EnrichmentCatalog,
|
||||
models_dev_id: &str,
|
||||
model_id: &str,
|
||||
) -> Option<&'a EnrichmentModel> {
|
||||
catalog.get(models_dev_id)?.get(model_id)
|
||||
}
|
||||
|
||||
/// Fill metadata gaps on a live listing entry. WIRE WINS: a field the
|
||||
/// provider served is never overwritten — enrichment only supplies what the
|
||||
/// wire left absent/zero. Never changes the model id (availability stays
|
||||
/// wire-truth).
|
||||
pub fn enrich_wire_model(wire: &mut crate::WireModel, meta: &EnrichmentModel) {
|
||||
if wire.context_length == 0 && meta.context > 0 {
|
||||
wire.context_length = meta.context;
|
||||
}
|
||||
if meta.reasoning {
|
||||
wire.supports_reasoning = true;
|
||||
}
|
||||
if meta.image_in {
|
||||
wire.supports_image_in = true;
|
||||
}
|
||||
if wire.display_name.is_none() {
|
||||
wire.display_name = meta.name.clone();
|
||||
}
|
||||
if wire.think_efforts.is_none() && !meta.efforts.is_empty() {
|
||||
wire.think_efforts = Some(crate::WireThinkEfforts {
|
||||
support: true,
|
||||
valid_efforts: meta.efforts.clone(),
|
||||
// The provider's implicit default applies when the user doesn't
|
||||
// pick a level; models.dev doesn't record one.
|
||||
default_effort: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The bundled snapshot parses and covers the providers this expansion
|
||||
/// references; kimi entries cross-check against the live wire values the
|
||||
/// registry already knows (guards against a corrupted regeneration).
|
||||
#[test]
|
||||
fn bundled_snapshot_parses_and_cross_checks_kimi() {
|
||||
let catalog = bundled_enrichment();
|
||||
assert!(
|
||||
catalog.len() >= 25,
|
||||
"snapshot lost providers: {}",
|
||||
catalog.len()
|
||||
);
|
||||
let k3 = lookup(catalog, "kimi-for-coding", "k3").expect("k3 present");
|
||||
assert_eq!(k3.context, 1_048_576, "k3 context must match the live wire");
|
||||
assert_eq!(k3.efforts, ["low", "high", "max"]);
|
||||
assert!(k3.reasoning);
|
||||
let opus = lookup(catalog, "anthropic", "claude-opus-4-8").expect("opus present");
|
||||
assert_eq!(opus.context, 1_000_000);
|
||||
assert_eq!(opus.efforts, ["low", "medium", "high", "xhigh", "max"]);
|
||||
}
|
||||
|
||||
/// Wire-served fields are never overwritten; absent fields are filled.
|
||||
#[test]
|
||||
fn enrich_fills_gaps_and_never_overwrites_wire() {
|
||||
let meta = EnrichmentModel {
|
||||
context: 400_000,
|
||||
reasoning: true,
|
||||
efforts: vec!["low".into(), "high".into()],
|
||||
image_in: true,
|
||||
name: Some("GPT Test".into()),
|
||||
..Default::default()
|
||||
};
|
||||
// Bare listing entry (OpenAI-style: id only).
|
||||
let mut bare: crate::WireModel =
|
||||
serde_json::from_value(serde_json::json!({ "id": "gpt-test" })).unwrap();
|
||||
enrich_wire_model(&mut bare, &meta);
|
||||
assert_eq!(bare.context_length, 400_000);
|
||||
assert!(bare.supports_reasoning);
|
||||
assert!(bare.supports_image_in);
|
||||
assert_eq!(bare.display_name.as_deref(), Some("GPT Test"));
|
||||
let efforts = bare.think_efforts.expect("efforts filled");
|
||||
assert!(efforts.support);
|
||||
assert_eq!(efforts.valid_efforts, ["low", "high"]);
|
||||
assert_eq!(efforts.default_effort, None);
|
||||
|
||||
// Wire-served entry: nothing may change.
|
||||
let mut served: crate::WireModel = serde_json::from_value(serde_json::json!({
|
||||
"id": "gpt-test",
|
||||
"context_length": 123,
|
||||
"display_name": "Wire Name",
|
||||
"think_efforts": { "support": true, "valid_efforts": ["max"] }
|
||||
}))
|
||||
.unwrap();
|
||||
enrich_wire_model(&mut served, &meta);
|
||||
assert_eq!(served.context_length, 123, "wire context wins");
|
||||
assert_eq!(served.display_name.as_deref(), Some("Wire Name"));
|
||||
assert_eq!(
|
||||
served.think_efforts.unwrap().valid_efforts,
|
||||
["max"],
|
||||
"wire efforts win"
|
||||
);
|
||||
}
|
||||
|
||||
/// The runtime refresh parses the FULL api.json (extra fields like cost/
|
||||
/// env/doc present) and filters to the registry's provider ids.
|
||||
#[test]
|
||||
fn parse_api_json_filters_and_tolerates_unknown_fields() {
|
||||
let full = serde_json::json!({
|
||||
"openai": {
|
||||
"id": "openai", "env": ["OPENAI_API_KEY"], "doc": "https://x",
|
||||
"models": {
|
||||
"gpt-test": {
|
||||
"name": "GPT Test",
|
||||
"reasoning": true,
|
||||
"reasoning_options": [
|
||||
{"type": "effort", "values": ["low", "high", "max"]}
|
||||
],
|
||||
"limit": {"context": 400000, "output": 128000},
|
||||
"modalities": {"input": ["text", "image"], "output": ["text"]},
|
||||
"tool_call": true,
|
||||
"cost": {"input": 1.25, "output": 10}
|
||||
}
|
||||
}
|
||||
},
|
||||
"unwanted-provider": { "models": { "m": {} } }
|
||||
})
|
||||
.to_string();
|
||||
let keep: std::collections::BTreeSet<&str> = ["openai"].into();
|
||||
let catalog = parse_api_json(&full, Some(&keep)).unwrap();
|
||||
assert!(!catalog.contains_key("unwanted-provider"));
|
||||
let m = lookup(&catalog, "openai", "gpt-test").unwrap();
|
||||
assert_eq!(m.context, 400_000);
|
||||
assert_eq!(m.output, 128_000);
|
||||
assert_eq!(m.efforts, ["low", "high", "max"]);
|
||||
assert!(m.reasoning && m.image_in && m.tool_call);
|
||||
assert_eq!(m.name.as_deref(), Some("GPT Test"));
|
||||
// Malformed document errors — callers must fall back loudly, never
|
||||
// proceed with a silently empty catalog.
|
||||
assert!(parse_api_json("not json", None).is_err());
|
||||
}
|
||||
|
||||
/// Every registry `models_dev_id` must be covered by the bundled
|
||||
/// snapshot — a registry row added without updating the script's
|
||||
/// TARGETS would silently diverge bundled-vs-refresh behavior.
|
||||
#[test]
|
||||
fn bundled_snapshot_covers_every_registry_models_dev_id() {
|
||||
let catalog = bundled_enrichment();
|
||||
for platform in crate::PlatformId::ALL {
|
||||
if let Some(dev_id) = platform.models_dev_id() {
|
||||
assert!(
|
||||
catalog.contains_key(dev_id),
|
||||
"{}: models_dev_id {dev_id:?} missing from the bundled \
|
||||
snapshot — add it to scripts/gen_enrichment_snapshot.py \
|
||||
TARGETS and regenerate",
|
||||
platform.as_str(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Field-coverage guard against script/parser drift: for every field
|
||||
/// `RawModel` reads, at least one bundled model must carry a non-default
|
||||
/// value. A regeneration that dropped a MODEL_KEYS entry (e.g.
|
||||
/// `modalities`) would zero that field across the whole snapshot and
|
||||
/// fail here, instead of silently diverging from runtime refreshes.
|
||||
#[test]
|
||||
fn bundled_snapshot_carries_every_parsed_field() {
|
||||
let all: Vec<&EnrichmentModel> = bundled_enrichment()
|
||||
.values()
|
||||
.flat_map(|models| models.values())
|
||||
.collect();
|
||||
assert!(all.iter().any(|m| m.context > 0), "no context anywhere");
|
||||
assert!(all.iter().any(|m| m.output > 0), "no output anywhere");
|
||||
assert!(all.iter().any(|m| m.reasoning), "no reasoning anywhere");
|
||||
assert!(
|
||||
all.iter().any(|m| !m.efforts.is_empty()),
|
||||
"no efforts anywhere"
|
||||
);
|
||||
assert!(all.iter().any(|m| m.image_in), "no image_in anywhere");
|
||||
assert!(all.iter().any(|m| m.tool_call), "no tool_call anywhere");
|
||||
assert!(all.iter().any(|m| m.name.is_some()), "no names anywhere");
|
||||
}
|
||||
|
||||
/// One malformed provider (schema drift on models.dev) must not kill
|
||||
/// the whole refresh — kept siblings still parse; only a non-object
|
||||
/// document errors.
|
||||
#[test]
|
||||
fn malformed_provider_is_skipped_not_fatal() {
|
||||
let doc = serde_json::json!({
|
||||
"good": { "models": { "m": { "limit": {"context": 7} } } },
|
||||
"drifted": { "models": "this is not an object" },
|
||||
"_meta": { "source": "stamp, must be ignored" }
|
||||
})
|
||||
.to_string();
|
||||
let catalog = parse_api_json(&doc, None).expect("document parses");
|
||||
assert_eq!(
|
||||
lookup(&catalog, "good", "m").map(|m| m.context),
|
||||
Some(7),
|
||||
"sibling providers must survive one drifted provider"
|
||||
);
|
||||
assert!(!catalog.contains_key("drifted"));
|
||||
assert!(!catalog.contains_key("_meta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lookup_misses_are_none() {
|
||||
let catalog = bundled_enrichment();
|
||||
assert!(lookup(catalog, "no-such-provider", "x").is_none());
|
||||
assert!(lookup(catalog, "openai", "no-such-model").is_none());
|
||||
}
|
||||
}
|
||||
@@ -16,6 +16,8 @@
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
pub mod enrichment;
|
||||
|
||||
// ── Platform registry (PRD F2) ──────────────────────────────────────────────
|
||||
|
||||
/// Env var holding the moonshot-cn API key (wins over the generic name).
|
||||
@@ -82,6 +84,13 @@ struct PlatformSpec {
|
||||
console_host: Option<&'static str>,
|
||||
/// Interactive login-picker label. `None` = fall back to `display_name`.
|
||||
login_label: Option<&'static str>,
|
||||
/// This platform's provider id on models.dev, for metadata enrichment.
|
||||
/// `None` = not covered there (enrichment silently skips).
|
||||
models_dev_id: Option<&'static str>,
|
||||
/// True when the platform's `/models` listing itself serves context
|
||||
/// window / thinking metadata — enrichment (and its network refresh) is
|
||||
/// skipped entirely for such platforms.
|
||||
wire_serves_metadata: bool,
|
||||
}
|
||||
|
||||
const KIMI_CODE_SPEC: PlatformSpec = PlatformSpec {
|
||||
@@ -94,6 +103,8 @@ const KIMI_CODE_SPEC: PlatformSpec = PlatformSpec {
|
||||
vendor: "Kimi",
|
||||
console_host: None,
|
||||
login_label: None,
|
||||
models_dev_id: Some("kimi-for-coding"),
|
||||
wire_serves_metadata: true,
|
||||
};
|
||||
|
||||
const MOONSHOT_CN_SPEC: PlatformSpec = PlatformSpec {
|
||||
@@ -109,6 +120,8 @@ const MOONSHOT_CN_SPEC: PlatformSpec = PlatformSpec {
|
||||
vendor: "Moonshot",
|
||||
console_host: Some("platform.moonshot.cn"),
|
||||
login_label: Some("Moonshot Open Platform (API key \u{b7} moonshot.cn)"),
|
||||
models_dev_id: Some("moonshotai-cn"),
|
||||
wire_serves_metadata: true,
|
||||
};
|
||||
|
||||
const MOONSHOT_AI_SPEC: PlatformSpec = PlatformSpec {
|
||||
@@ -124,6 +137,8 @@ const MOONSHOT_AI_SPEC: PlatformSpec = PlatformSpec {
|
||||
vendor: "Moonshot",
|
||||
console_host: Some("platform.moonshot.ai"),
|
||||
login_label: Some("Moonshot Open Platform (API key \u{b7} moonshot.ai)"),
|
||||
models_dev_id: Some("moonshotai"),
|
||||
wire_serves_metadata: true,
|
||||
};
|
||||
|
||||
/// The platform registry. Platforms are compiled-in spec rows; there is no
|
||||
@@ -217,6 +232,17 @@ impl PlatformId {
|
||||
pub fn login_label(self) -> &'static str {
|
||||
self.spec().login_label.unwrap_or(self.spec().display_name)
|
||||
}
|
||||
|
||||
/// This platform's provider id on models.dev (metadata enrichment).
|
||||
pub fn models_dev_id(self) -> Option<&'static str> {
|
||||
self.spec().models_dev_id
|
||||
}
|
||||
|
||||
/// True when the live `/models` wire serves metadata itself — enrichment
|
||||
/// and its network refresh are skipped for such platforms.
|
||||
pub fn wire_serves_metadata(self) -> bool {
|
||||
self.spec().wire_serves_metadata
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a managed catalog key `{platform_id}/{model_id}` back into its
|
||||
|
||||
@@ -0,0 +1,417 @@
|
||||
//! models.dev enrichment catalog loading (agent-side IO).
|
||||
//!
|
||||
//! Providers whose `/models` wire serves no context/thinking metadata get it
|
||||
//! from models.dev (see `kigi_models::enrichment`). This module owns the IO:
|
||||
//! a 24h-TTL disk cache under `~/.kigi`, a runtime refresh of
|
||||
//! `https://models.dev/api.json` (filtered to registry providers before
|
||||
//! caching), and the bundled-snapshot fallback. NO NETWORK unless some
|
||||
//! enabled platform actually needs enrichment (`wire_serves_metadata` false)
|
||||
//! — with only Kimi/Moonshot configured this module never leaves disk.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use kigi_models::enrichment::{EnrichmentCatalog, bundled_enrichment, parse_api_json};
|
||||
|
||||
/// Override the refresh URL (e2e mock), or disable refresh entirely with
|
||||
/// `0`/`off` (bundled snapshot + existing cache only).
|
||||
pub(crate) const MODELS_DEV_URL_ENV: &str = "KIGI_MODELS_DEV_URL";
|
||||
const DEFAULT_MODELS_DEV_URL: &str = "https://models.dev/api.json";
|
||||
const CACHE_FILE: &str = "models_dev_cache.json";
|
||||
const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct DiskCache {
|
||||
/// Unix seconds of the successful fetch.
|
||||
fetched_at: u64,
|
||||
/// The kigi version that wrote the cache — a different binary (upgrade
|
||||
/// OR downgrade) refetches rather than trusting old filtering rules.
|
||||
#[serde(default)]
|
||||
kigi_version: String,
|
||||
/// The keep-set the catalog was filtered to. A registry change (new
|
||||
/// provider row, models_dev_id rename) invalidates the cache instead of
|
||||
/// silently serving a catalog missing the new provider for up to 24h.
|
||||
#[serde(default)]
|
||||
keep_set: Vec<String>,
|
||||
/// Already filtered + transformed catalog.
|
||||
catalog: EnrichmentCatalog,
|
||||
}
|
||||
|
||||
fn current_keep_set() -> Vec<String> {
|
||||
registry_models_dev_ids()
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether any of `platforms` needs enrichment at all.
|
||||
pub(crate) fn any_platform_needs_enrichment(platforms: &[kigi_models::PlatformId]) -> bool {
|
||||
platforms.iter().any(|p| !p.wire_serves_metadata())
|
||||
}
|
||||
|
||||
/// The registry's models.dev provider ids (the refresh filter).
|
||||
fn registry_models_dev_ids() -> BTreeSet<&'static str> {
|
||||
kigi_models::PlatformId::ALL
|
||||
.into_iter()
|
||||
.filter_map(|p| p.models_dev_id())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn cache_path() -> std::path::PathBuf {
|
||||
crate::util::kigi_home::kigi_home().join(CACHE_FILE)
|
||||
}
|
||||
|
||||
fn refresh_url() -> Option<String> {
|
||||
match std::env::var(MODELS_DEV_URL_ENV) {
|
||||
Ok(v)
|
||||
if matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"0" | "off" | "false"
|
||||
) =>
|
||||
{
|
||||
None
|
||||
}
|
||||
Ok(v) if !v.trim().is_empty() => Some(v.trim().to_string()),
|
||||
_ => Some(DEFAULT_MODELS_DEV_URL.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Load the enrichment catalog for a fetch pass over `enabled` platforms.
|
||||
///
|
||||
/// Fast path: nothing needs enrichment → empty catalog, zero IO (the merge
|
||||
/// branch is never taken for wire-served platforms, and NOT forcing the
|
||||
/// bundled parse keeps its cost/panic surface off the kimi/moonshot path).
|
||||
/// Otherwise: fresh valid disk cache → use it; else refresh over HTTP
|
||||
/// (filter + transform + best-effort cache write); on refresh failure fall
|
||||
/// back to a STALE cache, then the bundled snapshot — each step logged,
|
||||
/// never silent.
|
||||
pub(crate) fn load_enrichment_catalog(
|
||||
enabled: &[kigi_models::PlatformId],
|
||||
) -> std::borrow::Cow<'static, EnrichmentCatalog> {
|
||||
if !any_platform_needs_enrichment(enabled) {
|
||||
return std::borrow::Cow::Owned(EnrichmentCatalog::new());
|
||||
}
|
||||
load_enrichment_catalog_at(&cache_path())
|
||||
}
|
||||
|
||||
/// Path-injectable core (tests use a tempdir path directly — no env, no
|
||||
/// `kigi_home()` OnceLock interaction).
|
||||
fn load_enrichment_catalog_at(
|
||||
path: &std::path::Path,
|
||||
) -> std::borrow::Cow<'static, EnrichmentCatalog> {
|
||||
use std::borrow::Cow;
|
||||
|
||||
let cached: Option<DiskCache> =
|
||||
std::fs::read_to_string(path)
|
||||
.ok()
|
||||
.and_then(|s| match serde_json::from_str(&s) {
|
||||
Ok(c) => Some(c),
|
||||
Err(e) => {
|
||||
tracing::warn!(path = %path.display(), error = %e,
|
||||
"models.dev cache unreadable; refetching");
|
||||
None
|
||||
}
|
||||
});
|
||||
let now = now_unix();
|
||||
let cache_is_fresh = cached.as_ref().is_some_and(|c| {
|
||||
// A future fetched_at (clock jump backwards, corrupt stamp) is
|
||||
// stale, not fresh-forever; a different binary or keep-set means the
|
||||
// cache was filtered under other rules — refetch instead of serving
|
||||
// a catalog that may miss newly-registered providers for 24h.
|
||||
c.fetched_at <= now
|
||||
&& now - c.fetched_at < CACHE_TTL.as_secs()
|
||||
&& c.kigi_version == kigi_version::VERSION
|
||||
&& c.keep_set == current_keep_set()
|
||||
});
|
||||
if cache_is_fresh {
|
||||
tracing::debug!("models.dev enrichment: fresh disk cache");
|
||||
return Cow::Owned(cached.expect("cache_is_fresh implies Some").catalog);
|
||||
}
|
||||
|
||||
match refresh_url() {
|
||||
Some(url) => match fetch_and_filter(&url) {
|
||||
Ok(catalog) => {
|
||||
let cache = DiskCache {
|
||||
fetched_at: now_unix(),
|
||||
kigi_version: kigi_version::VERSION.to_string(),
|
||||
keep_set: current_keep_set(),
|
||||
catalog,
|
||||
};
|
||||
// Best-effort write: a read-only home must not fail the fetch.
|
||||
match serde_json::to_string(&cache) {
|
||||
Ok(body) => {
|
||||
if let Err(e) = crate::util::config::atomic_write_string(path, &body) {
|
||||
tracing::warn!(path = %path.display(), error = %e,
|
||||
"models.dev cache write failed; continuing in-memory");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "models.dev cache serialize failed")
|
||||
}
|
||||
}
|
||||
tracing::info!("models.dev enrichment refreshed");
|
||||
Cow::Owned(cache.catalog)
|
||||
}
|
||||
Err(e) => {
|
||||
if let Some(c) = cached {
|
||||
tracing::warn!(error = %e,
|
||||
"models.dev refresh failed; using STALE cache");
|
||||
Cow::Owned(c.catalog)
|
||||
} else {
|
||||
tracing::warn!(error = %e,
|
||||
"models.dev refresh failed; using bundled snapshot");
|
||||
Cow::Borrowed(bundled_enrichment())
|
||||
}
|
||||
}
|
||||
},
|
||||
None => {
|
||||
tracing::info!("models.dev refresh disabled; using cache/bundled");
|
||||
match cached {
|
||||
Some(c) => Cow::Owned(c.catalog),
|
||||
None => Cow::Borrowed(bundled_enrichment()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fetch_and_filter(url: &str) -> anyhow::Result<EnrichmentCatalog> {
|
||||
let response = crate::http::shared_blocking_client().get(url).send()?;
|
||||
let status = response.status();
|
||||
anyhow::ensure!(status.is_success(), "GET {url}: HTTP {}", status.as_u16());
|
||||
let body = response.text()?;
|
||||
let keep = registry_models_dev_ids();
|
||||
Ok(parse_api_json(&body, Some(&keep))?)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_test_support::EnvGuard;
|
||||
use serial_test::serial;
|
||||
|
||||
/// Wire-served-only platform sets never trigger IO — and never force the
|
||||
/// bundled parse (empty owned catalog; the merge branch is gated off).
|
||||
#[test]
|
||||
fn wire_served_platforms_get_empty_catalog_without_io() {
|
||||
assert!(!any_platform_needs_enrichment(
|
||||
&kigi_models::PlatformId::ALL
|
||||
));
|
||||
let catalog = load_enrichment_catalog(&kigi_models::PlatformId::ALL);
|
||||
assert!(catalog.is_empty());
|
||||
}
|
||||
|
||||
fn cache_file_in(dir: &tempfile::TempDir) -> std::path::PathBuf {
|
||||
dir.path().join(CACHE_FILE)
|
||||
}
|
||||
|
||||
fn write_cache(path: &std::path::Path, fetched_at: u64, versioned: bool) {
|
||||
let cache = DiskCache {
|
||||
fetched_at,
|
||||
kigi_version: if versioned {
|
||||
kigi_version::VERSION.to_string()
|
||||
} else {
|
||||
"0.0.0-other".to_string()
|
||||
},
|
||||
keep_set: current_keep_set(),
|
||||
catalog: EnrichmentCatalog::from([(
|
||||
"moonshotai".to_string(),
|
||||
std::collections::BTreeMap::from([(
|
||||
"from-cache".to_string(),
|
||||
kigi_models::enrichment::EnrichmentModel {
|
||||
context: 111,
|
||||
..Default::default()
|
||||
},
|
||||
)]),
|
||||
)]),
|
||||
};
|
||||
std::fs::write(path, serde_json::to_string(&cache).unwrap()).unwrap();
|
||||
}
|
||||
|
||||
/// Fresh valid cache short-circuits — no HTTP (no mock server mounted:
|
||||
/// a fetch attempt would fail and fall to bundled, failing the assert).
|
||||
#[test]
|
||||
#[serial]
|
||||
fn fresh_valid_cache_is_served_without_refresh() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
write_cache(&path, now_unix(), true);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "http://127.0.0.1:1/api.json");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some(),
|
||||
"fresh cache must be served"
|
||||
);
|
||||
}
|
||||
|
||||
/// Version/keep-set/future-stamp guards: each invalidates a fresh-aged
|
||||
/// cache. Refresh is disabled here, so invalidation falls through to the
|
||||
/// STALE cache (resource degradation, not data loss) — proving both the
|
||||
/// guard firing and the fallback order.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn cache_guards_invalidate_and_fall_back_to_stale() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "0");
|
||||
// Wrong binary version → not fresh → (refresh disabled) → stale used.
|
||||
write_cache(&path, now_unix(), false);
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some(),
|
||||
"stale-fallback must still serve the cached data"
|
||||
);
|
||||
// Future fetched_at → same path (guard fired: debug-log absence is
|
||||
// not observable here; the behavioral pin is refresh-disabled + the
|
||||
// wiremock test below proving a fired guard refetches).
|
||||
write_cache(&path, now_unix() + 10_000, true);
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-cache").is_some());
|
||||
}
|
||||
|
||||
/// An invalidated cache (wrong version) REFETCHES when refresh is
|
||||
/// enabled: wiremock expect(1) proves the HTTP call happened; the new
|
||||
/// cache file carries the current version + keep-set and the fetched
|
||||
/// content replaces the stale entry.
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn invalidated_cache_refetches_and_rewrites() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
|
||||
.with_test_writer()
|
||||
.try_init();
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"moonshotai": { "models": { "from-refresh": {
|
||||
"limit": {"context": 222}
|
||||
}}}
|
||||
})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
write_cache(&path, now_unix(), false);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, format!("{}/api.json", server.uri()));
|
||||
let path2 = path.clone();
|
||||
let catalog =
|
||||
tokio::task::spawn_blocking(move || load_enrichment_catalog_at(&path2).into_owned())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "moonshotai", "from-refresh").is_some(),
|
||||
"guard-invalidated cache must refetch"
|
||||
);
|
||||
let rewritten: DiskCache =
|
||||
serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap();
|
||||
assert_eq!(rewritten.kigi_version, kigi_version::VERSION);
|
||||
assert_eq!(rewritten.keep_set, current_keep_set());
|
||||
assert!(rewritten.catalog.contains_key("moonshotai"));
|
||||
}
|
||||
|
||||
/// Corrupted cache file → refetch (not a crash, not trust-garbage).
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn corrupted_cache_refetches() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({ "moonshotai": { "models": {} } })),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
std::fs::write(&path, "not json {").unwrap();
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, format!("{}/api.json", server.uri()));
|
||||
let path2 = path.clone();
|
||||
let catalog =
|
||||
tokio::task::spawn_blocking(move || load_enrichment_catalog_at(&path2).into_owned())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(catalog.contains_key("moonshotai"));
|
||||
}
|
||||
|
||||
/// Refresh failure with NO cache → bundled snapshot fallback.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn refresh_failure_without_cache_falls_back_to_bundled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, "http://127.0.0.1:1/api.json");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(
|
||||
kigi_models::enrichment::lookup(&catalog, "kimi-for-coding", "k3").is_some(),
|
||||
"bundled snapshot must back a total refresh failure"
|
||||
);
|
||||
assert!(!path.exists(), "failed refresh must not write a cache");
|
||||
}
|
||||
|
||||
/// Kill switch through the FULL load path (not just refresh_url): no
|
||||
/// cache + refresh disabled → bundled, no HTTP attempted (an attempt
|
||||
/// against the sentinel URL would be a hang/refusal, not bundled data).
|
||||
#[test]
|
||||
#[serial]
|
||||
fn kill_switch_full_path_serves_bundled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = cache_file_in(&dir);
|
||||
for token in ["0", "off", "FALSE", " Off "] {
|
||||
let _url = EnvGuard::set(MODELS_DEV_URL_ENV, token);
|
||||
assert!(refresh_url().is_none(), "token {token:?} must disable");
|
||||
let catalog = load_enrichment_catalog_at(&path);
|
||||
assert!(kigi_models::enrichment::lookup(&catalog, "kimi-for-coding", "k3").is_some());
|
||||
}
|
||||
}
|
||||
|
||||
/// Refresh path: mock server → transform runs and the keep-set filters
|
||||
/// to registry providers (`moonshotai` is a real registry models_dev id;
|
||||
/// unknown providers are dropped).
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[serial]
|
||||
async fn refresh_transforms_and_filters_to_registry_ids() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.and(wiremock::matchers::path("/api.json"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"moonshotai": { "models": { "kimi-test": {
|
||||
"limit": {"context": 262144},
|
||||
"reasoning": true,
|
||||
"reasoning_options": [
|
||||
{"type": "effort", "values": ["low", "high"]}
|
||||
]
|
||||
}}},
|
||||
"not-in-registry": { "models": { "m": {} } }
|
||||
})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let url = format!("{}/api.json", server.uri());
|
||||
let catalog = tokio::task::spawn_blocking(move || fetch_and_filter(&url))
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("fetch must succeed");
|
||||
assert!(
|
||||
!catalog.contains_key("not-in-registry"),
|
||||
"keep-set must drop providers outside the registry"
|
||||
);
|
||||
let m = kigi_models::enrichment::lookup(&catalog, "moonshotai", "kimi-test")
|
||||
.expect("registry provider survives the filter");
|
||||
assert_eq!(m.context, 262_144);
|
||||
assert_eq!(m.efforts, ["low", "high"]);
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@ pub mod auth_method;
|
||||
pub mod chat_modes;
|
||||
pub mod config;
|
||||
pub mod config_model_override_parse;
|
||||
pub(crate) mod enrichment_fetch;
|
||||
mod ext_parsers;
|
||||
pub(crate) mod feedback_client;
|
||||
pub mod folder_trust;
|
||||
|
||||
@@ -4,9 +4,11 @@
|
||||
//! (the subscription platform via the OAuth session, the open platforms via
|
||||
//! their API keys), plus the custom-endpoint OpenAI-compatible listing path.
|
||||
//!
|
||||
//! This is the sole surviving network surface relocated out of the deleted
|
||||
//! xAI-proxy backend client (`remote/`); it talks only to the configured
|
||||
//! Kimi/Moonshot model endpoints, never to a proxy backend.
|
||||
//! This is the network surface relocated out of the deleted xAI-proxy
|
||||
//! backend client (`remote/`); it talks only to the configured platform
|
||||
//! model endpoints (plus the models.dev metadata refresh when an enabled
|
||||
//! platform needs enrichment — see `enrichment_fetch`), never to a proxy
|
||||
//! backend.
|
||||
use crate::auth::KimiAuth;
|
||||
use indexmap::IndexMap;
|
||||
use serde::Deserialize;
|
||||
@@ -188,6 +190,9 @@ fn fetch_platform_models_blocking(
|
||||
let mut oauth_unauthorized = false;
|
||||
let mut successes = 0usize;
|
||||
let mut last_error: Option<BackendError> = None;
|
||||
// Loaded once per fetch pass; zero IO while every enabled platform
|
||||
// serves its own metadata (kimi/moonshot today).
|
||||
let enrichment = crate::agent::enrichment_fetch::load_enrichment_catalog(&enabled);
|
||||
for platform in &enabled {
|
||||
let bearer = if platform.uses_oauth() {
|
||||
auth.map(|a| a.key.clone())
|
||||
@@ -198,7 +203,7 @@ fn fetch_platform_models_blocking(
|
||||
.expect("enabled_platforms gated on key presence")
|
||||
.to_owned()
|
||||
};
|
||||
match fetch_one_platform_models(*platform, endpoints, &bearer) {
|
||||
match fetch_one_platform_models(*platform, endpoints, &bearer, &enrichment) {
|
||||
Ok((platform_models, platform_etag)) => {
|
||||
tracing::info!(
|
||||
platform = platform.as_str(),
|
||||
@@ -258,6 +263,7 @@ fn fetch_one_platform_models(
|
||||
platform: kigi_models::PlatformId,
|
||||
endpoints: &crate::agent::config::EndpointsConfig,
|
||||
bearer: &str,
|
||||
enrichment: &kigi_models::enrichment::EnrichmentCatalog,
|
||||
) -> Result<(Vec<crate::agent::config::ModelEntryConfig>, Option<String>), BackendError> {
|
||||
let client = crate::http::shared_blocking_client();
|
||||
let url = platform_models_url(platform, endpoints);
|
||||
@@ -294,7 +300,23 @@ fn fetch_one_platform_models(
|
||||
};
|
||||
let models = filtered
|
||||
.into_iter()
|
||||
.map(|wire| platform_wire_model_to_entry(platform, wire, &base_url))
|
||||
.map(|mut wire| {
|
||||
// Metadata-poor listings (bare ids) get context window / thinking
|
||||
// levels from the models.dev catalog; wire-served platforms skip
|
||||
// this entirely and wire values always win (enrich_wire_model).
|
||||
if !platform.wire_serves_metadata()
|
||||
&& let Some(dev_id) = platform.models_dev_id()
|
||||
{
|
||||
match kigi_models::enrichment::lookup(enrichment, dev_id, &wire.id) {
|
||||
Some(meta) => kigi_models::enrichment::enrich_wire_model(&mut wire, meta),
|
||||
None => tracing::debug!(
|
||||
platform = platform.as_str(), model = %wire.id,
|
||||
"no enrichment entry; defaults will apply"
|
||||
),
|
||||
}
|
||||
}
|
||||
platform_wire_model_to_entry(platform, wire, &base_url)
|
||||
})
|
||||
.collect();
|
||||
Ok((models, etag))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user