feat(login): custom provider logins fetch and select models

After the base URL, name, and key steps, the login now opens a model
selection screen instead of blindly saving. The listing fetch runs
concurrently with a manual-id input — a WAF challenge page or a 404 no
longer passes validation silently: a base without a /vN tail is probed
once as {base}/v1 and the working base is adopted, an anthropic-wire
listing that 401s on x-api-key retries once with Bearer (relays gate
GET /models on Bearer while /messages takes the key header), and a
failed fetch shows its reason while manual add stays usable. The chosen
ids persist as [platforms.<name>] models = [...]; catalog assembly
synthesizes selected ids the listing never served (before [model.*]
overrides so they compose) and filters unselected provider entries
(after overrides, authoritative in the provider namespace). Provider-
owned entries keep their declared wire against donor backend
inheritance, and listing errors quote line/column only — a reflected
body could carry the key.
This commit is contained in:
雷电芽衣 committed 2026-09-29 22:00:20 -04:00
1 parent 1d631a07fa
commit 9a2a9eac99
20 files changed
+1655 -68

No files matched your search

@@ -812,6 +812,9 @@ pub struct PlatformCredentialConfig {
/// `"openai"` or `"anthropic"`.
#[serde(skip_serializing_if = "Option::is_none")]
pub api: Option<String>,
/// A custom provider's login-curated model allowlist; absent = whole listing.
#[serde(skip_serializing_if = "Option::is_none")]
pub models: Option<Vec<String>>,
}
/// Resolve the API key for an API-key registry platform with injected env
@@ -2628,6 +2631,11 @@ pub(crate) fn resolve_model_list(
}
resolved = prefetched;
}
crate::agent::custom_providers::synthesize_selected(
&mut resolved,
&cfg.platforms,
platform_keys.custom(),
);
let mut explicit_api_backend_keys = std::collections::HashSet::new();
for (key, model_override) in &cfg.config_models {
let had_base = resolved.contains_key(key);
@@ -2656,8 +2664,18 @@ pub(crate) fn resolve_model_list(
);
resolved.insert(key.clone(), entry);
}
crate::agent::custom_providers::filter_unselected(
&mut resolved,
&cfg.platforms,
platform_keys.custom(),
);
{
let default_cw = DEFAULT_CONTEXT_WINDOW;
let custom_prefixes: Vec<String> = platform_keys
.custom()
.iter()
.map(|c| format!("{}/", c.provider.name))
.collect();
let donors: std::collections::HashMap<String, (std::num::NonZeroU64, ApiBackend)> =
resolved
.values()
@@ -2680,7 +2698,11 @@ pub(crate) fn resolve_model_list(
entry.info.context_window = *donor_cw;
}
// An explicit `[model.X] api_backend` wins even when it is the default.
// A provider-owned entry keeps its declared wire too.
if !explicit_api_backend_keys.contains(key.as_str())
&& !custom_prefixes
.iter()
.any(|p| key.as_str().starts_with(p.as_str()))
&& entry.info.api_backend == ApiBackend::default()
&& *donor_backend != ApiBackend::default()
{
@@ -13,32 +13,17 @@ pub(crate) fn fetch_models_blocking(
credentialed: &CredentialedProvider,
) -> Result<Vec<ModelEntryConfig>, BackendError> {
let provider = &credentialed.provider;
let url = provider.models_url();
tracing::info!(provider = %provider.name, api = provider.api.as_str(), %url, "fetching custom provider models");
let client = crate::http::shared_blocking_client();
let request = match provider.api {
CustomApi::OpenAi => client
.get(&url)
.header("Authorization", format!("Bearer {}", credentialed.key())),
CustomApi::Anthropic => client
.get(&url)
.header("x-api-key", credentialed.key())
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
};
let response = request.send()?;
if !response.status().is_success() {
let status = response.status().as_u16();
let body = response.text().unwrap_or_default();
return Err(BackendError::RequestFailed { status, body });
}
let body = response.text()?;
let listing = match provider.api {
CustomApi::OpenAi => kigi_models::parse_openai_listing(&body),
CustomApi::Anthropic => kigi_models::parse_anthropic_listing(&body),
};
let wire = listing.map_err(|e| BackendError::RequestFailed {
tracing::info!(provider = %provider.name, api = provider.api.as_str(), url = %provider.models_url(), "fetching custom provider models");
let body = listing_body(provider, credentialed.key())?;
let wire = parse_listing(provider, &body).map_err(|e| BackendError::RequestFailed {
// No error content: serde quotes the offending value, a reflected key risk.
status: 200,
body: format!("{} listing parse failed: {e}", provider.api.as_str()),
body: format!(
"{} listing parse failed at line {} column {}",
provider.api.as_str(),
e.line(),
e.column()
),
})?;
if wire.is_empty() {
tracing::warn!(provider = %provider.name, "custom provider listed no models");
@@ -46,6 +31,153 @@ pub(crate) fn fetch_models_blocking(
Ok(wire.into_iter().map(|w| entry(provider, w)).collect())
}
/// The listing GET with the dialect's key header; status-gated body.
fn listing_body(provider: &CustomProvider, key: &str) -> Result<String, BackendError> {
let url = provider.models_url();
let client = crate::http::shared_blocking_client();
let send = |bearer_fallback: bool| {
let request = match (provider.api, bearer_fallback) {
(CustomApi::OpenAi, _) => client
.get(&url)
.header("Authorization", format!("Bearer {key}")),
(CustomApi::Anthropic, false) => client
.get(&url)
.header("x-api-key", key)
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
(CustomApi::Anthropic, true) => client
.get(&url)
.header("Authorization", format!("Bearer {key}"))
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
};
request.send()
};
let mut response = send(false)?;
// Relays may gate the listing on Bearer even for the Messages wire.
if response.status().as_u16() == 401 && provider.api == CustomApi::Anthropic {
tracing::info!(provider = %provider.name, "listing 401 with x-api-key; retrying Bearer");
response = send(true)?;
}
if !response.status().is_success() {
let status = response.status().as_u16();
let body = response.text().unwrap_or_default();
return Err(BackendError::RequestFailed { status, body });
}
Ok(response.text()?)
}
fn parse_listing(
provider: &CustomProvider,
body: &str,
) -> Result<Vec<WireModel>, serde_json::Error> {
match provider.api {
CustomApi::OpenAi => kigi_models::parse_openai_listing(body),
CustomApi::Anthropic => kigi_models::parse_anthropic_listing(body),
}
}
/// The login-step listing result; `provider` carries the adopted base URL.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LoginFetch {
pub provider: CustomProvider,
pub model_ids: Vec<String>,
}
/// Login listing fetch; probes `{base}/v1` once on a 404 or unparsable body and adopts it.
pub fn fetch_listing_for_login(provider: &CustomProvider, key: &str) -> Result<LoginFetch, String> {
match list_wire(provider, key) {
Ok(model_ids) => Ok(LoginFetch {
provider: provider.clone(),
model_ids,
}),
Err(first) => {
if !first.retryable() || has_version_tail(&provider.base_url) {
return Err(first.into_message());
}
let probed = CustomProvider::new(
&provider.name,
provider.api,
&format!("{}/v1", provider.base_url),
)
.map_err(|e| e.to_string())?;
tracing::info!(
provider = %provider.name,
base_url = %probed.base_url,
"custom provider listing retrying with a /v1 base"
);
match list_wire(&probed, key) {
Ok(model_ids) => Ok(LoginFetch {
provider: probed,
model_ids,
}),
Err(second) => Err(format!(
"{}; {} also failed: {}",
first.into_message(),
probed.models_url(),
second.into_message()
)),
}
}
}
}
/// A listing failure; `retryable` marks wrong-base-URL shapes only.
enum ListFail {
Retryable(String),
Fatal(String),
}
impl ListFail {
fn retryable(&self) -> bool {
matches!(self, Self::Retryable(_))
}
fn into_message(self) -> String {
match self {
Self::Retryable(m) | Self::Fatal(m) => m,
}
}
}
fn list_wire(provider: &CustomProvider, key: &str) -> Result<Vec<String>, ListFail> {
let url = provider.models_url();
let body = listing_body(provider, key).map_err(|e| match &e {
BackendError::RequestFailed { status, .. } => {
// No body: a reflected error could echo the key.
let msg = format!("{url} answered HTTP {status}");
if *status == 404 {
ListFail::Retryable(msg)
} else {
ListFail::Fatal(msg)
}
}
BackendError::Network(_) => ListFail::Fatal(format!("Couldn't reach {url}: {e}")),
BackendError::Auth(_) => ListFail::Fatal(e.to_string()),
})?;
let wire = parse_listing(provider, &body).map_err(|e| {
ListFail::Retryable(format!(
"{url} did not serve a {} model listing (invalid JSON at line {} column {})",
provider.api.as_str(),
e.line(),
e.column()
))
})?;
Ok(wire.into_iter().map(|w| w.id).collect())
}
/// True when the base URL's last path segment is a version like `v1`.
fn has_version_tail(base_url: &str) -> bool {
let Some(tail) = base_url.rsplit('/').next() else {
return false;
};
let digits = tail.strip_prefix('v').unwrap_or("");
!digits.is_empty() && digits.chars().all(|c| c.is_ascii_digit())
}
/// A selected id the listing never served becomes an entry with wire defaults.
pub(crate) fn manual_entry(provider: &CustomProvider, id: &str) -> ModelEntryConfig {
entry(provider, WireModel::bare(id.to_owned()))
}
fn entry(provider: &CustomProvider, wire: WireModel) -> ModelEntryConfig {
let mut entry = wire_model_to_entry(
provider.managed_model_key(&wire.id),
@@ -129,6 +261,133 @@ mod tests {
.unwrap()
}
async fn login_fetch_on_thread(
provider: CustomProvider,
key: &str,
) -> Result<LoginFetch, String> {
let key = key.to_owned();
tokio::task::spawn_blocking(move || fetch_listing_for_login(&provider, &key))
.await
.unwrap()
}
#[test]
fn version_tail_detection() {
assert!(has_version_tail("https://h.example/v1"));
assert!(has_version_tail("https://h.example/api/v2"));
assert!(!has_version_tail("https://h.example"));
assert!(!has_version_tail("https://h.example/coding"));
assert!(!has_version_tail("https://h.example/v"));
assert!(!has_version_tail("https://h.example/v1x"));
}
#[tokio::test(flavor = "multi_thread")]
async fn login_fetch_returns_the_served_ids() {
let server = MockServer::start().await;
Mock::given(path("/v1/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": [{"id": "m1"}, {"id": "m2"}]
})))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/v1", server.uri());
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &base).unwrap();
let fetch = login_fetch_on_thread(provider, "sk-k").await.unwrap();
assert_eq!(fetch.provider.base_url, base);
assert_eq!(fetch.model_ids, ["m1", "m2"]);
}
#[tokio::test(flavor = "multi_thread")]
async fn login_fetch_probes_v1_on_a_404_and_adopts_the_base() {
let server = MockServer::start().await;
Mock::given(path("/coding/models"))
.respond_with(ResponseTemplate::new(404))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/coding/v1/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": [{"id": "m1"}]
})))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/coding", server.uri());
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &base).unwrap();
let fetch = login_fetch_on_thread(provider, "sk-k").await.unwrap();
assert_eq!(fetch.provider.base_url, format!("{base}/v1"));
assert_eq!(fetch.model_ids, ["m1"]);
}
#[tokio::test(flavor = "multi_thread")]
async fn login_fetch_probes_v1_when_the_200_body_is_not_a_listing() {
// A WAF answers every path with an HTML challenge page at HTTP 200.
let server = MockServer::start().await;
Mock::given(path("/models"))
.respond_with(ResponseTemplate::new(200).set_body_string("<html>challenge</html>"))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/v1/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": [{"id": "m1"}]
})))
.expect(1)
.mount(&server)
.await;
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &server.uri()).unwrap();
let fetch = login_fetch_on_thread(provider, "sk-k").await.unwrap();
assert_eq!(fetch.provider.base_url, format!("{}/v1", server.uri()));
assert_eq!(fetch.model_ids, ["m1"]);
}
#[tokio::test(flavor = "multi_thread")]
async fn login_fetch_does_not_probe_auth_failures_or_versioned_bases() {
let server = MockServer::start().await;
Mock::given(path("/models"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &server.uri()).unwrap();
let err = login_fetch_on_thread(provider, "bad").await.unwrap_err();
assert!(err.contains("401"), "{err}");
let server = MockServer::start().await;
Mock::given(path("/v2/models"))
.respond_with(ResponseTemplate::new(404))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/v2", server.uri());
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &base).unwrap();
let err = login_fetch_on_thread(provider, "sk-k").await.unwrap_err();
assert!(err.contains("404") && !err.contains("also failed"), "{err}");
}
#[tokio::test(flavor = "multi_thread")]
async fn login_fetch_reports_both_hops_when_the_probe_fails_too() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(404))
.expect(2)
.mount(&server)
.await;
let provider = CustomProvider::new("proxy", CustomApi::OpenAi, &server.uri()).unwrap();
let err = login_fetch_on_thread(provider, "sk-k").await.unwrap_err();
assert!(err.contains("also failed"), "{err}");
assert!(!err.contains("sk-k"), "{err}");
}
#[tokio::test(flavor = "multi_thread")]
async fn openai_listing_becomes_chat_completions_entries() {
let server = MockServer::start().await;
@@ -204,6 +463,54 @@ mod tests {
assert_eq!(m.auth_scheme, Some(kigi_sampler::AuthScheme::XApiKey));
}
#[tokio::test(flavor = "multi_thread")]
async fn anthropic_listing_falls_back_to_bearer_on_a_401() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/models"))
.and(header("x-api-key", "sk-gw"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/models"))
.and(header("authorization", "Bearer sk-gw"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": [{"id": "c1"}], "has_more": false
})))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/v1", server.uri());
let models = fetch_on_thread(credentialed("gw", CustomApi::Anthropic, &base, "sk-gw"))
.await
.unwrap();
assert_eq!(models.len(), 1);
assert_eq!(models[0].id.as_deref(), Some("gw/c1"));
}
#[tokio::test(flavor = "multi_thread")]
async fn openai_listing_never_retries_on_a_401() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/models"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/v1", server.uri());
let err = fetch_on_thread(credentialed("proxy", CustomApi::OpenAi, &base, "bad"))
.await
.unwrap_err();
assert!(matches!(
err,
BackendError::RequestFailed { status: 401, .. }
));
}
#[tokio::test(flavor = "multi_thread")]
async fn listing_failures_surface_status_and_parse_errors() {
let server = MockServer::start().await;
@@ -46,19 +46,32 @@ pub(crate) async fn validate_key(credentialed: &CredentialedProvider) -> Result<
let provider = &credentialed.provider;
let url = provider.models_url();
let client = crate::http::shared_client();
let request = match provider.api {
CustomApi::OpenAi => client
.get(&url)
.header("Authorization", format!("Bearer {}", credentialed.key())),
CustomApi::Anthropic => client
.get(&url)
.header("x-api-key", credentialed.key())
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
let send = |bearer_fallback: bool| {
let request = match (provider.api, bearer_fallback) {
(CustomApi::OpenAi, _) => client
.get(&url)
.header("Authorization", format!("Bearer {}", credentialed.key())),
(CustomApi::Anthropic, false) => client
.get(&url)
.header("x-api-key", credentialed.key())
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
(CustomApi::Anthropic, true) => client
.get(&url)
.header("Authorization", format!("Bearer {}", credentialed.key()))
.header("anthropic-version", kigi_sampling_types::ANTHROPIC_VERSION),
};
request.send()
};
let response = request
.send()
let mut response = send(false)
.await
.map_err(|e| format!("Couldn't reach {url}: {e}"))?;
// Mirrors listing_body: relays may gate the listing on Bearer.
if response.status().as_u16() == 401 && provider.api == CustomApi::Anthropic {
tracing::info!(provider = %provider.name, "validation 401 with x-api-key; retrying Bearer");
response = send(true)
.await
.map_err(|e| format!("Couldn't reach {url}: {e}"))?;
}
let status = response.status().as_u16();
match status {
200..=299 => Ok(()),
@@ -136,6 +149,29 @@ mod tests {
.unwrap();
}
#[tokio::test]
async fn validation_falls_back_to_bearer_for_anthropic_relays() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/models"))
.and(header("x-api-key", "sk-gw"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/models"))
.and(header("authorization", "Bearer sk-gw"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.expect(1)
.mount(&server)
.await;
let base = format!("{}/v1", server.uri());
validate_key(&credentialed(CustomApi::Anthropic, &base, "sk-gw"))
.await
.unwrap();
}
#[tokio::test]
async fn validation_names_the_failure_without_the_key() {
let server = MockServer::start().await;
@@ -7,11 +7,14 @@ use crate::agent::config::{PlatformCredentialConfig, PlatformsConfig};
mod fetch;
mod login;
mod save;
mod selection;
pub use fetch::{LoginFetch, fetch_listing_for_login};
pub(crate) use fetch::{drop_orphans, fetch_models_blocking, stamp_credentials};
pub(crate) use login::validate_key;
pub use login::{login_method_id, login_provider_name, picker_api, picker_methods};
pub use save::save_custom_provider;
pub(crate) use selection::{filter_unselected, synthesize_selected};
/// Tables setting `base_url` or `api`, valid or not.
fn declarations(platforms: &PlatformsConfig) -> Vec<(&str, Result<CustomProvider, String>)> {
@@ -4,9 +4,14 @@ use toml::Value as TomlValue;
use toml::map::Map as TomlMap;
/// Writes `[platforms.<name>]` to config.toml and the key to auth.json.
pub async fn save_custom_provider(provider: &CustomProvider, api_key: &str) -> Result<()> {
pub async fn save_custom_provider(
provider: &CustomProvider,
api_key: &str,
models: &[String],
) -> Result<()> {
let provider = provider.clone();
let api_key = api_key.to_owned();
let models = models.to_owned();
let _guard = crate::util::config::SAVE_LOCK.lock().await;
tokio::task::spawn_blocking(move || {
save_custom_provider_in(
@@ -14,6 +19,7 @@ pub async fn save_custom_provider(provider: &CustomProvider, api_key: &str) -> R
&crate::util::kigi_home::kigi_home(),
&provider,
&api_key,
&models,
)
})
.await
@@ -26,6 +32,7 @@ pub(crate) fn save_custom_provider_in(
kigi_home: &std::path::Path,
provider: &CustomProvider,
api_key: &str,
models: &[String],
) -> Result<()> {
let api_key = api_key.trim();
anyhow::ensure!(!api_key.is_empty(), "API key must not be empty");
@@ -37,7 +44,7 @@ pub(crate) fn save_custom_provider_in(
config_path.display()
)
})?;
upsert_definition(&mut root, provider)?;
upsert_definition(&mut root, provider, models)?;
let dest = crate::util::config::config_write_dest(config_path)?;
crate::util::config::atomic_write_string(&dest, &toml::to_string_pretty(&root)?)
.with_context(|| format!("writing {}", dest.display()))?;
@@ -47,12 +54,17 @@ pub(crate) fn save_custom_provider_in(
provider = %provider.name,
api = provider.api.as_str(),
base_url = %provider.base_url,
models = models.len(),
"custom provider saved"
);
Ok(())
}
fn upsert_definition(root: &mut TomlValue, provider: &CustomProvider) -> Result<()> {
fn upsert_definition(
root: &mut TomlValue,
provider: &CustomProvider,
models: &[String],
) -> Result<()> {
let top = root
.as_table_mut()
.ok_or_else(|| anyhow!("config root is not a table"))?;
@@ -68,6 +80,15 @@ fn upsert_definition(root: &mut TomlValue, provider: &CustomProvider) -> Result<
.ok_or_else(|| anyhow!("[platforms.{}] is not a table", provider.name))?;
entry.insert("api".into(), provider.api.as_str().into());
entry.insert("base_url".into(), provider.base_url.as_str().into());
entry.insert(
"models".into(),
TomlValue::Array(
models
.iter()
.map(|m| TomlValue::String(m.clone()))
.collect(),
),
);
Ok(())
}
@@ -91,7 +112,7 @@ mod tests {
.unwrap();
let p = provider("proxy", CustomApi::OpenAi, "https://h.example/v1/");
save_custom_provider_in(&config, dir.path(), &p, " sk-new ").unwrap();
save_custom_provider_in(&config, dir.path(), &p, " sk-new ", &["m1".to_owned()]).unwrap();
let saved: TomlValue = toml::from_str(&std::fs::read_to_string(&config).unwrap()).unwrap();
assert_eq!(saved["ui"]["theme"].as_str(), Some("kiginight"));
@@ -103,6 +124,10 @@ mod tests {
assert_eq!(entry["api"].as_str(), Some("openai"));
assert_eq!(entry["base_url"].as_str(), Some("https://h.example/v1"));
assert_eq!(entry["api_key"].as_str(), Some("hand-written"));
assert_eq!(
entry["models"].as_array().unwrap(),
&[TomlValue::String("m1".into())]
);
let stored = crate::auth::read_auth_json(&dir.path().join("auth.json")).unwrap();
assert_eq!(stored["proxy"].key, "sk-new");
}
@@ -116,6 +141,7 @@ mod tests {
dir.path(),
&provider("gw", CustomApi::OpenAi, "https://a.example/v1"),
"k1",
&["a1".to_owned()],
)
.unwrap();
save_custom_provider_in(
@@ -123,6 +149,7 @@ mod tests {
dir.path(),
&provider("gw", CustomApi::Anthropic, "https://b.example/v1"),
"k2",
&["b1".to_owned()],
)
.unwrap();
@@ -132,6 +159,10 @@ mod tests {
saved["platforms"]["gw"]["base_url"].as_str(),
Some("https://b.example/v1")
);
assert_eq!(
saved["platforms"]["gw"]["models"].as_array().unwrap(),
&[TomlValue::String("b1".into())]
);
let stored = crate::auth::read_auth_json(&dir.path().join("auth.json")).unwrap();
assert_eq!(stored["gw"].key, "k2");
}
@@ -142,12 +173,12 @@ mod tests {
let config = dir.path().join("config.toml");
let p = provider("proxy", CustomApi::OpenAi, "https://h.example/v1");
assert!(save_custom_provider_in(&config, dir.path(), &p, " ").is_err());
assert!(save_custom_provider_in(&config, dir.path(), &p, " ", &[]).is_err());
assert!(!config.exists());
assert!(!dir.path().join("auth.json").exists());
std::fs::write(&config, "[platforms\nnope").unwrap();
let err = save_custom_provider_in(&config, dir.path(), &p, "sk").unwrap_err();
let err = save_custom_provider_in(&config, dir.path(), &p, "sk", &[]).unwrap_err();
assert!(err.to_string().contains("unparseable"), "{err}");
assert_eq!(
std::fs::read_to_string(&config).unwrap(),
@@ -162,7 +193,7 @@ mod tests {
let config = dir.path().join("config.toml");
std::fs::write(&config, "platforms = \"x\"\n").unwrap();
let p = provider("proxy", CustomApi::OpenAi, "https://h.example/v1");
let err = save_custom_provider_in(&config, dir.path(), &p, "sk").unwrap_err();
let err = save_custom_provider_in(&config, dir.path(), &p, "sk", &[]).unwrap_err();
assert!(err.to_string().contains("not a table"), "{err}");
assert!(!dir.path().join("auth.json").exists());
}
@@ -0,0 +1,233 @@
//! The `[platforms.<name>] models` allowlist: keep the selected, synthesize the manual.
use indexmap::IndexMap;
use super::CredentialedProvider;
use crate::agent::config::{ModelEntry, PlatformsConfig};
/// Synthesizes selected ids the listing never served; runs before `[model.*]` overrides.
pub(crate) fn synthesize_selected(
resolved: &mut IndexMap<String, ModelEntry>,
platforms: &PlatformsConfig,
keyed: &[CredentialedProvider],
) {
for credentialed in keyed {
let provider = &credentialed.provider;
let Some(list) = platforms
.entries
.get(&provider.name)
.and_then(|e| e.models.as_ref())
else {
continue;
};
for id in list.iter().map(|m| m.trim()) {
if let Err(e) = kigi_models::custom::validate_model_id(id) {
tracing::warn!(provider = %provider.name, %e, "skipping an invalid selected model id");
continue;
}
let key = provider.managed_model_key(id);
if resolved.contains_key(&key) {
continue;
}
let entry = super::fetch::manual_entry(provider, id);
resolved.insert(key, ModelEntry::from_config_entry(&entry));
}
}
}
/// Drops `{name}/*` entries outside the provider's selection; runs after `[model.*]` overrides.
pub(crate) fn filter_unselected(
resolved: &mut IndexMap<String, ModelEntry>,
platforms: &PlatformsConfig,
keyed: &[CredentialedProvider],
) {
for credentialed in keyed {
let provider = &credentialed.provider;
let Some(list) = platforms
.entries
.get(&provider.name)
.and_then(|e| e.models.as_ref())
else {
continue;
};
let list: Vec<&str> = list.iter().map(|m| m.trim()).collect();
let prefix = format!("{}/", provider.name);
resolved.retain(|key, entry| {
let id = entry.info.id.as_deref().unwrap_or(key.as_str());
id.strip_prefix(&prefix)
.is_none_or(|bare| list.contains(&bare))
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use kigi_models::custom::{CustomApi, CustomProvider};
use crate::agent::models::PlatformApiKeys;
fn provider(name: &str, base: &str) -> CustomProvider {
CustomProvider::new(name, CustomApi::OpenAi, base).unwrap()
}
fn fetched_entry(provider: &CustomProvider, id: &str) -> (String, ModelEntry) {
let cfg = super::super::fetch::manual_entry(provider, id);
let key = cfg.id.clone().unwrap();
(key, ModelEntry::from_config_entry(&cfg))
}
fn keyed(provider: CustomProvider) -> Vec<CredentialedProvider> {
vec![CredentialedProvider::for_test(provider, "sk-x")]
}
#[test]
fn selection_filters_the_listing_and_synthesizes_the_manual() {
let proxy = provider("proxy", "https://p.example/v1");
let mut resolved = IndexMap::new();
for id in ["m1", "m2"] {
let (key, entry) = fetched_entry(&proxy, id);
resolved.insert(key, entry);
}
let platforms: PlatformsConfig = toml::from_str(
"[proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"m2\", \"manual-1\"]\n",
)
.unwrap();
synthesize_selected(&mut resolved, &platforms, &keyed(proxy.clone()));
filter_unselected(&mut resolved, &platforms, &keyed(proxy));
let ids: Vec<_> = resolved.keys().map(String::as_str).collect();
assert_eq!(ids, ["proxy/m2", "proxy/manual-1"]);
assert_eq!(resolved["proxy/manual-1"].info.model, "manual-1");
}
#[test]
fn a_provider_without_the_models_key_keeps_its_whole_listing() {
let proxy = provider("proxy", "https://p.example/v1");
let mut resolved = IndexMap::new();
let (key, entry) = fetched_entry(&proxy, "m1");
resolved.insert(key, entry);
let platforms: PlatformsConfig =
toml::from_str("[proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\n")
.unwrap();
synthesize_selected(&mut resolved, &platforms, &keyed(proxy.clone()));
filter_unselected(&mut resolved, &platforms, &keyed(proxy));
assert!(resolved.contains_key("proxy/m1"));
}
#[test]
fn an_empty_selection_drops_every_fetched_entry() {
let proxy = provider("proxy", "https://p.example/v1");
let mut resolved = IndexMap::new();
let (key, entry) = fetched_entry(&proxy, "m1");
resolved.insert(key, entry);
let platforms: PlatformsConfig = toml::from_str(
"[proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = []\n",
)
.unwrap();
synthesize_selected(&mut resolved, &platforms, &keyed(proxy.clone()));
filter_unselected(&mut resolved, &platforms, &keyed(proxy));
assert!(resolved.is_empty());
}
#[test]
fn invalid_selected_ids_are_skipped_with_a_warning() {
let proxy = provider("proxy", "https://p.example/v1");
let mut resolved = IndexMap::new();
let platforms: PlatformsConfig = toml::from_str(
"[proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"ok\", \"has space\"]\n",
)
.unwrap();
synthesize_selected(&mut resolved, &platforms, &keyed(proxy.clone()));
filter_unselected(&mut resolved, &platforms, &keyed(proxy));
assert!(resolved.contains_key("proxy/ok"));
assert_eq!(resolved.len(), 1);
}
#[test]
fn selection_applies_through_resolve_model_list_and_gets_the_key() {
let raw: toml::Value = toml::from_str(
"[platforms.proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"manual-1\"]\n",
)
.unwrap();
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
let proxy = provider("proxy", "https://p.example/v1");
let keys =
PlatformApiKeys::test_custom(vec![CredentialedProvider::for_test(proxy, "sk-proxy")]);
let resolved = crate::agent::config::resolve_model_list(&cfg, None, &keys);
let entry = resolved
.get("proxy/manual-1")
.expect("a manual model survives without any fetched catalog");
assert_eq!(entry.api_key.as_deref(), Some("sk-proxy"));
assert!(entry.has_own_credentials());
assert_eq!(entry.info.context_window.get(), 256_000);
}
#[test]
fn a_model_override_composes_with_the_synthesized_entry() {
let raw: toml::Value = toml::from_str(
"[platforms.proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"manual-1\"]\n\n[model.\"proxy/manual-1\"]\ncontext_window = 1000000\n",
)
.unwrap();
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
let proxy = provider("proxy", "https://p.example/v1");
let keys =
PlatformApiKeys::test_custom(vec![CredentialedProvider::for_test(proxy, "sk-proxy")]);
let resolved = crate::agent::config::resolve_model_list(&cfg, None, &keys);
let entry = &resolved["proxy/manual-1"];
assert_eq!(entry.info.context_window.get(), 1_000_000);
assert_eq!(entry.info.base_url, "https://p.example/v1");
assert_eq!(entry.api_key.as_deref(), Some("sk-proxy"));
}
#[test]
fn the_selection_is_authoritative_in_the_provider_namespace() {
// A context-only override cannot resurrect a deselected model.
let raw: toml::Value = toml::from_str(
"[platforms.proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"m2\"]\n\n[model.\"proxy/m1\"]\ncontext_window = 123456\n",
)
.unwrap();
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
let proxy = provider("proxy", "https://p.example/v1");
let (key, entry) = fetched_entry(&proxy, "m1");
let prefetched = IndexMap::from([(key, entry)]);
let keys =
PlatformApiKeys::test_custom(vec![CredentialedProvider::for_test(proxy, "sk-proxy")]);
let resolved = crate::agent::config::resolve_model_list(&cfg, Some(prefetched), &keys);
assert!(!resolved.contains_key("proxy/m1"));
assert!(resolved.contains_key("proxy/m2"));
}
#[test]
fn provider_owned_entries_keep_their_declared_wire_against_donors() {
let raw: toml::Value = toml::from_str(
"[platforms.proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\nmodels = [\"shared\"]\n\n[model.donorx]\nmodel = \"shared\"\napi_backend = \"messages\"\ncontext_window = 500000\n",
)
.unwrap();
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
let proxy = provider("proxy", "https://p.example/v1");
let keys =
PlatformApiKeys::test_custom(vec![CredentialedProvider::for_test(proxy, "sk-proxy")]);
let resolved = crate::agent::config::resolve_model_list(&cfg, None, &keys);
let entry = &resolved["proxy/shared"];
assert_eq!(
entry.info.api_backend,
crate::sampling::ApiBackend::ChatCompletions
);
}
}