diff --git a/.gitignore b/.gitignore index 59a26ed..be5032c 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,6 @@ node_modules/ # Runner-setup prompt; holds a Gitea registration token once filled in. /prompt.txt + +# Local test credentials +cred.md diff --git a/AGENTS.md b/AGENTS.md index 3959c3d..897d4f8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -485,8 +485,27 @@ client-side, no backend surface. instead of `x-api-key` (this also repaired `anthropic` and `minimax`). Login: picker rows `custom-openai` / `custom-anthropic` are not loginable (`AuthMethodKind::Unknown`); the TUI asks base URL, name (default from the - host), key, then `save_custom_provider` writes config.toml (through - `config_write_dest`) and auth.json, and `authenticate("custom:")` + host), key, then opens a MODEL-SELECTION screen (`custom_entry.rs` + `CustomSelectState`, welcome `AuthMode::CustomProviderSelect`): the listing + fetch (`fetch_listing_for_login`) runs concurrently with the manual-id input + — both visible at once, failure shows its reason — space toggles fetched + rows, enter adds a typed id or finishes (a finish pressed mid-fetch waits for + it: `finish_requested`), and the chosen ids persist as `[platforms.] + models = [...]` via `save_custom_provider`. The listing fetch probes + `{base}/v1` once when `{base}/models` 404s or answers unparsable (a WAF + challenge is HTTP 200 HTML) and the base has no `/vN` tail, adopting the + probed base for the saved definition; anthropic-wire listings that 401 on + `x-api-key` retry once with `Bearer` (relays gate `GET /models` on Bearer + while `/messages` takes `x-api-key`); listing errors never quote response + content (line/column only — a reflected body could carry the key). Catalog + assembly honors the allowlist in `selection.rs`: `synthesize_selected` (BEFORE + `[model.*]` overrides, so an override composes with a synthesized entry) + builds wire-default entries for selected ids the listing never served, and + `filter_unselected` (AFTER overrides) is authoritative in the provider's + namespace — even a `[model."name/x"]` entry for an unselected id is dropped. + Provider-owned entries never take donor `api_backend` inheritance in + `resolve_model_list` (the declared wire wins); no `models` key = the whole + listing (the declarative path). `authenticate("custom:")` then validates `GET {base}/models` (401/403 = rejected key, other non-2xx and network errors surface). Limits: context window 256K unless the listing serves `context_length` (Anthropic: `max_input_tokens`); no models.dev enrichment; a diff --git a/crates/codegen/kigi-models/src/custom.rs b/crates/codegen/kigi-models/src/custom.rs index e440318..6e18105 100644 --- a/crates/codegen/kigi-models/src/custom.rs +++ b/crates/codegen/kigi-models/src/custom.rs @@ -165,6 +165,18 @@ pub fn normalize_base_url(raw: &str) -> Result { Ok(url.to_owned()) } +/// A manually added model id: no whitespace, no controls, bounded length. +pub fn validate_model_id(id: &str) -> Result<(), CustomProviderError> { + let err = |reason: &str| Err(CustomProviderError::Name(reason.to_owned())); + if id.is_empty() || id.len() > 128 { + return err("model id must be 1 to 128 characters"); + } + if id.chars().any(|c| c.is_whitespace() || c.is_control()) { + return err("model id must not contain whitespace"); + } + Ok(()) +} + /// A valid provider name derived from a host. pub fn default_name(base_url: &str) -> String { let host = base_url diff --git a/crates/codegen/kigi-models/src/lib.rs b/crates/codegen/kigi-models/src/lib.rs index 50fadc1..7e95550 100644 --- a/crates/codegen/kigi-models/src/lib.rs +++ b/crates/codegen/kigi-models/src/lib.rs @@ -1756,6 +1756,23 @@ pub struct WireThinkEfforts { pub default_effort: Option, } +impl WireModel { + /// An id-only model, as if the listing served no metadata. + pub fn bare(id: String) -> Self { + Self { + id, + context_length: 0, + supports_reasoning: false, + supports_image_in: false, + supports_video_in: false, + display_name: None, + max_output_tokens: 0, + supports_thinking_type: None, + think_efforts: None, + } + } +} + /// `GET {base}/models` response envelope. #[derive(Debug, Clone, serde::Deserialize)] pub struct WireModelsResponse { diff --git a/crates/codegen/kigi-shell/src/agent/config.rs b/crates/codegen/kigi-shell/src/agent/config.rs index 0bf220d..f0ba17b 100644 --- a/crates/codegen/kigi-shell/src/agent/config.rs +++ b/crates/codegen/kigi-shell/src/agent/config.rs @@ -812,6 +812,9 @@ pub struct PlatformCredentialConfig { /// `"openai"` or `"anthropic"`. #[serde(skip_serializing_if = "Option::is_none")] pub api: Option, + /// A custom provider's login-curated model allowlist; absent = whole listing. + #[serde(skip_serializing_if = "Option::is_none")] + pub models: Option>, } /// 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 = platform_keys + .custom() + .iter() + .map(|c| format!("{}/", c.provider.name)) + .collect(); let donors: std::collections::HashMap = 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() { diff --git a/crates/codegen/kigi-shell/src/agent/custom_providers/fetch.rs b/crates/codegen/kigi-shell/src/agent/custom_providers/fetch.rs index 7d2abb8..14e550a 100644 --- a/crates/codegen/kigi-shell/src/agent/custom_providers/fetch.rs +++ b/crates/codegen/kigi-shell/src/agent/custom_providers/fetch.rs @@ -13,32 +13,17 @@ pub(crate) fn fetch_models_blocking( credentialed: &CredentialedProvider, ) -> Result, 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 { + 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, 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, +} + +/// 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 { + 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, 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 { + 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("challenge")) + .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; diff --git a/crates/codegen/kigi-shell/src/agent/custom_providers/login.rs b/crates/codegen/kigi-shell/src/agent/custom_providers/login.rs index f7b58ba..ba37fee 100644 --- a/crates/codegen/kigi-shell/src/agent/custom_providers/login.rs +++ b/crates/codegen/kigi-shell/src/agent/custom_providers/login.rs @@ -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; diff --git a/crates/codegen/kigi-shell/src/agent/custom_providers/mod.rs b/crates/codegen/kigi-shell/src/agent/custom_providers/mod.rs index 1c29dc2..7801be8 100644 --- a/crates/codegen/kigi-shell/src/agent/custom_providers/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/custom_providers/mod.rs @@ -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)> { diff --git a/crates/codegen/kigi-shell/src/agent/custom_providers/save.rs b/crates/codegen/kigi-shell/src/agent/custom_providers/save.rs index 4231b0b..e2a718c 100644 --- a/crates/codegen/kigi-shell/src/agent/custom_providers/save.rs +++ b/crates/codegen/kigi-shell/src/agent/custom_providers/save.rs @@ -4,9 +4,14 @@ use toml::Value as TomlValue; use toml::map::Map as TomlMap; /// Writes `[platforms.]` 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()); } diff --git a/crates/codegen/kigi-shell/src/agent/custom_providers/selection.rs b/crates/codegen/kigi-shell/src/agent/custom_providers/selection.rs new file mode 100644 index 0000000..0e3a6d6 --- /dev/null +++ b/crates/codegen/kigi-shell/src/agent/custom_providers/selection.rs @@ -0,0 +1,233 @@ +//! The `[platforms.] 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, + 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, + 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 { + 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 + ); + } +} diff --git a/crates/codegen/kigi-tui/src/app/actions.rs b/crates/codegen/kigi-tui/src/app/actions.rs index dcf908a..27d05f1 100644 --- a/crates/codegen/kigi-tui/src/app/actions.rs +++ b/crates/codegen/kigi-tui/src/app/actions.rs @@ -574,6 +574,8 @@ pub enum Action { BeginCustomProviderEntry(kigi_shell::models::custom::CustomApi), /// Enter in a custom provider step. SubmitCustomProviderInput(String), + /// Enter on an empty manual-model box: finish the model selection. + SubmitCustomProviderSelection, /// Copy the auth URL to the clipboard during authentication. CopyAuthUrl, /// Show the raw auth URL with mouse capture disabled for manual copy. @@ -1623,6 +1625,13 @@ pub enum Effect { request_seq: u64, provider: kigi_shell::models::custom::CustomProvider, key: String, + models: Vec, + }, + /// List the provider's models for the login selection screen. + FetchCustomProviderModels { + request_seq: u64, + provider: kigi_shell::models::custom::CustomProvider, + key: String, }, /// Fetch MCP server list from the shell (kigi/mcp/list). FetchMcpsList { @@ -2166,6 +2175,11 @@ pub enum TaskResult { AuthCodeSubmitted { request_seq: u64, }, + /// The login-step model listing came back (or failed). + CustomProviderModelsListed { + request_seq: u64, + result: Result, + }, /// MCP server list fetched from shell. McpsListLoaded { agent_id: AgentId, diff --git a/crates/codegen/kigi-tui/src/app/app_view.rs b/crates/codegen/kigi-tui/src/app/app_view.rs index 2a337a4..5eb0ff0 100644 --- a/crates/codegen/kigi-tui/src/app/app_view.rs +++ b/crates/codegen/kigi-tui/src/app/app_view.rs @@ -264,6 +264,8 @@ pub enum AuthMode { ApiKeyEntry(PlatformLogin), /// Custom provider login, one field per step. CustomProviderEntry(crate::app::custom_entry::CustomEntryStep), + /// Custom provider login, the model-selection screen. + CustomProviderSelect, } /// API-key platform login target, selected from the welcome picker. Wraps a /// non-OAuth registry platform ([`kigi_shell::models::PlatformId`]); @@ -924,6 +926,8 @@ pub struct AppView { pub auth_code_input: String, /// Values collected by the custom provider login steps. pub custom_entry: crate::app::custom_entry::CustomEntryDraft, + /// Model-selection state of a custom provider login; present in that mode. + pub custom_select: Option, /// Monotonically increasing sequence number for auth requests. pub next_auth_request_seq: u64, /// Every session/chat/worktree/prompt action deferred behind startup gates. @@ -1128,6 +1132,7 @@ impl AppView { auth_start_mode: AuthMode::Pending, auth_code_input: String::new(), custom_entry: Default::default(), + custom_select: None, next_auth_request_seq: 1, deferred_startup: Default::default(), auth_use_oauth: false, @@ -1663,6 +1668,7 @@ impl AppView { mid_session_login: self.auth_return_view.is_some(), auth_methods: &self.auth_methods, auth_code_input: &mut self.auth_code_input, + custom_select: &mut self.custom_select, prompt: &mut self.welcome_prompt, prompt_focused: &mut self.welcome_prompt_focused, new_worktree_dialog: &mut self.new_worktree_dialog, @@ -2188,6 +2194,7 @@ struct WelcomeInputCtx<'a> { /// the `AuthState::Pending` welcome menu. auth_methods: &'a [acp::AuthMethod], auth_code_input: &'a mut String, + custom_select: &'a mut Option, prompt: &'a mut PromptWidget, prompt_focused: &'a mut bool, new_worktree_dialog: &'a mut Option, @@ -2712,6 +2719,70 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco return InputOutcome::Changed; } } + AuthState::Authenticating { + mode: AuthMode::CustomProviderSelect, + .. + } => { + // Esc cancels back to the picker, like the entry steps. + if key!(Esc).matches(key) { + return InputOutcome::Action(Action::CancelPlatformKeyEntry); + } + if key!('q', CONTROL).matches(key) || key!('c', CONTROL).matches(key) { + if ctx.mid_session_login { + return InputOutcome::Action(Action::CancelLogin); + } + return InputOutcome::Action(Action::QuitConfirmed); + } + let Some(select) = ctx.custom_select.as_mut() else { + return InputOutcome::Unchanged; + }; + if key!(Up).matches(key) { + select.move_highlight(-1); + return InputOutcome::Changed; + } + if key!(Down).matches(key) { + select.move_highlight(1); + return InputOutcome::Changed; + } + // Model ids hold no spaces, so the spacebar toggles rows. + if key!(Char(' ')).matches(key) { + select.toggle_highlighted(); + return InputOutcome::Changed; + } + if crate::input::key::is_paste_key(key) + || crate::input::key::is_inline_paste_key(key) + { + if let Some(text) = crate::clipboard::system_clipboard_get() { + let cleaned: String = + text.chars().filter(|c| *c != '\n' && *c != '\r').collect(); + ctx.auth_code_input.push_str(&cleaned); + return InputOutcome::Changed; + } + crate::clipboard::log_paste_key_empty_host_clipboard("auth input"); + return InputOutcome::Unchanged; + } + if key!(Enter).matches(key) { + let trimmed = ctx.auth_code_input.trim().to_string(); + if trimmed.is_empty() { + return InputOutcome::Action(Action::SubmitCustomProviderSelection); + } + select.add_manual(&trimmed); + if select.error.is_none() { + ctx.auth_code_input.clear(); + } + return InputOutcome::Changed; + } + if key!(Backspace).matches(key) { + ctx.auth_code_input.pop(); + return InputOutcome::Changed; + } + if let crossterm::event::KeyCode::Char(c) = key.code + && crate::input::key::is_text_input_key(key) + { + ctx.auth_code_input.push(c); + return InputOutcome::Changed; + } + } AuthState::Authenticating { mode: AuthMode::Loopback, .. @@ -2775,7 +2846,10 @@ fn handle_welcome_input(ev: &Event, ctx: &mut WelcomeInputCtx<'_>) -> InputOutco } AuthState::Authenticating { mode: - AuthMode::Loopback | AuthMode::ApiKeyEntry(_) | AuthMode::CustomProviderEntry(_), + AuthMode::Loopback + | AuthMode::ApiKeyEntry(_) + | AuthMode::CustomProviderEntry(_) + | AuthMode::CustomProviderSelect, .. } => { let cleaned: String = text.chars().filter(|c| *c != '\n' && *c != '\r').collect(); @@ -3276,6 +3350,7 @@ impl AppView { login_label: self.login_label.as_deref(), auth_code_input: &self.auth_code_input, custom_entry: &self.custom_entry, + custom_select: self.custom_select.as_ref(), clipboard_copied: self.auth_clipboard_copied, show_raw_url: self.auth_show_raw_url, tip, @@ -4412,6 +4487,7 @@ pub(crate) mod tests { auth_start_mode: AuthMode::Pending, auth_code_input: String::new(), custom_entry: Default::default(), + custom_select: None, next_auth_request_seq: 1, deferred_startup: Default::default(), auth_use_oauth: false, @@ -7632,12 +7708,24 @@ pub(crate) mod tests { /// Ctrl+V / Super+V paste the host clipboard into every auth input. #[test] fn auth_inputs_paste_key_chords_insert_clipboard_text() { + let select_state = || { + crate::app::custom_entry::CustomSelectState::new( + kigi_shell::models::custom::CustomProvider::new( + "gw", + kigi_shell::models::custom::CustomApi::OpenAi, + "https://h.example/v1", + ) + .unwrap(), + "sk".to_owned(), + ) + }; let modes = [ AuthMode::Loopback, AuthMode::ApiKeyEntry(PlatformLogin(kigi_shell::models::PlatformId::MoonshotCn)), AuthMode::CustomProviderEntry(crate::app::custom_entry::CustomEntryStep::first( kigi_shell::models::custom::CustomApi::OpenAi, )), + AuthMode::CustomProviderSelect, ]; for mode in modes { for mods in [KeyModifiers::CONTROL, KeyModifiers::SUPER] { @@ -7648,12 +7736,13 @@ pub(crate) mod tests { auth_url: None, mode, }; - crate::clipboard::set_clipboard_probe_hook( - crate::clipboard::ClipboardProbeHook { - text: Some("sk-pasted\nkey".to_owned()), - ..Default::default() - }, - ); + if matches!(mode, AuthMode::CustomProviderSelect) { + app.custom_select = Some(select_state()); + } + crate::clipboard::set_clipboard_probe_hook(crate::clipboard::ClipboardProbeHook { + text: Some("sk-pasted\nkey".to_owned()), + ..Default::default() + }); let outcome = app.handle_input(&key_event(KeyCode::Char('v'), mods)); crate::clipboard::clear_clipboard_probe_hook(); assert!( diff --git a/crates/codegen/kigi-tui/src/app/custom_entry.rs b/crates/codegen/kigi-tui/src/app/custom_entry.rs index cbf877d..cbc3e06 100644 --- a/crates/codegen/kigi-tui/src/app/custom_entry.rs +++ b/crates/codegen/kigi-tui/src/app/custom_entry.rs @@ -1,7 +1,7 @@ //! Custom provider login steps: base URL, name, key. use kigi_shell::models::custom::{ - CustomApi, CustomProvider, default_name, normalize_base_url, validate_name, + CustomApi, CustomProvider, default_name, normalize_base_url, validate_model_id, validate_name, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -141,6 +141,141 @@ impl CustomEntryDraft { } } +/// Fetch status of the model-selection screen. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SelectFetch { + InFlight, + Listed, + Failed(String), +} + +/// One selectable model row; `manual` marks user-typed ids. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SelectRow { + pub id: String, + pub selected: bool, + pub manual: bool, +} + +/// Model selection after the key step: fetched rows plus manual adds. +#[derive(Clone)] +pub struct CustomSelectState { + /// The fetch may adopt a `/v1` base URL. + pub provider: CustomProvider, + key: String, + pub rows: Vec, + pub highlight: usize, + pub fetch: SelectFetch, + pub error: Option, + /// Enter hit while the listing is in flight; it completes on arrival. + pub finish_requested: bool, +} + +impl CustomSelectState { + pub fn new(provider: CustomProvider, key: String) -> Self { + Self { + provider, + key, + rows: Vec::new(), + highlight: 0, + fetch: SelectFetch::InFlight, + error: None, + finish_requested: false, + } + } + + pub fn take_key(self) -> String { + self.key + } + + pub fn apply_fetch(&mut self, fetch: kigi_shell::agent::custom_providers::LoginFetch) { + self.provider = fetch.provider; + for id in fetch.model_ids { + if self.rows.iter().any(|r| r.id == id) { + continue; + } + self.rows.push(SelectRow { + id, + selected: false, + manual: false, + }); + } + self.fetch = SelectFetch::Listed; + } + + pub fn fail_fetch(&mut self, reason: String) { + self.fetch = SelectFetch::Failed(reason); + } + + pub fn move_highlight(&mut self, delta: isize) { + if self.rows.is_empty() { + return; + } + let max = self.rows.len() - 1; + self.highlight = self.highlight.saturating_add_signed(delta).min(max); + } + + pub fn toggle_highlighted(&mut self) { + if let Some(row) = self.rows.get_mut(self.highlight) { + row.selected = !row.selected; + } + } + + /// Adds a typed id selected; an id already listed just gets selected. + pub fn add_manual(&mut self, raw: &str) { + let id = raw.trim(); + if let Err(e) = validate_model_id(id) { + self.error = Some(e.to_string()); + return; + } + if let Some(row) = self.rows.iter_mut().find(|r| r.id == id) { + row.selected = true; + } else { + self.rows.push(SelectRow { + id: id.to_owned(), + selected: true, + manual: true, + }); + } + self.error = None; + } + + /// The chosen ids, or a refusal when nothing is chosen. + pub fn try_finish(&mut self) -> Option> { + // Wait for the listing: it may adopt a probed `/v1` base URL. + if self.fetch == SelectFetch::InFlight { + self.finish_requested = true; + return None; + } + let ids: Vec = self + .rows + .iter() + .filter(|r| r.selected) + .map(|r| r.id.clone()) + .collect(); + if ids.is_empty() { + self.finish_requested = false; + self.error = Some("Select at least one model, or type one and press enter".to_owned()); + return None; + } + Some(ids) + } +} + +impl std::fmt::Debug for CustomSelectState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CustomSelectState") + .field("provider", &self.provider) + .field("key", &"") + .field("rows", &self.rows) + .field("highlight", &self.highlight) + .field("fetch", &self.fetch) + .field("error", &self.error) + .field("finish_requested", &self.finish_requested) + .finish() + } +} + #[cfg(test)] mod tests { use super::*; @@ -236,4 +371,69 @@ mod tests { assert!(!step(CustomEntryField::BaseUrl).masks_input()); assert!(!step(CustomEntryField::Name).masks_input()); } + + fn select_state() -> CustomSelectState { + let provider = + CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/v1").unwrap(); + CustomSelectState::new(provider, "sk-secret".to_owned()) + } + + #[test] + fn select_state_toggles_adds_and_finishes() { + let mut state = select_state(); + assert!(state.try_finish().is_none(), "in flight: finish waits"); + assert!(state.finish_requested && state.error.is_none()); + state.finish_requested = false; + + state.apply_fetch(kigi_shell::agent::custom_providers::LoginFetch { + provider: state.provider.clone(), + model_ids: vec!["m1".to_owned(), "m2".to_owned()], + }); + assert!(state.try_finish().is_none()); + assert!(state.error.is_some(), "listed: an empty selection refuses"); + assert_eq!(state.rows.len(), 2); + + state.toggle_highlighted(); + state.move_highlight(1); + state.toggle_highlighted(); + state.move_highlight(9); + assert_eq!(state.highlight, 1, "highlight clamps at the last row"); + + state.add_manual(" claude-opus-5-5[1m] "); + assert_eq!(state.rows.len(), 3); + state.add_manual("m1"); + assert_eq!( + state.rows.len(), + 3, + "a listed id is selected, not duplicated" + ); + state.add_manual("has space"); + assert!(state.error.is_some()); + assert_eq!(state.rows.len(), 3); + + let ids = state.try_finish().expect("three selected"); + assert_eq!(ids, ["m1", "m2", "claude-opus-5-5[1m]"]); + } + + #[test] + fn select_state_apply_fetch_adopts_the_probed_base_and_dedups() { + let mut state = select_state(); + state.add_manual("m1"); + let adopted = + CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/coding/v1").unwrap(); + state.apply_fetch(kigi_shell::agent::custom_providers::LoginFetch { + provider: adopted, + model_ids: vec!["m1".to_owned(), "m2".to_owned()], + }); + assert_eq!(state.provider.base_url, "https://h.example/coding/v1"); + assert_eq!(state.rows.len(), 2, "the manual row covers the fetched id"); + assert!(state.rows[0].manual && state.rows[0].selected); + assert!(!state.rows[1].selected); + } + + #[test] + fn select_state_debug_hides_the_key() { + let state = select_state(); + assert!(!format!("{state:?}").contains("sk-secret")); + } } diff --git a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs index c6fc8c4..3ac8fa3 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/auth.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/auth.rs @@ -282,6 +282,7 @@ pub(super) fn dispatch_cancel_login(app: &mut AppView) -> Vec { app.auth_state = AuthState::Done; app.auth_show_raw_url = false; app.auth_code_input.clear(); + app.custom_select = None; app.auth_in_flight_method = None; restore_auth_return_view(app, return_view); // The user bailed out of re-auth — drop stashed prompts and strip the @@ -335,6 +336,7 @@ pub(super) fn dispatch_begin_custom_provider_entry( app.next_auth_request_seq += 1; app.auth_code_input.clear(); app.custom_entry.clear(); + app.custom_select = None; app.auth_state = AuthState::Authenticating { request_seq, handle: None, @@ -373,18 +375,20 @@ pub(super) fn dispatch_submit_custom_provider_input( vec![] } CustomEntryOutcome::Done { provider, key } => { - app.auth_in_flight_method = Some(kigi_shell::agent::custom_providers::login_method_id( - &provider.name, - )); + // Key accepted: open the model selection and fetch concurrently. app.custom_entry.clear(); app.auth_code_input.clear(); + app.custom_select = Some(crate::app::custom_entry::CustomSelectState::new( + provider.clone(), + key.clone(), + )); app.auth_state = AuthState::Authenticating { request_seq, handle: None, auth_url: None, - mode: AuthMode::Pending, + mode: AuthMode::CustomProviderSelect, }; - vec![Effect::PersistCustomProviderAndAuthenticate { + vec![Effect::FetchCustomProviderModels { request_seq, provider, key, @@ -393,6 +397,81 @@ pub(super) fn dispatch_submit_custom_provider_input( } } +/// Enter on an empty manual-model box: persist the chosen models, then auth. +pub(super) fn dispatch_submit_custom_provider_selection(app: &mut AppView) -> Vec { + let request_seq = match &app.auth_state { + AuthState::Authenticating { + request_seq, + mode: AuthMode::CustomProviderSelect, + .. + } => *request_seq, + _ => return vec![], + }; + complete_custom_provider_selection(app, request_seq) +} + +/// Persists the selection once the listing no longer gates the base URL. +fn complete_custom_provider_selection(app: &mut AppView, request_seq: u64) -> Vec { + let Some(state) = app.custom_select.as_mut() else { + return vec![]; + }; + let Some(models) = state.try_finish() else { + return vec![]; + }; + let state = app.custom_select.take().expect("checked above"); + let provider = state.provider.clone(); + let key = state.take_key(); + app.auth_in_flight_method = Some(kigi_shell::agent::custom_providers::login_method_id( + &provider.name, + )); + app.auth_code_input.clear(); + app.auth_state = AuthState::Authenticating { + request_seq, + handle: None, + auth_url: None, + mode: AuthMode::Pending, + }; + vec![Effect::PersistCustomProviderAndAuthenticate { + request_seq, + provider, + key, + models, + }] +} + +/// The login-step model listing came back; fill or flag the selection. +pub(super) fn handle_custom_provider_models_listed( + app: &mut AppView, + request_seq: u64, + result: Result, +) -> Vec { + let in_flight = matches!( + &app.auth_state, + AuthState::Authenticating { + request_seq: current_seq, + mode: AuthMode::CustomProviderSelect, + .. + } if *current_seq == request_seq + ); + if !in_flight { + return vec![]; + } + match (app.custom_select.as_mut(), result) { + (Some(state), Ok(fetch)) => state.apply_fetch(fetch), + (Some(state), Err(reason)) => state.fail_fetch(reason), + (None, _) => {} + } + // A finish pressed during the fetch completes now, with the adopted base. + if app + .custom_select + .as_ref() + .is_some_and(|s| s.finish_requested) + { + return complete_custom_provider_selection(app, request_seq); + } + vec![] +} + /// Esc in the API-key paste box: back to the login picker (no error line). /// Bumps the request seq so any stale in-flight auth result is dropped by /// the `AuthComplete`/`AuthFailed` guards. @@ -400,7 +479,9 @@ pub(super) fn dispatch_cancel_platform_key_entry(app: &mut AppView) -> Vec Vec Vec { Action::SubmitCustomProviderInput(input) => { dispatch_submit_custom_provider_input(app, input) } + Action::SubmitCustomProviderSelection => dispatch_submit_custom_provider_selection(app), Action::CopyAuthUrl => { if let AuthState::Authenticating { auth_url: Some(url), diff --git a/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs b/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs index 8913992..a4481e1 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/task_result.rs @@ -509,9 +509,14 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec { app.auth_state = AuthState::Pending { error: Some(error) }; app.auth_code_input.clear(); + app.custom_select = None; } vec![] } + TaskResult::CustomProviderModelsListed { + request_seq, + result, + } => super::auth::handle_custom_provider_models_listed(app, request_seq, result), TaskResult::AuthUrlReady { request_seq, auth_url, @@ -798,6 +803,7 @@ pub(super) fn dispatch_task_result(result: TaskResult, app: &mut AppView) -> Vec app.login_method_id = None; ensure_login_method(app); app.auth_clipboard_copied = false; + app.custom_select = None; let effects = dispatch_exit_session(app); app.welcome_prompt_focused = false; effects diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs index 8f5b330..cf0fe0d 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/auth.rs @@ -582,9 +582,9 @@ fn custom_field(app: &AppView) -> crate::app::custom_entry::CustomEntryField { } } -/// URL, name, key; the last Enter emits one persist effect. +/// URL, name, key; the last Enter opens the selection screen and fetches. #[test] -fn custom_provider_steps_end_in_one_persist_effect() { +fn custom_provider_steps_end_in_the_model_selection() { use crate::app::custom_entry::CustomEntryField; use kigi_shell::models::custom::CustomApi; @@ -619,7 +619,7 @@ fn custom_provider_steps_end_in_one_persist_effect() { let effects = dispatch(Action::SubmitCustomProviderInput("sk-gw".into()), &mut app); match effects.as_slice() { [ - Effect::PersistCustomProviderAndAuthenticate { + Effect::FetchCustomProviderModels { request_seq, provider, key, @@ -631,21 +631,236 @@ fn custom_provider_steps_end_in_one_persist_effect() { assert_eq!(provider.base_url, "https://api.gw.example/v1"); assert_eq!(key, "sk-gw"); } - other => panic!("expected exactly the persist effect, got {other:?}"), + other => panic!("expected exactly the fetch effect, got {other:?}"), } assert!(matches!( app.auth_state, AuthState::Authenticating { request_seq, - mode: AuthMode::Pending, + mode: AuthMode::CustomProviderSelect, .. } if request_seq == seq )); + assert!(app.custom_select.is_some()); + assert!(app.auth_code_input.is_empty()); +} + +/// The listing fills rows; the selection persists only the chosen models. +#[test] +fn custom_provider_selection_persists_the_chosen_models() { + use kigi_shell::models::custom::{CustomApi, CustomProvider}; + + let mut app = test_app(); + app.auth_state = AuthState::Pending { error: None }; + dispatch( + Action::BeginCustomProviderEntry(CustomApi::OpenAi), + &mut app, + ); + dispatch( + Action::SubmitCustomProviderInput("https://h.example".into()), + &mut app, + ); + dispatch(Action::SubmitCustomProviderInput("gw".into()), &mut app); + dispatch(Action::SubmitCustomProviderInput("sk-gw".into()), &mut app); + let seq = match &app.auth_state { + AuthState::Authenticating { request_seq, .. } => *request_seq, + _ => unreachable!(), + }; + + // The listing lands: rows appear, the probed base is adopted. + let adopted = CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/v1").unwrap(); + let listed = kigi_shell::agent::custom_providers::LoginFetch { + provider: adopted, + model_ids: vec!["m1".to_owned(), "m2".to_owned()], + }; + let effects = dispatch( + Action::TaskComplete(TaskResult::CustomProviderModelsListed { + request_seq: seq, + result: Ok(listed), + }), + &mut app, + ); + assert!(effects.is_empty()); + let state = app.custom_select.as_ref().expect("selection lives"); + assert_eq!(state.provider.base_url, "https://h.example/v1"); + assert_eq!(state.rows.len(), 2); + assert!(!state.rows[0].selected); + + // An empty selection refuses to finish once the listing resolved. + assert!(dispatch(Action::SubmitCustomProviderSelection, &mut app).is_empty()); + assert!( + app.custom_select + .as_ref() + .and_then(|s| s.error.as_deref()) + .is_some() + ); + + // A late manual id typed before the fetch survives the fill. + let state = app.custom_select.as_mut().expect("selection lives"); + state.add_manual("claude-opus-5-5[1m]"); + assert_eq!(state.rows.len(), 3); + assert!(state.rows[2].manual && state.rows[2].selected); + + // Select m2 as well, then finish. + let state = app.custom_select.as_mut().expect("selection lives"); + state.highlight = 1; + state.toggle_highlighted(); + let effects = dispatch(Action::SubmitCustomProviderSelection, &mut app); + match effects.as_slice() { + [ + Effect::PersistCustomProviderAndAuthenticate { + request_seq, + provider, + key, + models, + }, + ] => { + assert_eq!(*request_seq, seq); + assert_eq!(provider.base_url, "https://h.example/v1"); + assert_eq!(key, "sk-gw"); + assert_eq!(models, &["m2".to_owned(), "claude-opus-5-5[1m]".to_owned()]); + } + other => panic!("expected the persist effect, got {other:?}"), + } + assert!(app.custom_select.is_none()); + assert!(matches!( + app.auth_state, + AuthState::Authenticating { + mode: AuthMode::Pending, + .. + } + )); assert_eq!( app.auth_in_flight_method.as_ref().map(|m| m.0.as_ref()), Some("custom:gw") ); - assert!(app.auth_code_input.is_empty()); +} + +/// Finishing while the listing is in flight completes when it lands. +#[test] +fn custom_provider_selection_finish_waits_for_the_fetch() { + use kigi_shell::models::custom::{CustomApi, CustomProvider}; + + let mut app = test_app(); + app.auth_state = AuthState::Pending { error: None }; + dispatch( + Action::BeginCustomProviderEntry(CustomApi::OpenAi), + &mut app, + ); + dispatch( + Action::SubmitCustomProviderInput("https://h.example".into()), + &mut app, + ); + dispatch(Action::SubmitCustomProviderInput("gw".into()), &mut app); + dispatch(Action::SubmitCustomProviderInput("sk-gw".into()), &mut app); + let seq = match &app.auth_state { + AuthState::Authenticating { request_seq, .. } => *request_seq, + _ => unreachable!(), + }; + + let state = app.custom_select.as_mut().expect("selection lives"); + state.add_manual("m1"); + // Enter with an empty box while the fetch runs: no persist yet. + assert!(dispatch(Action::SubmitCustomProviderSelection, &mut app).is_empty()); + assert!(app.custom_select.as_ref().expect("lives").finish_requested); + assert!(matches!( + app.auth_state, + AuthState::Authenticating { + mode: AuthMode::CustomProviderSelect, + .. + } + )); + + // The listing lands with an adopted /v1 base; the finish fires now. + let adopted = CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/v1").unwrap(); + let effects = dispatch( + Action::TaskComplete(TaskResult::CustomProviderModelsListed { + request_seq: seq, + result: Ok(kigi_shell::agent::custom_providers::LoginFetch { + provider: adopted, + model_ids: vec!["m2".to_owned()], + }), + }), + &mut app, + ); + match effects.as_slice() { + [ + Effect::PersistCustomProviderAndAuthenticate { + provider, models, .. + }, + ] => { + assert_eq!(provider.base_url, "https://h.example/v1"); + assert_eq!(models, &["m1".to_owned()]); + } + other => panic!("expected the persist effect, got {other:?}"), + } + assert!(app.custom_select.is_none()); +} + +/// A failed listing keeps the manual input usable and shows the reason. +#[test] +fn custom_provider_selection_survives_a_failed_fetch() { + use kigi_shell::models::custom::CustomApi; + + let mut app = test_app(); + app.auth_state = AuthState::Pending { error: None }; + dispatch( + Action::BeginCustomProviderEntry(CustomApi::OpenAi), + &mut app, + ); + dispatch( + Action::SubmitCustomProviderInput("https://h.example/v1".into()), + &mut app, + ); + dispatch(Action::SubmitCustomProviderInput("gw".into()), &mut app); + dispatch(Action::SubmitCustomProviderInput("sk-gw".into()), &mut app); + let seq = match &app.auth_state { + AuthState::Authenticating { request_seq, .. } => *request_seq, + _ => unreachable!(), + }; + + let effects = dispatch( + Action::TaskComplete(TaskResult::CustomProviderModelsListed { + request_seq: seq, + result: Err("https://h.example/v1/models answered HTTP 404".to_owned()), + }), + &mut app, + ); + assert!(effects.is_empty()); + let state = app.custom_select.as_ref().expect("selection lives"); + assert!(matches!( + &state.fetch, + crate::app::custom_entry::SelectFetch::Failed(reason) if reason.contains("404") + )); + + // A stale result for an older request never clobbers the screen. + let effects = dispatch( + Action::TaskComplete(TaskResult::CustomProviderModelsListed { + request_seq: seq + 99, + result: Ok(kigi_shell::agent::custom_providers::LoginFetch { + provider: state.provider.clone(), + model_ids: vec!["stale".to_owned()], + }), + }), + &mut app, + ); + assert!(effects.is_empty()); + assert!( + app.custom_select + .as_ref() + .expect("selection lives") + .rows + .is_empty() + ); + + // Manual add still finishes the flow. + let state = app.custom_select.as_mut().expect("selection lives"); + state.add_manual("m1"); + let effects = dispatch(Action::SubmitCustomProviderSelection, &mut app); + assert!(matches!( + effects.as_slice(), + [Effect::PersistCustomProviderAndAuthenticate { models, .. }] if models == &["m1".to_owned()] + )); } /// A refused field keeps the user on it. diff --git a/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs b/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs index fdae6f0..0e3c154 100644 --- a/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs +++ b/crates/codegen/kigi-tui/src/app/dispatch/tests/mod.rs @@ -127,6 +127,7 @@ fn test_app() -> AppView { auth_start_mode: AuthMode::Pending, auth_code_input: String::new(), custom_entry: Default::default(), + custom_select: None, next_auth_request_seq: 1, deferred_startup: Default::default(), auth_use_oauth: false, diff --git a/crates/codegen/kigi-tui/src/app/effects/mod.rs b/crates/codegen/kigi-tui/src/app/effects/mod.rs index 063323d..96dda4a 100644 --- a/crates/codegen/kigi-tui/src/app/effects/mod.rs +++ b/crates/codegen/kigi-tui/src/app/effects/mod.rs @@ -1915,13 +1915,15 @@ pub(crate) fn execute( request_seq, provider, key, + models, } => { let tx = acp_tx.clone(); let abort_handle = tasks.spawn(async move { // Persist first so authenticate finds the definition and key. - if let Err(e) = - kigi_shell::agent::custom_providers::save_custom_provider(&provider, &key) - .await + if let Err(e) = kigi_shell::agent::custom_providers::save_custom_provider( + &provider, &key, &models, + ) + .await { let error = format!("Couldn't save the provider: {e:#}"); ulog::error( @@ -1940,6 +1942,23 @@ pub(crate) fn execute( }); meta.auth_abort_handle = Some((request_seq, abort_handle)); } + Effect::FetchCustomProviderModels { + request_seq, + provider, + key, + } => { + tasks.spawn(async move { + let result = tokio::task::spawn_blocking(move || { + kigi_shell::agent::custom_providers::fetch_listing_for_login(&provider, &key) + }) + .await + .unwrap_or_else(|e| Err(format!("listing task failed: {e}"))); + TaskResult::CustomProviderModelsListed { + request_seq, + result, + } + }); + } Effect::SubmitAuthCode { request_seq, code } => { let tx = acp_tx.clone(); tasks diff --git a/crates/codegen/kigi-tui/src/views/welcome/mod.rs b/crates/codegen/kigi-tui/src/views/welcome/mod.rs index 3aa0654..625e654 100644 --- a/crates/codegen/kigi-tui/src/views/welcome/mod.rs +++ b/crates/codegen/kigi-tui/src/views/welcome/mod.rs @@ -462,6 +462,7 @@ pub struct WelcomeRenderParams<'a> { pub login_label: Option<&'a str>, pub auth_code_input: &'a str, pub custom_entry: &'a crate::app::custom_entry::CustomEntryDraft, + pub custom_select: Option<&'a crate::app::custom_entry::CustomSelectState>, pub clipboard_copied: bool, pub show_raw_url: bool, pub tip: Option<&'a str>, @@ -605,6 +606,7 @@ pub fn render_welcome( *mode, params.auth_code_input, params.custom_entry, + params.custom_select, params.clipboard_copied, params.show_raw_url, ); @@ -1179,6 +1181,7 @@ fn render_welcome_authenticating( mode: AuthMode, auth_code_input: &str, custom_entry: &crate::app::custom_entry::CustomEntryDraft, + custom_select: Option<&crate::app::custom_entry::CustomSelectState>, clipboard_copied: bool, show_raw_url: bool, ) -> (Option, Option) { @@ -1336,6 +1339,21 @@ fn render_welcome_authenticating( (None, None) } + AuthMode::CustomProviderSelect => { + if let Some(select) = custom_select { + render_custom_select( + content_area, + buf, + theme, + top_pad, + logo_line_count, + select, + auth_code_input, + ); + } + (None, None) + } + AuthMode::Command => render_browser_status_arm( content_area, buf, @@ -2034,6 +2052,182 @@ fn render_entry_box( Paragraph::new(hints).render(hint_area, buf); } +/// The custom provider model-selection screen: fetched rows + manual input. +fn render_custom_select( + content_area: Rect, + buf: &mut Buffer, + theme: &Theme, + top_pad: u16, + logo_line_count: u16, + select: &crate::app::custom_entry::CustomSelectState, + input: &str, +) { + use crate::app::custom_entry::SelectFetch; + let h_pad: u16 = content_area.width / 8; + let inner_width = content_area.width.saturating_sub(h_pad * 2).max(1); + let instruction = format!( + "Select the models to add for {} ({})", + select.provider.name, select.provider.base_url + ); + let status: Option<(String, Style)> = match &select.fetch { + SelectFetch::InFlight if select.finish_requested => Some(( + "Fetching models... finishing when the list lands".to_owned(), + Style::default().fg(theme.gray), + )), + SelectFetch::InFlight => Some(( + "Fetching models...".to_owned(), + Style::default().fg(theme.gray), + )), + SelectFetch::Listed if select.rows.is_empty() => Some(( + "The provider served no models; add one below".to_owned(), + Style::default().fg(theme.warning), + )), + SelectFetch::Listed => None, + SelectFetch::Failed(reason) => Some(( + format!("Fetch failed: {reason}"), + Style::default().fg(theme.warning), + )), + }; + // Word-wrap-aware rows; byte math undercounts and clips the failure reason. + let rows_of = |text: &str| -> u16 { + crate::render::wrapping::word_wrap_line( + &Line::from(Span::raw(text.to_owned())), + inner_width as usize, + ) + .len() as u16 + }; + let msg_height = rows_of(&instruction) + + status.as_ref().map_or(0, |(s, _)| rows_of(s)) + + select.error.as_deref().map_or(0, rows_of); + let [ + _, + logo_area, + _, + msg_area, + _, + rows_area, + _, + prompt_area, + _, + hint_area, + _, + ] = Layout::vertical([ + Constraint::Length(top_pad), + Constraint::Length(logo_line_count), + Constraint::Length(1), + // Min: spare rows land here, so a mis-measured wrap never clips. + Constraint::Min(msg_height), + Constraint::Length(1), + Constraint::Min(3), + Constraint::Length(1), + Constraint::Length(5), + Constraint::Length(1), + Constraint::Length(1), + Constraint::Min(0), + ]) + .areas(content_area); + + render_logo(logo_area, buf, theme, content_area.height); + + let mut lines = vec![ + Line::from(Span::styled( + instruction, + Style::default().fg(theme.gray_bright), + )) + .alignment(Alignment::Center), + ]; + if let Some((text, style)) = status { + lines.push(Line::from(Span::styled(text, style)).alignment(Alignment::Center)); + } + if let Some(error) = select.error.as_deref() { + lines.push( + Line::from(Span::styled(error, Style::default().fg(theme.warning))) + .alignment(Alignment::Center), + ); + } + Paragraph::new(lines) + .wrap(Wrap { trim: false }) + .block(Block::default().padding(Padding::horizontal(h_pad))) + .render(msg_area, buf); + + let visible = rows_area.height as usize; + let start = select + .highlight + .saturating_add(1) + .saturating_sub(visible) + .min(select.rows.len().saturating_sub(visible)); + let mut row_lines = Vec::new(); + for (idx, row) in select.rows.iter().enumerate().skip(start).take(visible) { + let highlighted = idx == select.highlight; + let check = if row.selected { "[x]" } else { "[ ]" }; + let suffix = if row.manual { " (manual)" } else { "" }; + let style = if highlighted { + Style::default() + .fg(theme.accent_user) + .add_modifier(Modifier::BOLD) + } else { + Style::default().fg(theme.gray_bright) + }; + let marker = if highlighted { "› " } else { " " }; + row_lines.push(Line::from(Span::styled( + format!("{marker}{check} {}{suffix}", row.id), + style, + ))); + } + Paragraph::new(row_lines) + .block(Block::default().padding(Padding::horizontal(h_pad))) + .render(rows_area, buf); + + let [_, prompt_centered, _] = Layout::horizontal([ + Constraint::Min(0), + Constraint::Length(content_area.width), + Constraint::Min(0), + ]) + .flex(Flex::Center) + .areas(prompt_area); + render_auth_input_box( + prompt_centered, + buf, + theme, + input, + "Type a model id and press enter to add", + false, + ); + + let hints = Line::from(vec![ + Span::styled( + "↑↓", + Style::default() + .fg(theme.accent_user) + .add_modifier(Modifier::BOLD), + ), + Span::styled(" navigate ", Style::default().fg(theme.gray)), + Span::styled( + "space", + Style::default() + .fg(theme.accent_user) + .add_modifier(Modifier::BOLD), + ), + Span::styled(" toggle ", Style::default().fg(theme.gray)), + Span::styled( + "enter", + Style::default() + .fg(theme.accent_user) + .add_modifier(Modifier::BOLD), + ), + Span::styled(" add model / done ", Style::default().fg(theme.gray)), + Span::styled( + "esc", + Style::default() + .fg(theme.accent_user) + .add_modifier(Modifier::BOLD), + ), + Span::styled(" back", Style::default().fg(theme.gray)), + ]) + .alignment(Alignment::Center); + Paragraph::new(hints).render(hint_area, buf); +} + /// Render the auth token input box (loopback mode). fn render_auth_input_box( area: Rect, @@ -2197,6 +2391,7 @@ mod tests { login_label: None, auth_code_input: "", custom_entry: &EMPTY_CUSTOM_ENTRY, + custom_select: None, clipboard_copied: false, show_raw_url: false, tip: None, @@ -2409,12 +2604,87 @@ mod tests { AuthMode::CustomProviderEntry(step), input, draft, + None, false, false, ); buffer_text(&buf) } + fn render_custom_select( + select: &crate::app::custom_entry::CustomSelectState, + input: &str, + ) -> String { + let area = Rect::new(0, 0, 100, 40); + let mut buf = Buffer::empty(area); + render_welcome_authenticating( + area, + &mut buf, + &Theme::current(), + logo_line_count(area.height), + None, + AuthMode::CustomProviderSelect, + input, + &crate::app::custom_entry::CustomEntryDraft::new(), + Some(select), + false, + false, + ); + buffer_text(&buf) + } + + /// The selection screen lists rows, status, and the manual input together. + #[test] + fn custom_select_arm_lists_rows_status_and_input() { + use crate::app::custom_entry::CustomSelectState; + use kigi_shell::models::custom::{CustomApi, CustomProvider}; + + let provider = + CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/v1").unwrap(); + let mut select = CustomSelectState::new(provider, "sk".to_owned()); + let fetching = render_custom_select(&select, ""); + assert!(fetching.contains("Fetching models"), "{fetching}"); + assert!(fetching.contains("add model / done"), "{fetching}"); + + select.apply_fetch(kigi_shell::agent::custom_providers::LoginFetch { + provider: select.provider.clone(), + model_ids: vec!["m1".to_owned(), "m2".to_owned()], + }); + select.toggle_highlighted(); + select.add_manual("claude-opus-5-5[1m]"); + let listed = render_custom_select(&select, "typed-so-far"); + assert!(listed.contains("[x] m1"), "{listed}"); + assert!(listed.contains("[ ] m2"), "{listed}"); + assert!(listed.contains("claude-opus-5-5[1m] (manual)"), "{listed}"); + assert!(listed.contains("typed-so-far"), "{listed}"); + + select.fail_fetch("boom".to_owned()); + let failed = render_custom_select(&select, ""); + assert!(failed.contains("Fetch failed: boom"), "{failed}"); + + // A two-hop failure reason must render whole, not clip at the wrap. + let reason = "https://h.example/models answered HTTP 404; https://h.example/v1/models also failed: https://h.example/v1/models answered HTTP 500"; + select.fail_fetch(reason.to_owned()); + let text = render_custom_select(&select, ""); + assert!(text.contains("Fetch failed:"), "{text}"); + assert!(text.contains("HTTP 500"), "{text}"); + } + + /// An empty selection refuses once the fetch resolved. + #[test] + fn custom_select_arm_shows_the_finish_refusal() { + use crate::app::custom_entry::CustomSelectState; + use kigi_shell::models::custom::{CustomApi, CustomProvider}; + + let provider = + CustomProvider::new("gw", CustomApi::OpenAi, "https://h.example/v1").unwrap(); + let mut select = CustomSelectState::new(provider, "sk".to_owned()); + select.fail_fetch("boom".to_owned()); + assert!(select.try_finish().is_none()); + let text = render_custom_select(&select, ""); + assert!(text.contains("Select at least one model"), "{text}"); + } + /// Only the key step masks its input. #[test] fn custom_entry_arm_masks_only_the_key() { @@ -2479,6 +2749,7 @@ mod tests { // auth_code_input "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, // clipboard_copied false, // show_raw_url @@ -3147,6 +3418,7 @@ mod tests { // auth_code_input — unused in device mode "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, // clipboard_copied false, // show_raw_url @@ -3204,6 +3476,7 @@ mod tests { AuthMode::Device, "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, false, // show_raw_url true, @@ -3232,6 +3505,7 @@ mod tests { AuthMode::Device, "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, false, // show_raw_url true, @@ -3271,6 +3545,7 @@ mod tests { AuthMode::Device, "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, false, // show_raw_url true, @@ -3313,6 +3588,7 @@ mod tests { // auth_code_input — unused "", &crate::app::custom_entry::CustomEntryDraft::new(), + None, // clipboard_copied false, // show_raw_url