feat(shell): custom providers declared in config.toml join the model catalog
[platforms.<name>] with base_url and api (openai or anthropic) declares a provider. Its key lives in auth.json under the name or in the table; /models is fetched per provider, entries are keyed <name>/<model>, and the key is stamped in memory only.
This commit is contained in:
1 parent
4514b66ffe
commit
2ca1f3baa7
11 files changed
+1048
-23
No files matched your search
@@ -178,6 +178,9 @@ pub struct EndpointsConfig {
|
||||
/// Read by `load_gcs_service_account_key_sync()`. Declared for `serde_ignored`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub gcs_service_account_key: Option<String>,
|
||||
/// Base URLs of declared custom providers; derived, never serialized.
|
||||
#[serde(skip)]
|
||||
pub custom_provider_bases: Vec<String>,
|
||||
}
|
||||
/// A blank or whitespace-only override counts as unset. Single source of truth
|
||||
/// for the "empty value = not configured" rule shared by the endpoint resolvers.
|
||||
@@ -210,7 +213,13 @@ impl EndpointsConfig {
|
||||
if let Some(endpoints) = config.get("endpoints") {
|
||||
crate::config::deep_merge_toml(&mut base, endpoints);
|
||||
}
|
||||
base.try_into().unwrap_or_default()
|
||||
let mut endpoints: Self = base.try_into().unwrap_or_default();
|
||||
if let Some(platforms) = config.get("platforms").cloned()
|
||||
&& let Ok(platforms) = platforms.try_into::<PlatformsConfig>()
|
||||
{
|
||||
endpoints.custom_provider_bases = super::custom_providers::bases(&platforms);
|
||||
}
|
||||
endpoints
|
||||
}
|
||||
/// The subscription proxy base URL through which all auxiliary services (and
|
||||
/// OAuth/session inference) resolve: explicit `coding_api_base_url`, else
|
||||
@@ -265,6 +274,7 @@ impl Default for EndpointsConfig {
|
||||
managed_config_url: env_string("KIGI_MANAGED_CONFIG_URL"),
|
||||
management_api_key: None,
|
||||
gcs_service_account_key: None,
|
||||
custom_provider_bases: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -771,8 +781,11 @@ impl PlatformsConfig {
|
||||
/// registry platform, so a typo like `moonshot_cn` fails loudly instead of
|
||||
/// silently never matching. Key values are not logged.
|
||||
pub fn warn_unknown_platforms(&self) {
|
||||
for id in self.entries.keys() {
|
||||
if kigi_models::PlatformId::parse(id).is_none() {
|
||||
super::custom_providers::warn_invalid_declarations(self);
|
||||
for (id, entry) in &self.entries {
|
||||
if kigi_models::PlatformId::parse(id).is_none()
|
||||
&& !super::custom_providers::is_declared(entry)
|
||||
{
|
||||
tracing::warn!(
|
||||
platform = %id,
|
||||
known = ?kigi_models::PlatformId::ALL
|
||||
@@ -793,6 +806,12 @@ pub struct PlatformCredentialConfig {
|
||||
/// API key for this platform. NEVER logged; never re-serialized.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub api_key: Option<String>,
|
||||
/// With `api`, declares a custom provider outside the registry.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub base_url: Option<String>,
|
||||
/// `"openai"` or `"anthropic"`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub api: Option<String>,
|
||||
}
|
||||
|
||||
/// Resolve the API key for an API-key registry platform with injected env
|
||||
@@ -1582,6 +1601,7 @@ impl Config {
|
||||
config.config_models = config_models;
|
||||
config.model_override_warnings = model_override_warnings;
|
||||
config.platforms.warn_unknown_platforms();
|
||||
config.endpoints.custom_provider_bases = super::custom_providers::bases(&config.platforms);
|
||||
if config.client_version.is_none() {
|
||||
config.client_version = Self::default().client_version;
|
||||
}
|
||||
@@ -2569,6 +2589,13 @@ pub(crate) fn resolve_model_list(
|
||||
tracing::debug!(count = defaults.len(), "loaded default models");
|
||||
resolved.extend(defaults);
|
||||
}
|
||||
let prefetched = prefetched.map(|models| {
|
||||
if cfg.endpoints.has_custom_endpoint() {
|
||||
models
|
||||
} else {
|
||||
crate::agent::custom_providers::drop_orphans(models, platform_keys.custom())
|
||||
}
|
||||
});
|
||||
if let Some(mut prefetched) = prefetched {
|
||||
tracing::debug!(count = prefetched.len(), "loaded prefetched models");
|
||||
let default_cw = DEFAULT_CONTEXT_WINDOW;
|
||||
@@ -2676,6 +2703,7 @@ pub(crate) fn resolve_model_list(
|
||||
apply_global_extra_headers(&mut resolved, &cfg.models);
|
||||
apply_global_scalar_defaults(&mut resolved, &cfg.models);
|
||||
apply_platform_credentials(&mut resolved, &cfg.platforms, platform_keys);
|
||||
crate::agent::custom_providers::stamp_credentials(&mut resolved, platform_keys.custom());
|
||||
for entry in resolved.values_mut() {
|
||||
entry.info.derive_reasoning_effort_fields();
|
||||
}
|
||||
@@ -9673,6 +9701,31 @@ default = "kigi-4.5"
|
||||
"non-BYOK config overlay must preserve bundled supported_in_api=false"
|
||||
);
|
||||
}
|
||||
/// Custom provider bases reach `EndpointsConfig` from both entry points.
|
||||
#[test]
|
||||
fn declared_custom_provider_bases_reach_the_endpoints() {
|
||||
let raw: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[platforms.proxy]
|
||||
base_url = "http://localhost:11434/v1/"
|
||||
api = "openai"
|
||||
[platforms.broken]
|
||||
api = "openai"
|
||||
[platforms.moonshot-cn]
|
||||
api_key = "sk-cn"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let cfg = Config::new_from_toml_cfg(&raw).expect("config should parse");
|
||||
assert_eq!(
|
||||
cfg.endpoints.custom_provider_bases,
|
||||
vec!["http://localhost:11434/v1".to_owned()]
|
||||
);
|
||||
assert_eq!(
|
||||
EndpointsConfig::from_config_value(&raw).custom_provider_bases,
|
||||
vec!["http://localhost:11434/v1".to_owned()]
|
||||
);
|
||||
}
|
||||
/// PRD F2: a `[platforms.<id>].api_key` from config.toml is stamped onto
|
||||
/// that platform's catalog entries (in-memory only), making them usable
|
||||
/// and API-key-visible — and only onto that platform.
|
||||
|
||||
@@ -0,0 +1,450 @@
|
||||
//! Listing fetch and key stamping for custom providers.
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use kigi_models::WireModel;
|
||||
use kigi_models::custom::{CustomApi, CustomProvider};
|
||||
|
||||
use super::CredentialedProvider;
|
||||
use crate::agent::config::{ModelEntry, ModelEntryConfig};
|
||||
use crate::agent::models_fetch::{BackendError, wire_model_to_entry};
|
||||
|
||||
/// `GET {base}/models` for one provider; every model is keyed `{name}/{id}`.
|
||||
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 {
|
||||
status: 200,
|
||||
body: format!("{} listing parse failed: {e}", provider.api.as_str()),
|
||||
})?;
|
||||
if wire.is_empty() {
|
||||
tracing::warn!(provider = %provider.name, "custom provider listed no models");
|
||||
}
|
||||
Ok(wire.into_iter().map(|w| entry(provider, w)).collect())
|
||||
}
|
||||
|
||||
fn entry(provider: &CustomProvider, wire: WireModel) -> ModelEntryConfig {
|
||||
let mut entry = wire_model_to_entry(
|
||||
provider.managed_model_key(&wire.id),
|
||||
provider.api.wire_api(),
|
||||
provider.api.key_header(),
|
||||
None,
|
||||
true,
|
||||
wire,
|
||||
&provider.base_url,
|
||||
);
|
||||
entry.description = Some(provider.name.clone());
|
||||
entry
|
||||
}
|
||||
|
||||
/// Drops fetched entries of a provider that is gone or keyless.
|
||||
pub(crate) fn drop_orphans(
|
||||
prefetched: IndexMap<String, ModelEntry>,
|
||||
keyed: &[CredentialedProvider],
|
||||
) -> IndexMap<String, ModelEntry> {
|
||||
prefetched
|
||||
.into_iter()
|
||||
.filter(|(key, entry)| {
|
||||
let id = entry.info.id.as_deref().unwrap_or(key.as_str());
|
||||
let Some((prefix, _)) = id.split_once('/') else {
|
||||
return true;
|
||||
};
|
||||
let known = kigi_models::PlatformId::parse(prefix).is_some()
|
||||
|| keyed.iter().any(|k| k.provider.name == prefix);
|
||||
if !known {
|
||||
tracing::warn!(model_key = %key, "dropping a fetched entry of an unknown or keyless provider");
|
||||
}
|
||||
known
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Stamps each entry with its provider's key, in memory only.
|
||||
pub(crate) fn stamp_credentials(
|
||||
resolved: &mut IndexMap<String, ModelEntry>,
|
||||
keyed: &[CredentialedProvider],
|
||||
) {
|
||||
for (key, entry) in resolved.iter_mut() {
|
||||
let id = entry.info.id.as_deref().unwrap_or(key.as_str());
|
||||
let Some((name, _)) = id.split_once('/') else {
|
||||
continue;
|
||||
};
|
||||
let Some(owner) = keyed.iter().find(|k| k.provider.name == name) else {
|
||||
continue;
|
||||
};
|
||||
if entry.api_key.is_some() {
|
||||
continue;
|
||||
}
|
||||
if !crate::util::matches_trusted_base_url(&entry.info.base_url, &owner.provider.base_url) {
|
||||
tracing::warn!(
|
||||
model_key = %key, provider = %name,
|
||||
"entry points away from its custom provider; key not stamped"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
entry.api_key = Some(owner.key().to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::agent::models::{ModelFetchAuth, PlatformApiKeys};
|
||||
use crate::agent::models_fetch::fetch_models_blocking as fetch_all;
|
||||
use wiremock::matchers::{header, method, path, query_param};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
fn credentialed(name: &str, api: CustomApi, base: &str, key: &str) -> CredentialedProvider {
|
||||
CredentialedProvider::for_test(CustomProvider::new(name, api, base).unwrap(), key)
|
||||
}
|
||||
|
||||
async fn fetch_on_thread(
|
||||
c: CredentialedProvider,
|
||||
) -> Result<Vec<ModelEntryConfig>, BackendError> {
|
||||
tokio::task::spawn_blocking(move || fetch_models_blocking(&c))
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn openai_listing_becomes_chat_completions_entries() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/v1/models"))
|
||||
.and(header("authorization", "Bearer sk-proxy"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"object": "list",
|
||||
"data": [{"id": "m1", "context_length": 32000}, {"id": "m2"}]
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let base = format!("{}/v1/", server.uri());
|
||||
|
||||
let models = fetch_on_thread(credentialed("proxy", CustomApi::OpenAi, &base, "sk-proxy"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let ids: Vec<_> = models.iter().map(|m| m.id.as_deref().unwrap()).collect();
|
||||
assert_eq!(ids, ["proxy/m1", "proxy/m2"]);
|
||||
assert_eq!(models[0].model, "m1");
|
||||
assert_eq!(models[0].context_window.get(), 32_000);
|
||||
assert_eq!(
|
||||
models[1].context_window.get(),
|
||||
crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW
|
||||
);
|
||||
assert_eq!(models[0].base_url, format!("{}/v1", server.uri()));
|
||||
assert_eq!(
|
||||
models[0].api_backend,
|
||||
crate::sampling::ApiBackend::ChatCompletions
|
||||
);
|
||||
assert_eq!(models[0].auth_scheme, None);
|
||||
assert_eq!(models[0].description.as_deref(), Some("proxy"));
|
||||
assert!(models[0].api_key.is_none() && models[0].env_key.is_none());
|
||||
assert!(models[0].supported_in_api);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn anthropic_listing_becomes_messages_entries() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/v1/models"))
|
||||
.and(query_param("limit", "1000"))
|
||||
.and(header("x-api-key", "sk-gw"))
|
||||
.and(header(
|
||||
"anthropic-version",
|
||||
kigi_sampling_types::ANTHROPIC_VERSION,
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"data": [{
|
||||
"id": "claude-x", "display_name": "Claude X",
|
||||
"max_input_tokens": 200000, "max_tokens": 64000
|
||||
}],
|
||||
"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);
|
||||
let m = &models[0];
|
||||
assert_eq!(m.id.as_deref(), Some("gw/claude-x"));
|
||||
assert_eq!(m.name.as_deref(), Some("Claude X"));
|
||||
assert_eq!(m.context_window.get(), 200_000);
|
||||
assert_eq!(m.max_completion_tokens, Some(64_000));
|
||||
assert_eq!(m.api_backend, crate::sampling::ApiBackend::Messages);
|
||||
assert_eq!(m.auth_scheme, Some(kigi_sampler::AuthScheme::XApiKey));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn listing_failures_surface_status_and_parse_errors() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path("/v1/models"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_string("nope"))
|
||||
.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, .. }),
|
||||
"{err}"
|
||||
);
|
||||
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path("/v1/models"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let base = format!("{}/v1", server.uri());
|
||||
let err = fetch_on_thread(credentialed("proxy", CustomApi::OpenAi, &base, "k"))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("openai listing parse failed"),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn custom_providers_alone_satisfy_the_platform_fetch() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path("/v1/models"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"data": [{"id": "m1"}]
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let base = format!("{}/v1", server.uri());
|
||||
let keys = PlatformApiKeys::test_custom(vec![credentialed(
|
||||
"proxy",
|
||||
CustomApi::OpenAi,
|
||||
&base,
|
||||
"sk-origin-secret",
|
||||
)]);
|
||||
let endpoints = crate::agent::config::EndpointsConfig::default();
|
||||
|
||||
let origin = crate::agent::models_fetch::models_fetch_origin(
|
||||
&endpoints,
|
||||
ModelFetchAuth::Platforms,
|
||||
false,
|
||||
&Default::default(),
|
||||
&keys,
|
||||
);
|
||||
assert!(origin.contains(&format!("proxy={base}/models")), "{origin}");
|
||||
assert!(
|
||||
!origin.contains("sk"),
|
||||
"the origin must never carry a key: {origin}"
|
||||
);
|
||||
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
fetch_all(
|
||||
&endpoints,
|
||||
None,
|
||||
&Default::default(),
|
||||
ModelFetchAuth::Platforms,
|
||||
&keys,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("a custom provider is a complete catalog source");
|
||||
assert_eq!(result.models.len(), 1);
|
||||
assert_eq!(result.models[0].id.as_deref(), Some("proxy/m1"));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn a_failing_custom_provider_alone_fails_the_fetch() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path("/v1/models"))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let base = format!("{}/v1", server.uri());
|
||||
let keys = PlatformApiKeys::test_custom(vec![credentialed(
|
||||
"proxy",
|
||||
CustomApi::OpenAi,
|
||||
&base,
|
||||
"sk-origin-secret",
|
||||
)]);
|
||||
let err = tokio::task::spawn_blocking(move || {
|
||||
fetch_all(
|
||||
&crate::agent::config::EndpointsConfig::default(),
|
||||
None,
|
||||
&Default::default(),
|
||||
ModelFetchAuth::Platforms,
|
||||
&keys,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.err()
|
||||
.expect("500 with no other source must fail");
|
||||
assert!(
|
||||
matches!(err, BackendError::RequestFailed { status: 500, .. }),
|
||||
"{err}"
|
||||
);
|
||||
}
|
||||
|
||||
fn resolved_entry(provider: &CustomProvider, id: &str, base: &str) -> (String, ModelEntry) {
|
||||
let wire: WireModel = serde_json::from_value(serde_json::json!({ "id": id })).unwrap();
|
||||
let mut cfg = entry(provider, wire);
|
||||
cfg.base_url = base.to_owned();
|
||||
let key = cfg.id.clone().unwrap();
|
||||
(key, ModelEntry::from_config_entry(&cfg))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stamp_gives_each_entry_only_its_own_providers_key_on_its_own_host() {
|
||||
let proxy =
|
||||
CustomProvider::new("proxy", CustomApi::OpenAi, "https://p.example/v1").unwrap();
|
||||
let gw = CustomProvider::new("gw", CustomApi::Anthropic, "https://g.example/v1").unwrap();
|
||||
let mut resolved = IndexMap::new();
|
||||
for (key, entry) in [
|
||||
resolved_entry(&proxy, "m1", "https://p.example/v1"),
|
||||
resolved_entry(&gw, "c1", "https://g.example/v1"),
|
||||
resolved_entry(&proxy, "moved", "https://evil.example/v1"),
|
||||
] {
|
||||
resolved.insert(key, entry);
|
||||
}
|
||||
let (own_key, mut own) = resolved_entry(&proxy, "own", "https://p.example/v1");
|
||||
own.api_key = Some("per-model".into());
|
||||
resolved.insert(own_key, own);
|
||||
let snapshot = vec![
|
||||
CredentialedProvider::for_test(proxy, "sk-proxy"),
|
||||
CredentialedProvider::for_test(gw, "sk-gw"),
|
||||
];
|
||||
|
||||
stamp_credentials(&mut resolved, &snapshot);
|
||||
|
||||
assert_eq!(resolved["proxy/m1"].api_key.as_deref(), Some("sk-proxy"));
|
||||
assert_eq!(resolved["gw/c1"].api_key.as_deref(), Some("sk-gw"));
|
||||
assert_eq!(
|
||||
resolved["proxy/moved"].api_key, None,
|
||||
"off-host entry gets nothing"
|
||||
);
|
||||
assert_eq!(resolved["proxy/own"].api_key.as_deref(), Some("per-model"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_model_list_stamps_fetched_custom_entries() {
|
||||
let raw: toml::Value = toml::from_str(
|
||||
"[platforms.proxy]\nbase_url = \"https://p.example/v1\"\napi = \"openai\"\n",
|
||||
)
|
||||
.unwrap();
|
||||
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
|
||||
let proxy =
|
||||
CustomProvider::new("proxy", CustomApi::OpenAi, "https://p.example/v1").unwrap();
|
||||
let (key, entry) = resolved_entry(&proxy, "m1", "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,
|
||||
Some(IndexMap::from([(key, entry)])),
|
||||
&keys,
|
||||
);
|
||||
|
||||
let m = resolved
|
||||
.get("proxy/m1")
|
||||
.expect("fetched entry survives resolution");
|
||||
assert_eq!(m.api_key.as_deref(), Some("sk-proxy"));
|
||||
assert!(m.has_own_credentials());
|
||||
assert!(m.visible_for_auth(false), "API-key users see the entry");
|
||||
assert_eq!(
|
||||
m.info.api_backend,
|
||||
crate::sampling::ApiBackend::ChatCompletions
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn orphaned_custom_entries_never_enter_the_catalog() {
|
||||
let proxy =
|
||||
CustomProvider::new("proxy", CustomApi::OpenAi, "https://p.example/v1").unwrap();
|
||||
let gone = CustomProvider::new("gone", CustomApi::OpenAi, "https://g.example/v1").unwrap();
|
||||
let mut prefetched = IndexMap::new();
|
||||
for (key, entry) in [
|
||||
resolved_entry(&proxy, "m1", "https://p.example/v1"),
|
||||
resolved_entry(&gone, "m2", "https://g.example/v1"),
|
||||
] {
|
||||
prefetched.insert(key, entry);
|
||||
}
|
||||
let registry: WireModel =
|
||||
serde_json::from_value(serde_json::json!({ "id": "kimi-k2-turbo-preview" })).unwrap();
|
||||
let registry = ModelEntry::from_config_entry(
|
||||
&crate::agent::models_fetch::platform_wire_model_to_entry(
|
||||
kigi_models::PlatformId::MoonshotCn,
|
||||
registry,
|
||||
"https://api.moonshot.cn/v1",
|
||||
),
|
||||
);
|
||||
prefetched.insert("moonshot-cn/kimi-k2-turbo-preview".into(), registry);
|
||||
let keyed = vec![CredentialedProvider::for_test(proxy, "sk-proxy")];
|
||||
let cfg = crate::agent::config::Config::default();
|
||||
|
||||
let resolved = crate::agent::config::resolve_model_list(
|
||||
&cfg,
|
||||
Some(prefetched.clone()),
|
||||
&PlatformApiKeys::test_custom(keyed.clone()),
|
||||
);
|
||||
let ids: Vec<_> = resolved.keys().map(String::as_str).collect();
|
||||
assert_eq!(ids, ["proxy/m1", "moonshot-cn/kimi-k2-turbo-preview"]);
|
||||
|
||||
let unkeyed =
|
||||
crate::agent::config::resolve_model_list(&cfg, Some(prefetched), &Default::default());
|
||||
assert!(
|
||||
!unkeyed.contains_key("proxy/m1") && !unkeyed.contains_key("gone/m2"),
|
||||
"a keyless provider leaves no entries behind"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_endpoint_mode_keeps_slashed_remote_ids() {
|
||||
let raw: toml::Value =
|
||||
toml::from_str("[endpoints]\nmodels_base_url = \"https://m.example/v1\"\n").unwrap();
|
||||
let cfg = crate::agent::config::Config::new_from_toml_cfg(&raw).unwrap();
|
||||
let other = CustomProvider::new("org", CustomApi::OpenAi, "https://o.example/v1").unwrap();
|
||||
let (key, entry) = resolved_entry(&other, "llama", "https://o.example/v1");
|
||||
|
||||
let resolved = crate::agent::config::resolve_model_list(
|
||||
&cfg,
|
||||
Some(IndexMap::from([(key, entry)])),
|
||||
&Default::default(),
|
||||
);
|
||||
|
||||
assert!(
|
||||
resolved.contains_key("org/llama"),
|
||||
"{:?}",
|
||||
resolved.keys().collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
//! `[platforms.<name>]` tables that declare custom providers.
|
||||
|
||||
use kigi_models::custom::{CustomApi, CustomProvider};
|
||||
|
||||
use crate::agent::config::{PlatformCredentialConfig, PlatformsConfig};
|
||||
|
||||
mod fetch;
|
||||
mod save;
|
||||
|
||||
pub(crate) use fetch::{drop_orphans, fetch_models_blocking, stamp_credentials};
|
||||
pub use save::save_custom_provider;
|
||||
|
||||
/// Tables setting `base_url` or `api`, valid or not.
|
||||
fn declarations(platforms: &PlatformsConfig) -> Vec<(&str, Result<CustomProvider, String>)> {
|
||||
platforms
|
||||
.entries
|
||||
.iter()
|
||||
.filter(|(_, entry)| entry.base_url.is_some() || entry.api.is_some())
|
||||
.map(|(name, entry)| (name.as_str(), parse_declaration(name, entry)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_declaration(
|
||||
name: &str,
|
||||
entry: &PlatformCredentialConfig,
|
||||
) -> Result<CustomProvider, String> {
|
||||
let api = entry
|
||||
.api
|
||||
.as_deref()
|
||||
.ok_or("api is required (\"openai\" or \"anthropic\")")?;
|
||||
let api = CustomApi::parse(api)
|
||||
.ok_or_else(|| format!("api {api:?} is not \"openai\" or \"anthropic\""))?;
|
||||
let base_url = entry.base_url.as_deref().ok_or("base_url is required")?;
|
||||
CustomProvider::new(name, api, base_url).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
/// The valid declarations.
|
||||
pub(crate) fn providers(platforms: &PlatformsConfig) -> Vec<CustomProvider> {
|
||||
declarations(platforms)
|
||||
.into_iter()
|
||||
.filter_map(|(_, parsed)| parsed.ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Base URLs of the valid declarations.
|
||||
pub(crate) fn bases(platforms: &PlatformsConfig) -> Vec<String> {
|
||||
providers(platforms)
|
||||
.into_iter()
|
||||
.map(|p| p.base_url)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// True when the table declares a provider, valid or not.
|
||||
pub(crate) fn is_declared(entry: &PlatformCredentialConfig) -> bool {
|
||||
entry.base_url.is_some() || entry.api.is_some()
|
||||
}
|
||||
|
||||
/// Warns once per refused declaration; the table stays inert.
|
||||
pub(crate) fn warn_invalid_declarations(platforms: &PlatformsConfig) {
|
||||
for (name, parsed) in declarations(platforms) {
|
||||
if let Err(reason) = parsed {
|
||||
tracing::warn!(platform = %name, %reason, "[platforms.{name}] is not a valid custom provider; ignored");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A declared provider with a key; `Debug` hides the key.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CredentialedProvider {
|
||||
pub(crate) provider: CustomProvider,
|
||||
key: String,
|
||||
}
|
||||
|
||||
impl CredentialedProvider {
|
||||
pub(crate) fn key(&self) -> &str {
|
||||
&self.key
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn for_test(provider: CustomProvider, key: &str) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
key: key.to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CredentialedProvider {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("CredentialedProvider")
|
||||
.field("provider", &self.provider)
|
||||
.field("key", &"<set>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Key order: auth.json scope, then the table's `api_key`.
|
||||
pub(crate) fn credentialed(
|
||||
platforms: &PlatformsConfig,
|
||||
stored: impl Fn(&str) -> Option<String>,
|
||||
) -> Vec<CredentialedProvider> {
|
||||
providers(platforms)
|
||||
.into_iter()
|
||||
.filter_map(|provider| {
|
||||
let key = stored(&provider.name)
|
||||
.filter(|k| !k.trim().is_empty())
|
||||
.or_else(|| {
|
||||
platforms
|
||||
.entries
|
||||
.get(&provider.name)
|
||||
.and_then(|e| e.api_key.clone())
|
||||
.filter(|k| !k.trim().is_empty())
|
||||
})?;
|
||||
Some(CredentialedProvider { provider, key })
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn platforms(toml_src: &str) -> PlatformsConfig {
|
||||
toml::from_str(toml_src).expect("platforms table parses")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declarations_read_api_and_base_url() {
|
||||
let cfg = platforms(
|
||||
r#"
|
||||
[proxy]
|
||||
base_url = "https://h.example/v1/"
|
||||
api = "openai"
|
||||
[gw]
|
||||
base_url = "https://g.example/v1"
|
||||
api = "Anthropic"
|
||||
"#,
|
||||
);
|
||||
let found = providers(&cfg);
|
||||
assert_eq!(found.len(), 2);
|
||||
assert_eq!(found[0].name, "proxy");
|
||||
assert_eq!(found[0].api, CustomApi::OpenAi);
|
||||
assert_eq!(found[0].base_url, "https://h.example/v1");
|
||||
assert_eq!(found[1].api, CustomApi::Anthropic);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_declarations_are_refused_with_a_reason() {
|
||||
let cfg = platforms(
|
||||
r#"
|
||||
[no-api]
|
||||
base_url = "https://h.example/v1"
|
||||
[no-url]
|
||||
api = "openai"
|
||||
[bad-api]
|
||||
base_url = "https://h.example/v1"
|
||||
api = "responses"
|
||||
[bad-url]
|
||||
base_url = "h.example"
|
||||
api = "openai"
|
||||
[openai]
|
||||
base_url = "https://h.example/v1"
|
||||
api = "openai"
|
||||
[Upper]
|
||||
base_url = "https://h.example/v1"
|
||||
api = "openai"
|
||||
"#,
|
||||
);
|
||||
assert!(providers(&cfg).is_empty());
|
||||
let reasons: Vec<String> = declarations(&cfg)
|
||||
.into_iter()
|
||||
.map(|(_, parsed)| parsed.expect_err("every declaration is invalid"))
|
||||
.collect();
|
||||
assert_eq!(reasons.len(), 6);
|
||||
assert!(reasons[0].contains("api is required"));
|
||||
assert!(reasons[1].contains("base_url is required"));
|
||||
assert!(reasons[2].contains("responses"));
|
||||
assert!(reasons[3].contains("base_url"));
|
||||
assert!(reasons[4].contains("built-in platform id"));
|
||||
assert!(reasons[5].contains("provider name"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_table_with_only_a_key_is_not_a_declaration() {
|
||||
let cfg = platforms("[moonshot-cn]\napi_key = \"sk-x\"\n");
|
||||
assert!(declarations(&cfg).is_empty());
|
||||
assert!(!is_declared(&cfg.entries["moonshot-cn"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_prefers_auth_json_over_config_and_blank_is_unset() {
|
||||
let cfg = platforms(
|
||||
r#"
|
||||
[proxy]
|
||||
base_url = "https://h.example/v1"
|
||||
api = "openai"
|
||||
api_key = "from-config"
|
||||
[blank]
|
||||
base_url = "https://b.example/v1"
|
||||
api = "openai"
|
||||
api_key = " "
|
||||
[keyless]
|
||||
base_url = "https://k.example/v1"
|
||||
api = "openai"
|
||||
"#,
|
||||
);
|
||||
let stored = |name: &str| (name == "proxy").then(|| "from-auth-json".to_owned());
|
||||
let found = credentialed(&cfg, stored);
|
||||
assert_eq!(found.len(), 1);
|
||||
assert_eq!(found[0].provider.name, "proxy");
|
||||
assert_eq!(found[0].key(), "from-auth-json");
|
||||
|
||||
let found = credentialed(&cfg, |_| None);
|
||||
assert_eq!(found.len(), 1);
|
||||
assert_eq!(found[0].key(), "from-config");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn debug_never_prints_the_key() {
|
||||
let cfg = platforms(
|
||||
"[proxy]\nbase_url = \"https://h.example/v1\"\napi = \"openai\"\napi_key = \"sk-secret\"\n",
|
||||
);
|
||||
let found = credentialed(&cfg, |_| None);
|
||||
assert!(!format!("{found:?}").contains("sk-secret"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use kigi_models::custom::CustomProvider;
|
||||
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<()> {
|
||||
let provider = provider.clone();
|
||||
let api_key = api_key.to_owned();
|
||||
let _guard = crate::util::config::SAVE_LOCK.lock().await;
|
||||
tokio::task::spawn_blocking(move || {
|
||||
save_custom_provider_in(
|
||||
&crate::util::config::user_config_path(),
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&provider,
|
||||
&api_key,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| anyhow!("custom provider save task: {e}"))?
|
||||
}
|
||||
|
||||
/// Config first; a keyless definition is inert.
|
||||
pub(crate) fn save_custom_provider_in(
|
||||
config_path: &std::path::Path,
|
||||
kigi_home: &std::path::Path,
|
||||
provider: &CustomProvider,
|
||||
api_key: &str,
|
||||
) -> Result<()> {
|
||||
let api_key = api_key.trim();
|
||||
anyhow::ensure!(!api_key.is_empty(), "API key must not be empty");
|
||||
let raw = crate::util::config::read_to_string_or_empty(config_path)
|
||||
.with_context(|| format!("reading {}", config_path.display()))?;
|
||||
let mut root: TomlValue = toml::from_str(&raw).map_err(|e| {
|
||||
anyhow!(
|
||||
"refusing to overwrite unparseable {}: {e}; fix the syntax error first",
|
||||
config_path.display()
|
||||
)
|
||||
})?;
|
||||
upsert_definition(&mut root, provider)?;
|
||||
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()))?;
|
||||
crate::auth::store_scoped_api_key(kigi_home, &provider.name, api_key)
|
||||
.with_context(|| format!("saving the key for {} to auth.json", provider.name))?;
|
||||
tracing::info!(
|
||||
provider = %provider.name,
|
||||
api = provider.api.as_str(),
|
||||
base_url = %provider.base_url,
|
||||
"custom provider saved"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn upsert_definition(root: &mut TomlValue, provider: &CustomProvider) -> Result<()> {
|
||||
let top = root
|
||||
.as_table_mut()
|
||||
.ok_or_else(|| anyhow!("config root is not a table"))?;
|
||||
let platforms = top
|
||||
.entry("platforms")
|
||||
.or_insert_with(|| TomlValue::Table(TomlMap::new()))
|
||||
.as_table_mut()
|
||||
.ok_or_else(|| anyhow!("[platforms] is not a table"))?;
|
||||
let entry = platforms
|
||||
.entry(provider.name.as_str())
|
||||
.or_insert_with(|| TomlValue::Table(TomlMap::new()))
|
||||
.as_table_mut()
|
||||
.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());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_models::custom::CustomApi;
|
||||
|
||||
fn provider(name: &str, api: CustomApi, url: &str) -> CustomProvider {
|
||||
CustomProvider::new(name, api, url).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_writes_definition_and_key_and_keeps_the_rest() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let config = dir.path().join("config.toml");
|
||||
std::fs::write(
|
||||
&config,
|
||||
"[ui]\ntheme = \"kiginight\"\n\n[platforms.moonshot-cn]\napi_key = \"sk-cn\"\n\n[platforms.proxy]\napi_key = \"hand-written\"\n",
|
||||
)
|
||||
.unwrap();
|
||||
let p = provider("proxy", CustomApi::OpenAi, "https://h.example/v1/");
|
||||
|
||||
save_custom_provider_in(&config, dir.path(), &p, " sk-new ").unwrap();
|
||||
|
||||
let saved: TomlValue = toml::from_str(&std::fs::read_to_string(&config).unwrap()).unwrap();
|
||||
assert_eq!(saved["ui"]["theme"].as_str(), Some("kiginight"));
|
||||
assert_eq!(
|
||||
saved["platforms"]["moonshot-cn"]["api_key"].as_str(),
|
||||
Some("sk-cn")
|
||||
);
|
||||
let entry = &saved["platforms"]["proxy"];
|
||||
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"));
|
||||
let stored = crate::auth::read_auth_json(&dir.path().join("auth.json")).unwrap();
|
||||
assert_eq!(stored["proxy"].key, "sk-new");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_creates_config_and_a_second_save_rewrites_the_definition() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let config = dir.path().join("config.toml");
|
||||
save_custom_provider_in(
|
||||
&config,
|
||||
dir.path(),
|
||||
&provider("gw", CustomApi::OpenAi, "https://a.example/v1"),
|
||||
"k1",
|
||||
)
|
||||
.unwrap();
|
||||
save_custom_provider_in(
|
||||
&config,
|
||||
dir.path(),
|
||||
&provider("gw", CustomApi::Anthropic, "https://b.example/v1"),
|
||||
"k2",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let saved: TomlValue = toml::from_str(&std::fs::read_to_string(&config).unwrap()).unwrap();
|
||||
assert_eq!(saved["platforms"]["gw"]["api"].as_str(), Some("anthropic"));
|
||||
assert_eq!(
|
||||
saved["platforms"]["gw"]["base_url"].as_str(),
|
||||
Some("https://b.example/v1")
|
||||
);
|
||||
let stored = crate::auth::read_auth_json(&dir.path().join("auth.json")).unwrap();
|
||||
assert_eq!(stored["gw"].key, "k2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_key_and_broken_config_fail_before_any_write() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
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!(!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();
|
||||
assert!(err.to_string().contains("unparseable"), "{err}");
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&config).unwrap(),
|
||||
"[platforms\nnope"
|
||||
);
|
||||
assert!(!dir.path().join("auth.json").exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_non_table_platforms_value_is_refused() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
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();
|
||||
assert!(err.to_string().contains("not a table"), "{err}");
|
||||
assert!(!dir.path().join("auth.json").exists());
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@ pub mod auth_method;
|
||||
pub mod chat_modes;
|
||||
pub mod config;
|
||||
pub mod config_model_override_parse;
|
||||
pub mod custom_providers;
|
||||
pub(crate) mod enrichment_fetch;
|
||||
mod ext_parsers;
|
||||
pub(crate) mod feedback_client;
|
||||
|
||||
@@ -73,6 +73,7 @@ enum CacheAuthMethod {
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct PlatformApiKeys {
|
||||
keys: std::collections::BTreeMap<kigi_models::PlatformId, String>,
|
||||
custom: Vec<super::custom_providers::CredentialedProvider>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for PlatformApiKeys {
|
||||
@@ -83,6 +84,7 @@ impl std::fmt::Debug for PlatformApiKeys {
|
||||
s.field(platform.as_str(), &self.keys.contains_key(&platform));
|
||||
}
|
||||
}
|
||||
s.field("custom", &self.custom);
|
||||
s.finish()
|
||||
}
|
||||
}
|
||||
@@ -114,7 +116,13 @@ impl PlatformApiKeys {
|
||||
keys.insert(platform, key);
|
||||
}
|
||||
}
|
||||
Self { keys }
|
||||
let custom = super::custom_providers::credentialed(platforms, |name| {
|
||||
stored
|
||||
.as_ref()
|
||||
.and_then(|m| m.get(name))
|
||||
.map(|a| a.key.clone())
|
||||
});
|
||||
Self { keys, custom }
|
||||
}
|
||||
|
||||
/// Resolve from the effective on-disk config (startup paths that have no
|
||||
@@ -135,7 +143,11 @@ impl PlatformApiKeys {
|
||||
/// Any API-key platform credentialed? Drives "should we prefetch without
|
||||
/// a session" and the F2 acceptance path (platform key only, no login).
|
||||
pub(crate) fn any(&self) -> bool {
|
||||
!self.keys.is_empty()
|
||||
!self.keys.is_empty() || !self.custom.is_empty()
|
||||
}
|
||||
|
||||
pub(crate) fn custom(&self) -> &[super::custom_providers::CredentialedProvider] {
|
||||
&self.custom
|
||||
}
|
||||
|
||||
/// Test-only constructor (fields are private to this module).
|
||||
@@ -148,7 +160,18 @@ impl PlatformApiKeys {
|
||||
if let Some(k) = ai {
|
||||
keys.insert(kigi_models::PlatformId::MoonshotAi, k.to_owned());
|
||||
}
|
||||
Self { keys }
|
||||
Self {
|
||||
keys,
|
||||
custom: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn test_custom(custom: Vec<super::custom_providers::CredentialedProvider>) -> Self {
|
||||
Self {
|
||||
keys: Default::default(),
|
||||
custom,
|
||||
}
|
||||
}
|
||||
|
||||
/// Test-only constructor for a single API-key platform.
|
||||
@@ -156,6 +179,7 @@ impl PlatformApiKeys {
|
||||
pub(crate) fn test_single(platform: kigi_models::PlatformId, key: &str) -> Self {
|
||||
Self {
|
||||
keys: std::collections::BTreeMap::from([(platform, key.to_owned())]),
|
||||
custom: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3698,12 +3722,14 @@ mod tests {
|
||||
"moonshot-cn".into(),
|
||||
config::PlatformCredentialConfig {
|
||||
api_key: Some("cfg-cn".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
platforms.entries.insert(
|
||||
"moonshot-ai".into(),
|
||||
config::PlatformCredentialConfig {
|
||||
api_key: Some("cfg-ai".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -114,6 +114,12 @@ pub(crate) fn models_fetch_origin(
|
||||
let parts: Vec<String> = enabled_platforms(has_oauth, oauth_tokens, platform_keys)
|
||||
.into_iter()
|
||||
.map(|p| format!("{}={}", p.as_str(), platform_models_url(p, endpoints)))
|
||||
.chain(
|
||||
platform_keys
|
||||
.custom()
|
||||
.iter()
|
||||
.map(|c| format!("{}={}", c.provider.name, c.provider.models_url())),
|
||||
)
|
||||
.collect();
|
||||
format!("platforms[{}]", parts.join(";"))
|
||||
}
|
||||
@@ -265,7 +271,7 @@ fn fetch_platform_models_blocking(
|
||||
platform_keys: &crate::agent::models::PlatformApiKeys,
|
||||
) -> Result<FetchModelsResult, BackendError> {
|
||||
let enabled = enabled_platforms(auth.is_some(), oauth_tokens, platform_keys);
|
||||
if enabled.is_empty() {
|
||||
if enabled.is_empty() && platform_keys.custom().is_empty() {
|
||||
return Err(BackendError::Auth(
|
||||
"No platform credentials: log in with `kigi login`, paste a platform API key in \
|
||||
the login screen (stored in ~/.kigi/auth.json), or set a platform env var such \
|
||||
@@ -330,6 +336,28 @@ fn fetch_platform_models_blocking(
|
||||
}
|
||||
}
|
||||
|
||||
for credentialed in platform_keys.custom() {
|
||||
match crate::agent::custom_providers::fetch_models_blocking(credentialed) {
|
||||
Ok(custom_models) => {
|
||||
tracing::info!(
|
||||
provider = %credentialed.provider.name,
|
||||
count = custom_models.len(),
|
||||
"custom provider models fetch succeeded"
|
||||
);
|
||||
successes += 1;
|
||||
models.extend(custom_models);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
provider = %credentialed.provider.name,
|
||||
error = %e,
|
||||
"custom provider models fetch failed"
|
||||
);
|
||||
last_error = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if successes == 0 {
|
||||
// All enabled platforms failed. When the failure includes an OAuth
|
||||
// 401, return `Ok` with the flag set (and no models) so the async
|
||||
@@ -613,6 +641,37 @@ pub(crate) fn platform_wire_model_to_entry(
|
||||
platform: kigi_models::PlatformId,
|
||||
wire: kigi_models::WireModel,
|
||||
base_url: &str,
|
||||
) -> crate::agent::config::ModelEntryConfig {
|
||||
let env_key = (!platform.uses_oauth())
|
||||
.then(|| crate::agent::config::EnvKeys::new(platform.api_key_env_names().iter().copied()));
|
||||
// `supported_in_api: false` hides a model unless the PRIMARY session is
|
||||
// an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code` rides
|
||||
// that primary session, so only it may be gated on it. Every other OAuth
|
||||
// platform (claude-pro-max, openai-codex, github-copilot, xai-grok) carries
|
||||
// its OWN pooled credential and its models only enter the catalog once THAT
|
||||
// provider is signed in — gating them on the Kimi session would hide every
|
||||
// model from a user who signed in with only a Claude/ChatGPT/Copilot/Grok
|
||||
// subscription.
|
||||
wire_model_to_entry(
|
||||
platform.managed_model_key(&wire.id),
|
||||
platform.wire_api(),
|
||||
platform.key_header(),
|
||||
env_key,
|
||||
platform != kigi_models::PlatformId::KimiCode,
|
||||
wire,
|
||||
base_url,
|
||||
)
|
||||
}
|
||||
|
||||
/// Wire-to-catalog mapping shared by platforms and custom providers.
|
||||
pub(crate) fn wire_model_to_entry(
|
||||
id: String,
|
||||
wire_api: kigi_models::PlatformWireApi,
|
||||
key_header: kigi_models::PlatformKeyHeader,
|
||||
env_key: Option<crate::agent::config::EnvKeys>,
|
||||
supported_in_api: bool,
|
||||
wire: kigi_models::WireModel,
|
||||
base_url: &str,
|
||||
) -> crate::agent::config::ModelEntryConfig {
|
||||
let capabilities = wire.capabilities();
|
||||
// Selectable thinking levels (live wire `think_efforts`, e.g. K3's
|
||||
@@ -626,21 +685,19 @@ pub(crate) fn platform_wire_model_to_entry(
|
||||
);
|
||||
std::num::NonZeroU64::new(DEFAULT_CONTEXT_WINDOW).expect("non-zero")
|
||||
});
|
||||
let env_key = (!platform.uses_oauth())
|
||||
.then(|| crate::agent::config::EnvKeys::new(platform.api_key_env_names().iter().copied()));
|
||||
let api_backend = match platform.wire_api() {
|
||||
let api_backend = match wire_api {
|
||||
kigi_models::PlatformWireApi::ChatCompletions => {
|
||||
crate::sampling::ApiBackend::ChatCompletions
|
||||
}
|
||||
kigi_models::PlatformWireApi::Responses => crate::sampling::ApiBackend::Responses,
|
||||
kigi_models::PlatformWireApi::Messages => crate::sampling::ApiBackend::Messages,
|
||||
};
|
||||
let auth_scheme = match platform.key_header() {
|
||||
let auth_scheme = match key_header {
|
||||
kigi_models::PlatformKeyHeader::Bearer => None,
|
||||
kigi_models::PlatformKeyHeader::XApiKey => Some(kigi_sampler::AuthScheme::XApiKey),
|
||||
};
|
||||
crate::agent::config::ModelEntryConfig {
|
||||
id: Some(platform.managed_model_key(&wire.id)),
|
||||
id: Some(id),
|
||||
name: Some(wire.display_name.clone().unwrap_or_else(|| wire.id.clone())),
|
||||
model: wire.id,
|
||||
base_url: base_url.to_owned(),
|
||||
@@ -674,15 +731,7 @@ pub(crate) fn platform_wire_model_to_entry(
|
||||
inference_idle_timeout_secs: None,
|
||||
max_retries: None,
|
||||
hidden: false,
|
||||
// `supported_in_api: false` hides a model unless the PRIMARY session is
|
||||
// an OAuth session (`ModelInfo::visible_for_auth`). Only `kimi-code`
|
||||
// rides that primary session, so only it may be gated on it. Every
|
||||
// other OAuth platform (claude-pro-max, openai-codex, github-copilot,
|
||||
// xai-grok) carries its OWN pooled credential, and its models only
|
||||
// enter the catalog once THAT provider is signed in — gating them on
|
||||
// the Kimi session would hide every model from a user who signed in
|
||||
// with only a Claude/ChatGPT/Copilot/Grok subscription.
|
||||
supported_in_api: platform != kigi_models::PlatformId::KimiCode,
|
||||
supported_in_api,
|
||||
supports_backend_search: false,
|
||||
compactions_remaining: None,
|
||||
compaction_at_tokens: None,
|
||||
|
||||
@@ -202,6 +202,15 @@ impl CredentialAuthority {
|
||||
/// Deliberately NOT [`crate::util::is_first_party_url`], which is
|
||||
/// production-only and would break every custom deployment.
|
||||
fn is_session_coding_endpoint(&self, base_url: &str) -> bool {
|
||||
// Declared custom providers are never the session's endpoint.
|
||||
if self
|
||||
.endpoints
|
||||
.custom_provider_bases
|
||||
.iter()
|
||||
.any(|base| crate::util::matches_trusted_base_url(base_url, base))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if crate::util::is_effective_coding_endpoint_url(base_url) {
|
||||
return true;
|
||||
}
|
||||
@@ -422,6 +431,38 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// LEAK guard: a declared loopback custom provider never gets the primary.
|
||||
#[test]
|
||||
fn declared_custom_provider_hosts_never_receive_the_primary() {
|
||||
let (_d, kimi) = primary("kimi-tok");
|
||||
let endpoints = EndpointsConfig {
|
||||
custom_provider_bases: vec![
|
||||
"http://localhost:11434/v1".to_owned(),
|
||||
"https://gw.example/v1".to_owned(),
|
||||
],
|
||||
..EndpointsConfig::default()
|
||||
};
|
||||
let auth = authority(endpoints, kimi);
|
||||
for url in [
|
||||
"http://localhost:11434/v1",
|
||||
"http://localhost:11434/v1/chat/completions",
|
||||
"https://gw.example/v1",
|
||||
] {
|
||||
assert_eq!(
|
||||
auth.credential_class(None, url),
|
||||
CredentialClass::None,
|
||||
"{url}"
|
||||
);
|
||||
assert!(auth.credential_for(None, url).is_none(), "{url}");
|
||||
assert!(auth.manager_for(None, url).is_none(), "{url}");
|
||||
}
|
||||
assert_eq!(
|
||||
auth.credential_class(None, "http://localhost:3000/v1"),
|
||||
CredentialClass::Primary,
|
||||
"an undeclared loopback proxy keeps the session bearer"
|
||||
);
|
||||
}
|
||||
|
||||
/// LEAK guard: every API-key registry platform, and any platform-less model
|
||||
/// on a third-party host, gets NO session credential and NO manager.
|
||||
#[test]
|
||||
|
||||
@@ -32,6 +32,7 @@ pub use manager::{AuthManager, shared_api_key_provider};
|
||||
pub use meta::AuthMeta;
|
||||
pub use model::{AuthMode, KimiAuth, lookup_auth};
|
||||
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
|
||||
pub(crate) use storage::store_scoped_api_key;
|
||||
pub use storage::{
|
||||
clear_api_key, read_api_key, read_auth_json, read_platform_api_key, read_token_by_scope,
|
||||
store_api_key, store_platform_api_key,
|
||||
|
||||
@@ -619,6 +619,15 @@ pub fn store_platform_api_key(
|
||||
),
|
||||
));
|
||||
}
|
||||
store_scoped_api_key(kigi_home, platform.as_str(), api_key)
|
||||
}
|
||||
|
||||
/// Stores a key under `scope`; callers own the name.
|
||||
pub(crate) fn store_scoped_api_key(
|
||||
kigi_home: &Path,
|
||||
scope: &str,
|
||||
api_key: &str,
|
||||
) -> std::io::Result<()> {
|
||||
let path = kigi_home.join("auth.json");
|
||||
// Serialize with the manager's cross-process auth.json writers (token
|
||||
// refresh holds the same flock): an unlocked read-modify-write here
|
||||
@@ -641,7 +650,7 @@ pub fn store_platform_api_key(
|
||||
};
|
||||
let mut map = read_auth_json_or_empty_recovering_corrupt(&path)?;
|
||||
map.insert(
|
||||
platform.as_str().to_owned(),
|
||||
scope.to_owned(),
|
||||
KimiAuth {
|
||||
key: api_key.to_owned(),
|
||||
auth_mode: AuthMode::ApiKey,
|
||||
|
||||
@@ -9,7 +9,7 @@ use toml::map::Map as TomlMap;
|
||||
///
|
||||
/// Serializes the read-modify-write in `save_config` so two rapid
|
||||
/// settings toggles can't interleave and clobber each other.
|
||||
static SAVE_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
|
||||
pub(crate) static SAVE_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
|
||||
|
||||
pub async fn save_config(config: &Config) -> Result<()> {
|
||||
let _guard = SAVE_LOCK.lock().await;
|
||||
|
||||
Reference in new issue
Block a user