Fix provider model inventories and Bedrock verification
This commit is contained in:
+147
-40
@@ -605,6 +605,11 @@ pub struct LLMPreferences {
|
||||
catalog_refresh_generations: HashMap<Option<usize>, u64>,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
chatgpt_subscription_models_refresh_in_flight: bool,
|
||||
bedrock_models_refresh_in_flight: bool,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
bedrock_verified_config: Option<super::provider::client::BedrockClientConfig>,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
bedrock_verified_model_ids: HashSet<String>,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
acp_selections: HashMap<LLMId, AcpModelSelection>,
|
||||
}
|
||||
@@ -708,8 +713,23 @@ impl LLMPreferences {
|
||||
| AISettingsChangedEvent::AcpAgents { .. }
|
||||
| AISettingsChangedEvent::AcpAgentId { .. }
|
||||
| AISettingsChangedEvent::BedrockModels { .. }
|
||||
| AISettingsChangedEvent::BedrockAuthMethod { .. }
|
||||
| AISettingsChangedEvent::BedrockProfile { .. }
|
||||
| AISettingsChangedEvent::BedrockRegion { .. }
|
||||
| AISettingsChangedEvent::BedrockAccessKeyId { .. }
|
||||
| AISettingsChangedEvent::BedrockSecretAccessKey { .. }
|
||||
| AISettingsChangedEvent::BedrockCrossRegionInference { .. }
|
||||
) {
|
||||
let previous_ids: Vec<_> = me.openai_provider_routing.keys().cloned().collect();
|
||||
let previous_ids: Vec<_> = me
|
||||
.models_by_feature
|
||||
.agent_mode
|
||||
.choices
|
||||
.iter()
|
||||
.filter(|model| {
|
||||
matches!(model.provider, LLMProvider::LiteLLM | LLMProvider::Bedrock)
|
||||
})
|
||||
.map(|model| model.id.as_str().to_owned())
|
||||
.collect();
|
||||
me.inject_bedrock_models(ctx);
|
||||
me.inject_openai_models(ctx);
|
||||
if matches!(
|
||||
@@ -720,11 +740,6 @@ impl LLMPreferences {
|
||||
) {
|
||||
me.fetch_openai_models_from_endpoint(ctx);
|
||||
}
|
||||
if matches!(event, AISettingsChangedEvent::BedrockEnabled { .. })
|
||||
&& *AISettings::as_ref(ctx).bedrock_enabled.value()
|
||||
{
|
||||
me.refresh_bedrock_models(ctx);
|
||||
}
|
||||
me.refresh_changed_litellm_connections(ctx);
|
||||
// Safety: ensure the default model is still present in choices.
|
||||
// If all provider models were removed, the default_id would dangle.
|
||||
@@ -770,6 +785,11 @@ impl LLMPreferences {
|
||||
catalog_refresh_generations: HashMap::new(),
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
chatgpt_subscription_models_refresh_in_flight: false,
|
||||
bedrock_models_refresh_in_flight: false,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
bedrock_verified_config: None,
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
bedrock_verified_model_ids: HashSet::new(),
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
acp_selections: HashMap::new(),
|
||||
};
|
||||
@@ -789,9 +809,8 @@ impl LLMPreferences {
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
{
|
||||
// Make the persisted Bedrock catalog available immediately. Agreement
|
||||
// revalidation stays in the background and replaces this cache when it
|
||||
// completes, so startup never waits on the per-model AWS checks.
|
||||
// Saved toggle preferences survive startup, but access must be
|
||||
// verified against AWS before a Bedrock model becomes selectable.
|
||||
me.inject_bedrock_models(ctx);
|
||||
me.refresh_bedrock_models(ctx);
|
||||
me.inject_openai_models(ctx);
|
||||
@@ -807,40 +826,87 @@ impl LLMPreferences {
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn refresh_bedrock_models(&mut self, ctx: &mut ModelContext<Self>) {
|
||||
if self.bedrock_models_refresh_in_flight {
|
||||
return;
|
||||
}
|
||||
let settings = AISettings::as_ref(ctx);
|
||||
if !*settings.bedrock_enabled.value() {
|
||||
return;
|
||||
}
|
||||
let config = crate::ai::provider::client::BedrockClientConfig {
|
||||
auth_method: *settings.bedrock_auth_method.value(),
|
||||
profile: settings.bedrock_profile.value().clone(),
|
||||
region: settings.bedrock_region.value().clone(),
|
||||
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
||||
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
||||
session_token: None,
|
||||
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
|
||||
};
|
||||
self.bedrock_models_refresh_in_flight = true;
|
||||
let config = bedrock_discovery_config(settings);
|
||||
let requested_config = config.clone();
|
||||
self.bedrock_verified_config = None;
|
||||
self.bedrock_verified_model_ids.clear();
|
||||
self.inject_bedrock_models(ctx);
|
||||
|
||||
let _ = ctx.spawn(
|
||||
async move { crate::ai::provider::discovery::discover_available_models(config).await },
|
||||
|me, result, ctx| match result {
|
||||
Ok(models) => {
|
||||
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
||||
if let Err(error) = settings.bedrock_models.set_value(models, ctx) {
|
||||
log::warn!("[bedrock] Failed to persist discovered models: {error}");
|
||||
}
|
||||
});
|
||||
me.inject_bedrock_models(ctx);
|
||||
me.ensure_default_model_present();
|
||||
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
|
||||
move |me, result, ctx| {
|
||||
let result = result.map(|discovery| {
|
||||
for record in discovery.records {
|
||||
crate::ai::remote_logging::log_model_event(ctx, record);
|
||||
}
|
||||
discovery.models
|
||||
});
|
||||
me.bedrock_models_refresh_in_flight = false;
|
||||
if !*AISettings::as_ref(ctx).bedrock_enabled.value() {
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
log::debug!("[bedrock] Startup model discovery unavailable: {error}");
|
||||
if bedrock_discovery_config(AISettings::as_ref(ctx)) != requested_config {
|
||||
return;
|
||||
}
|
||||
match result {
|
||||
Ok(models) => {
|
||||
// Verification belongs to this live connection, not serialized settings.
|
||||
me.bedrock_verified_model_ids =
|
||||
models.iter().map(|model| model.model_id.clone()).collect();
|
||||
me.bedrock_verified_config = Some(requested_config);
|
||||
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
||||
let models = super::provider::discovery::preserve_model_toggles(
|
||||
settings.bedrock_models.value(),
|
||||
models,
|
||||
);
|
||||
if let Err(error) = settings.bedrock_models.set_value(models, ctx) {
|
||||
log::warn!(
|
||||
"[bedrock] Failed to persist discovered models: {error}"
|
||||
);
|
||||
}
|
||||
});
|
||||
me.inject_bedrock_models(ctx);
|
||||
me.ensure_default_model_present();
|
||||
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
|
||||
}
|
||||
Err(error) => {
|
||||
log::debug!("[bedrock] Startup model discovery unavailable: {error}");
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) fn record_bedrock_verification(
|
||||
&mut self,
|
||||
config: super::provider::client::BedrockClientConfig,
|
||||
models: &[BedrockModelConfig],
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
self.bedrock_verified_model_ids =
|
||||
models.iter().map(|model| model.model_id.clone()).collect();
|
||||
self.bedrock_verified_config = Some(config);
|
||||
self.inject_bedrock_models(ctx);
|
||||
self.ensure_default_model_present();
|
||||
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) fn bedrock_model_is_verified(&self, model_id: &str, ctx: &AppContext) -> bool {
|
||||
self.bedrock_verified_config.as_ref()
|
||||
== Some(&bedrock_discovery_config(AISettings::as_ref(ctx)))
|
||||
&& self.bedrock_verified_model_ids.contains(model_id)
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn inject_bedrock_models(&mut self, ctx: &AppContext) {
|
||||
// Remove only previously injected Bedrock entries. Other enabled provider
|
||||
@@ -906,6 +972,9 @@ impl LLMPreferences {
|
||||
|
||||
let effective = effective;
|
||||
for model in effective {
|
||||
if !model.enabled || !self.bedrock_model_is_verified(&model.model_id, ctx) {
|
||||
continue;
|
||||
}
|
||||
let metadata = model_metadata(&model.model_id);
|
||||
if metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false)) {
|
||||
log::debug!(
|
||||
@@ -1078,9 +1147,11 @@ impl LLMPreferences {
|
||||
configured_models
|
||||
};
|
||||
|
||||
// The registry replaces legacy settings. Otherwise stale legacy credentials
|
||||
// win duplicate model IDs before the configured provider is considered.
|
||||
if settings.openai_providers.value().is_empty() && !single_provider_models.is_empty() {
|
||||
// Only a registry entry for the same endpoint replaces the legacy
|
||||
// connection. Adding an unrelated subscription must not hide LiteLLM.
|
||||
let include_legacy =
|
||||
settings.legacy_openai_provider_enabled() && !single_provider_models.is_empty();
|
||||
if include_legacy {
|
||||
let base_url = settings.openai_base_url.value().clone();
|
||||
let api_key = {
|
||||
let key = settings.openai_api_key.value().clone();
|
||||
@@ -1132,7 +1203,7 @@ impl LLMPreferences {
|
||||
.is_none_or(|project| project.trim().is_empty()),
|
||||
OpenAIProviderKind::ChatGPTSubscription => false,
|
||||
};
|
||||
if !provider.enabled || missing_credentials || provider.models.is_empty() {
|
||||
if missing_credentials || provider.models.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some((
|
||||
@@ -1147,6 +1218,12 @@ impl LLMPreferences {
|
||||
))
|
||||
}),
|
||||
);
|
||||
// Keep existing raw-model selections routed to registry connections.
|
||||
// The legacy connection remains selectable under a qualified key when
|
||||
// it serves the same model as a registered provider.
|
||||
if include_legacy && provider_entries.len() > 1 {
|
||||
provider_entries.rotate_left(1);
|
||||
}
|
||||
|
||||
if provider_entries.is_empty() {
|
||||
return;
|
||||
@@ -1154,6 +1231,7 @@ impl LLMPreferences {
|
||||
|
||||
let mut total_injected = 0;
|
||||
let mut seen_model_ids: HashSet<String> = HashSet::new();
|
||||
let mut seen_connections = HashSet::new();
|
||||
for (
|
||||
provider_name,
|
||||
provider_kind,
|
||||
@@ -1165,14 +1243,13 @@ impl LLMPreferences {
|
||||
models,
|
||||
) in provider_entries
|
||||
{
|
||||
if !provider_enabled {
|
||||
let connection_identity = format!("{provider_kind:?}|{provider_name}|{base_url}|{provider_project_id:?}|{provider_location:?}");
|
||||
if !seen_connections.insert(connection_identity.clone()) {
|
||||
continue;
|
||||
}
|
||||
let mut connection_model_ids = HashSet::new();
|
||||
for model in &models {
|
||||
if !model.enabled {
|
||||
continue;
|
||||
}
|
||||
if !seen_model_ids.insert(model.model_id.clone()) {
|
||||
if !connection_model_ids.insert(model.model_id.clone()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -1189,10 +1266,25 @@ impl LLMPreferences {
|
||||
|
||||
for reasoning_effort in reasoning_efforts {
|
||||
let reasoning_effort = reasoning_effort.cloned();
|
||||
let model_key = reasoning_effort.as_deref().map_or_else(
|
||||
let mut model_key = reasoning_effort.as_deref().map_or_else(
|
||||
|| model.model_id.clone(),
|
||||
|effort| openai_model_variant_id(&model.model_id, effort),
|
||||
);
|
||||
if !seen_model_ids.insert(model_key.clone()) {
|
||||
use sha2::{Digest, Sha256};
|
||||
let connection_hash =
|
||||
format!("{:x}", Sha256::digest(connection_identity.as_bytes()));
|
||||
model_key = format!("{model_key}::connection::{connection_hash}");
|
||||
if !seen_model_ids.insert(model_key.clone()) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
// Disabled entries still reserve their selection keys so
|
||||
// toggling a connection cannot silently reroute another
|
||||
// connection's saved model selection.
|
||||
if !provider_enabled || !model.enabled {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Register the routing entry. Reasoning variants keep the provider's
|
||||
// actual model ID while using their synthetic key only for selection.
|
||||
@@ -3355,3 +3447,18 @@ fn get_cached_models(_app: &mut AppContext) -> Option<ModelsByFeature> {
|
||||
#[cfg(test)]
|
||||
#[path = "llms_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn bedrock_discovery_config(
|
||||
settings: &AISettings,
|
||||
) -> crate::ai::provider::client::BedrockClientConfig {
|
||||
crate::ai::provider::client::BedrockClientConfig {
|
||||
auth_method: *settings.bedrock_auth_method.value(),
|
||||
profile: settings.bedrock_profile.value().clone(),
|
||||
region: settings.bedrock_region.value().clone(),
|
||||
access_key_id: settings.bedrock_access_key_id.value().clone(),
|
||||
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
|
||||
session_token: None,
|
||||
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -161,6 +161,8 @@ fn openai_model(model_id: &str) -> OpenAIModelConfig {
|
||||
|
||||
fn bedrock_model(model_id: &str) -> BedrockModelConfig {
|
||||
BedrockModelConfig {
|
||||
enabled: true,
|
||||
access_verified: true,
|
||||
model_id: model_id.to_string(),
|
||||
display_name: model_id.to_string(),
|
||||
vision_supported: false,
|
||||
@@ -232,6 +234,9 @@ fn empty_preferences() -> LLMPreferences {
|
||||
litellm_connections: Vec::new(),
|
||||
catalog_refresh_generations: HashMap::new(),
|
||||
chatgpt_subscription_models_refresh_in_flight: false,
|
||||
bedrock_models_refresh_in_flight: false,
|
||||
bedrock_verified_config: None,
|
||||
bedrock_verified_model_ids: HashSet::new(),
|
||||
acp_selections: HashMap::new(),
|
||||
}
|
||||
}
|
||||
@@ -304,6 +309,14 @@ fn bedrock_and_openai_provider_models_coexist_in_picker() {
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.inject_openai_models(ctx);
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
});
|
||||
|
||||
@@ -323,6 +336,14 @@ fn bedrock_and_openai_provider_models_coexist_in_picker() {
|
||||
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
preferences.inject_openai_models(ctx);
|
||||
});
|
||||
@@ -429,6 +450,9 @@ fn acp_models_are_injected_only_while_acp_is_enabled() {
|
||||
litellm_connections: Vec::new(),
|
||||
catalog_refresh_generations: HashMap::new(),
|
||||
chatgpt_subscription_models_refresh_in_flight: false,
|
||||
bedrock_models_refresh_in_flight: false,
|
||||
bedrock_verified_config: None,
|
||||
bedrock_verified_model_ids: HashSet::new(),
|
||||
acp_selections: HashMap::new(),
|
||||
};
|
||||
app.read(|ctx| preferences.inject_acp_models(ctx));
|
||||
@@ -489,6 +513,14 @@ fn acp_is_selectable_before_protocol_model_discovery_completes() {
|
||||
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
preferences.inject_acp_models(ctx);
|
||||
});
|
||||
@@ -571,6 +603,14 @@ fn acp_models_expand_reasoning_levels_for_only_the_configured_agent() {
|
||||
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
preferences.inject_acp_models(ctx);
|
||||
});
|
||||
@@ -710,6 +750,14 @@ fn acp_bedrock_models_are_hidden_while_bedrock_is_disabled() {
|
||||
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
preferences.inject_acp_models(ctx);
|
||||
});
|
||||
@@ -751,6 +799,14 @@ fn disabled_providers_do_not_leave_models_in_the_runtime_inventory() {
|
||||
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences.bedrock_verified_model_ids = AISettings::as_ref(ctx)
|
||||
.bedrock_models
|
||||
.value()
|
||||
.iter()
|
||||
.map(|model| model.model_id.clone())
|
||||
.collect();
|
||||
preferences.inject_bedrock_models(ctx);
|
||||
preferences.inject_openai_models(ctx);
|
||||
});
|
||||
@@ -940,6 +996,94 @@ fn legacy_chatgpt_default_limits_routing() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscription_and_litellm_connections_keep_overlapping_models_and_independent_toggles() {
|
||||
App::test((), |mut app| async move {
|
||||
initialize_settings_for_tests(&mut app);
|
||||
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
||||
settings.openai_enabled.set_value(true, ctx).unwrap();
|
||||
settings
|
||||
.openai_base_url
|
||||
.set_value("https://legacy.test/v1".to_owned(), ctx)
|
||||
.unwrap();
|
||||
settings
|
||||
.openai_models
|
||||
.set_value(vec![openai_model("legacy-only")], ctx)
|
||||
.unwrap();
|
||||
settings
|
||||
.openai_providers
|
||||
.set_value(
|
||||
vec![
|
||||
OpenAIProviderConfig {
|
||||
kind: OpenAIProviderKind::ChatGPTSubscription,
|
||||
enabled: true,
|
||||
name: "Subscription".to_owned(),
|
||||
base_url: String::new(),
|
||||
api_key: None,
|
||||
project_id: None,
|
||||
location: None,
|
||||
models: vec![openai_model("shared-model")],
|
||||
},
|
||||
OpenAIProviderConfig {
|
||||
kind: OpenAIProviderKind::LiteLLM,
|
||||
enabled: true,
|
||||
name: "LiteLLM".to_owned(),
|
||||
base_url: "https://proxy.test/v1".to_owned(),
|
||||
api_key: None,
|
||||
project_id: None,
|
||||
location: None,
|
||||
models: vec![openai_model("shared-model")],
|
||||
},
|
||||
],
|
||||
ctx,
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| preferences.inject_openai_models(ctx));
|
||||
let models = &preferences.models_by_feature.agent_mode.choices;
|
||||
assert_eq!(models.len(), 3);
|
||||
for (label, kind) in [
|
||||
("Subscription", OpenAIProviderKind::ChatGPTSubscription),
|
||||
("LiteLLM", OpenAIProviderKind::LiteLLM),
|
||||
] {
|
||||
let model = models
|
||||
.iter()
|
||||
.find(|model| model.description.as_deref() == Some(label))
|
||||
.unwrap();
|
||||
let route = preferences
|
||||
.openai_client_config_for_model(model.id.as_str())
|
||||
.unwrap();
|
||||
assert_eq!(route.kind, kind);
|
||||
assert_eq!(route.model.as_deref(), Some("shared-model"));
|
||||
}
|
||||
assert!(models
|
||||
.iter()
|
||||
.any(|model| model.id.as_str() == "legacy-only"));
|
||||
let litellm_id = models
|
||||
.iter()
|
||||
.find(|model| model.description.as_deref() == Some("LiteLLM"))
|
||||
.unwrap()
|
||||
.id
|
||||
.clone();
|
||||
AISettings::handle(&app).update(&mut app, |settings, ctx| {
|
||||
let mut providers = settings.openai_providers.value().clone();
|
||||
providers[0].models[0].enabled = false;
|
||||
settings.openai_providers.set_value(providers, ctx).unwrap();
|
||||
});
|
||||
app.read(|ctx| preferences.inject_openai_models(ctx));
|
||||
let models = &preferences.models_by_feature.agent_mode.choices;
|
||||
assert_eq!(models.len(), 2);
|
||||
assert!(models
|
||||
.iter()
|
||||
.all(|model| model.description.as_deref() != Some("Subscription")));
|
||||
assert!(models
|
||||
.iter()
|
||||
.any(|model| model.description.as_deref() == Some("LiteLLM")));
|
||||
assert!(models.iter().any(|model| model.id == litellm_id));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_registry_overrides_legacy_endpoint_for_duplicate_models() {
|
||||
App::test((), |mut app| async move {
|
||||
@@ -1071,3 +1215,44 @@ async fn model_refresh_distinguishes_empty_catalog_from_failure() {
|
||||
response.remove_async().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_verification_survives_settings_reload_but_not_connection_changes() {
|
||||
App::test((), |mut app| async move {
|
||||
initialize_settings_for_tests(&mut app);
|
||||
let model = bedrock_model("anthropic.claude-test");
|
||||
let reloaded: BedrockModelConfig =
|
||||
serde_json::from_value(serde_json::to_value(&model).unwrap()).unwrap();
|
||||
assert!(!reloaded.access_verified);
|
||||
let mut preferences = empty_preferences();
|
||||
app.read(|ctx| {
|
||||
preferences.bedrock_verified_config =
|
||||
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
|
||||
preferences
|
||||
.bedrock_verified_model_ids
|
||||
.insert(model.model_id.clone());
|
||||
});
|
||||
app.update(|ctx| {
|
||||
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
||||
settings
|
||||
.bedrock_models
|
||||
.set_value(vec![reloaded], ctx)
|
||||
.unwrap();
|
||||
})
|
||||
});
|
||||
app.read(|ctx| assert!(preferences.bedrock_model_is_verified(&model.model_id, ctx)));
|
||||
// Another process/startup cannot inherit an in-memory approval.
|
||||
app.read(|ctx| {
|
||||
assert!(!empty_preferences().bedrock_model_is_verified(&model.model_id, ctx))
|
||||
});
|
||||
app.update(|ctx| {
|
||||
AISettings::handle(ctx).update(ctx, |settings, ctx| {
|
||||
settings
|
||||
.bedrock_region
|
||||
.set_value("different-region".to_owned(), ctx)
|
||||
.unwrap();
|
||||
})
|
||||
});
|
||||
app.read(|ctx| assert!(!preferences.bedrock_model_is_verified(&model.model_id, ctx)));
|
||||
});
|
||||
}
|
||||
|
||||
@@ -29,7 +29,7 @@ pub struct BedrockClient {
|
||||
region: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct BedrockClientConfig {
|
||||
pub auth_method: BedrockAuthMethod,
|
||||
pub profile: String,
|
||||
@@ -101,7 +101,7 @@ impl BedrockClient {
|
||||
BedrockAuthMethod::Profile | BedrockAuthMethod::Sso => {
|
||||
let mut loader = aws_config::defaults(BehaviorVersion::latest());
|
||||
|
||||
if !config.profile.is_empty() && config.profile != "default" {
|
||||
if !config.profile.is_empty() {
|
||||
loader = loader.profile_name(&config.profile);
|
||||
}
|
||||
|
||||
|
||||
@@ -8,34 +8,46 @@ use aws_config::BehaviorVersion;
|
||||
use aws_sdk_bedrock::Client;
|
||||
use aws_sdk_bedrockruntime::config::Region;
|
||||
use futures::{stream, StreamExt};
|
||||
use galaxy_bedrock_model_catalog::model_metadata;
|
||||
use galaxy_bedrock_model_catalog::models as catalog_models;
|
||||
use serde_json::json;
|
||||
|
||||
use super::client::{BedrockClientConfig, BedrockError};
|
||||
use crate::ai::remote_logging::{RemoteLogLevel, RemoteLogRecord};
|
||||
use crate::settings::ai::BedrockModelConfig;
|
||||
|
||||
pub(crate) struct BedrockDiscovery {
|
||||
pub models: Vec<BedrockModelConfig>,
|
||||
pub records: Vec<RemoteLogRecord>,
|
||||
}
|
||||
|
||||
const AVAILABILITY_CHECK_CONCURRENCY: usize = 8;
|
||||
|
||||
pub async fn discover_available_models(
|
||||
pub(crate) async fn discover_available_models(
|
||||
config: BedrockClientConfig,
|
||||
) -> Result<Vec<BedrockModelConfig>, String> {
|
||||
) -> Result<BedrockDiscovery, String> {
|
||||
let aws_config = load_aws_config(&config)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
let client = Client::new(&aws_config);
|
||||
let catalog = client
|
||||
.list_foundation_models()
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| format!("Could not list AWS Bedrock foundation models: {error}"))?;
|
||||
|
||||
// AWS exposes agreement/authorization state only through an individual
|
||||
// GetFoundationModelAvailability call. Keep using that control-plane API,
|
||||
// but bound the independent checks so a large catalog does not serialize
|
||||
// startup discovery one model at a time.
|
||||
let checks = catalog.model_summaries().iter().cloned().map(|summary| {
|
||||
// The bundled catalog defines candidates; AWS defines account access.
|
||||
// Never add arbitrary entries from ListFoundationModels to this catalog.
|
||||
let candidates = catalog_models().to_vec();
|
||||
let sweep_id = uuid::Uuid::new_v4().to_string();
|
||||
let region = aws_config
|
||||
.region()
|
||||
.map(|region| region.as_ref())
|
||||
.unwrap_or("UNKNOWN")
|
||||
.to_owned();
|
||||
log::info!(
|
||||
"[bedrock] Starting availability sweep {sweep_id}: {} catalog models, region={region}",
|
||||
candidates.len()
|
||||
);
|
||||
let checks = candidates.into_iter().map(|metadata| {
|
||||
let client = client.clone();
|
||||
let sweep_id = sweep_id.clone();
|
||||
let region = region.clone();
|
||||
async move {
|
||||
let model_id = summary.model_id();
|
||||
let model_id = metadata.model_id.as_str();
|
||||
let availability = match client
|
||||
.get_foundation_model_availability()
|
||||
.model_id(model_id)
|
||||
@@ -44,67 +56,59 @@ pub async fn discover_available_models(
|
||||
{
|
||||
Ok(availability) => availability,
|
||||
Err(error) => {
|
||||
log::debug!(
|
||||
"[bedrock] Availability check failed for {model_id}; excluding model: {error}"
|
||||
);
|
||||
return None;
|
||||
let error = crate::ai::remote_logging::sanitize_error(format!("{error:?}"));
|
||||
log::warn!("[bedrock] Availability sweep {sweep_id}: model={model_id}, region={region}, decision=excluded, error={error}");
|
||||
return (None, RemoteLogRecord {
|
||||
level: RemoteLogLevel::Warn,
|
||||
message: "Bedrock model availability check failed".to_owned(),
|
||||
context: json!({"event": "bedrock_model_availability", "sweep_id": sweep_id, "model_id": model_id, "aws_region": region, "decision": "excluded", "error": error}),
|
||||
});
|
||||
}
|
||||
};
|
||||
let agreement = availability.agreement_availability().map(|value| value.status().as_str());
|
||||
let authorization = availability.authorization_status().as_str();
|
||||
let entitlement = availability.entitlement_availability().as_str();
|
||||
let region_status = availability.region_availability().as_str();
|
||||
let access_allowed = model_availability_is_usable(agreement, authorization, entitlement, region_status);
|
||||
let runtime_supported = metadata.supports_agent_runtime() != Some(false);
|
||||
let allowed = access_allowed && runtime_supported;
|
||||
let decision = if allowed { "allowed" } else { "excluded" };
|
||||
log::info!("[bedrock] Availability sweep {sweep_id}: model={model_id}, agreement={}, authorization={authorization}, entitlement={entitlement}, region_status={region_status}, runtime_supported={runtime_supported}, decision={decision}", agreement.unwrap_or("MISSING"));
|
||||
let record = RemoteLogRecord {
|
||||
level: RemoteLogLevel::Info,
|
||||
message: "Bedrock model availability checked".to_owned(),
|
||||
context: json!({"event": "bedrock_model_availability", "sweep_id": sweep_id, "model_id": model_id, "aws_region": region, "agreement": agreement, "authorization": authorization, "entitlement": entitlement, "region_availability": region_status, "agreement_error": availability.agreement_availability().and_then(|value| value.error_message()).map(crate::ai::remote_logging::sanitize_error), "access_allowed": access_allowed, "runtime_supported": runtime_supported, "decision": decision}),
|
||||
};
|
||||
if !allowed { return (None, record); }
|
||||
|
||||
if !model_availability_is_usable(
|
||||
availability
|
||||
.agreement_availability()
|
||||
.map(|agreement| agreement.status().as_str()),
|
||||
availability.authorization_status().as_str(),
|
||||
availability.entitlement_availability().as_str(),
|
||||
availability.region_availability().as_str(),
|
||||
) {
|
||||
log::debug!(
|
||||
"[bedrock] Excluding {model_id}: agreement={}, authorization={}, entitlement={}, region={}",
|
||||
availability
|
||||
.agreement_availability()
|
||||
.map(|agreement| agreement.status().as_str())
|
||||
.unwrap_or("MISSING"),
|
||||
availability.authorization_status().as_str(),
|
||||
availability.entitlement_availability().as_str(),
|
||||
availability.region_availability().as_str(),
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
let metadata = model_metadata(model_id);
|
||||
let display_name = summary.model_name().map(str::to_owned).unwrap_or_else(|| {
|
||||
metadata
|
||||
.map(|metadata| metadata.model_name.clone())
|
||||
.unwrap_or_else(|| prettify_model_id(model_id))
|
||||
});
|
||||
let vision_supported = summary
|
||||
.input_modalities()
|
||||
.iter()
|
||||
.any(|modality| modality.as_str() == "IMAGE")
|
||||
|| metadata.is_some_and(|metadata| metadata.supports_vision_input());
|
||||
|
||||
Some(BedrockModelConfig {
|
||||
(Some(BedrockModelConfig {
|
||||
model_id: model_id.to_owned(),
|
||||
display_name,
|
||||
vision_supported,
|
||||
})
|
||||
display_name: metadata.model_name.clone(),
|
||||
vision_supported: metadata.supports_vision_input(),
|
||||
enabled: true,
|
||||
access_verified: true,
|
||||
}), record)
|
||||
}
|
||||
});
|
||||
let mut models = stream::iter(checks)
|
||||
let results = stream::iter(checks)
|
||||
.buffer_unordered(AVAILABILITY_CHECK_CONCURRENCY)
|
||||
.filter_map(futures::future::ready)
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
|
||||
let (models, mut records): (Vec<_>, Vec<_>) = results.into_iter().unzip();
|
||||
let mut models = models.into_iter().flatten().collect::<Vec<_>>();
|
||||
models.sort_by(|left, right| left.display_name.cmp(&right.display_name));
|
||||
if models.is_empty() {
|
||||
return Err(
|
||||
"AWS returned no Bedrock models that are authorized and available in this region."
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
Ok(models)
|
||||
log::info!(
|
||||
"[bedrock] Availability sweep {sweep_id} complete: checked={}, allowed={}",
|
||||
records.len(),
|
||||
models.len()
|
||||
);
|
||||
records.push(RemoteLogRecord {
|
||||
level: RemoteLogLevel::Info,
|
||||
message: "Bedrock availability sweep completed".to_owned(),
|
||||
context: json!({"event": "bedrock_availability_sweep_completed", "sweep_id": sweep_id, "aws_region": region, "catalog_count": catalog_models().len(), "checked_count": records.len(), "allowed_count": models.len()}),
|
||||
});
|
||||
Ok(BedrockDiscovery { models, records })
|
||||
}
|
||||
|
||||
fn model_availability_is_usable(
|
||||
@@ -119,6 +123,21 @@ fn model_availability_is_usable(
|
||||
&& region_status == "AVAILABLE"
|
||||
}
|
||||
|
||||
pub(crate) fn preserve_model_toggles(
|
||||
existing: &[BedrockModelConfig],
|
||||
mut discovered: Vec<BedrockModelConfig>,
|
||||
) -> Vec<BedrockModelConfig> {
|
||||
for model in &mut discovered {
|
||||
if let Some(previous) = existing
|
||||
.iter()
|
||||
.find(|previous| previous.model_id == model.model_id)
|
||||
{
|
||||
model.enabled = previous.enabled;
|
||||
}
|
||||
}
|
||||
discovered
|
||||
}
|
||||
|
||||
async fn load_aws_config(
|
||||
config: &BedrockClientConfig,
|
||||
) -> Result<aws_config::SdkConfig, BedrockError> {
|
||||
@@ -126,7 +145,7 @@ async fn load_aws_config(
|
||||
crate::settings::ai::BedrockAuthMethod::Profile
|
||||
| crate::settings::ai::BedrockAuthMethod::Sso => {
|
||||
let mut loader = aws_config::defaults(BehaviorVersion::latest());
|
||||
if !config.profile.is_empty() && config.profile != "default" {
|
||||
if !config.profile.is_empty() {
|
||||
loader = loader.profile_name(&config.profile);
|
||||
}
|
||||
if !config.region.is_empty() {
|
||||
@@ -162,14 +181,6 @@ async fn load_aws_config(
|
||||
Ok(sdk_config)
|
||||
}
|
||||
|
||||
fn prettify_model_id(model_id: &str) -> String {
|
||||
model_id
|
||||
.rsplit('.')
|
||||
.next()
|
||||
.unwrap_or(model_id)
|
||||
.replace(['-', ':'], " ")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "discovery_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -1,5 +1,26 @@
|
||||
use super::model_availability_is_usable;
|
||||
|
||||
#[test]
|
||||
fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() {
|
||||
let model = crate::settings::BedrockModelConfig {
|
||||
model_id: "test.model".to_owned(),
|
||||
display_name: "Test".to_owned(),
|
||||
vision_supported: false,
|
||||
enabled: false,
|
||||
access_verified: true,
|
||||
};
|
||||
let saved = serde_json::to_value(&model).unwrap();
|
||||
assert!(saved.get("access_verified").is_none());
|
||||
let restored: crate::settings::BedrockModelConfig = serde_json::from_value(saved).unwrap();
|
||||
assert!(!restored.access_verified);
|
||||
assert!(!restored.enabled);
|
||||
let mut discovered = model.clone();
|
||||
discovered.enabled = true;
|
||||
let refreshed = super::preserve_model_toggles(&[model], vec![discovered]);
|
||||
assert!(!refreshed[0].enabled);
|
||||
assert!(refreshed[0].access_verified);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requires_every_availability_status() {
|
||||
assert!(model_availability_is_usable(
|
||||
|
||||
@@ -141,6 +141,8 @@ fn parse_claude_code_model_map(
|
||||
.map(|(source_id, arn)| {
|
||||
let display_name = derive_display_name(&source_id);
|
||||
BedrockModelConfig {
|
||||
enabled: true,
|
||||
access_verified: false,
|
||||
model_id: arn,
|
||||
display_name,
|
||||
vision_supported: true,
|
||||
@@ -276,8 +278,14 @@ pub fn list_aws_profiles() -> Vec<String> {
|
||||
let mut profiles = BTreeSet::new();
|
||||
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
parse_aws_config_file(home.join(".aws").join("config"), &mut profiles);
|
||||
parse_aws_credentials_file(home.join(".aws").join("credentials"), &mut profiles);
|
||||
let config_path = std::env::var_os("AWS_CONFIG_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| home.join(".aws").join("config"));
|
||||
let credentials_path = std::env::var_os("AWS_SHARED_CREDENTIALS_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| home.join(".aws").join("credentials"));
|
||||
parse_aws_config_file(config_path, &mut profiles);
|
||||
parse_aws_credentials_file(credentials_path, &mut profiles);
|
||||
}
|
||||
|
||||
profiles.into_iter().collect()
|
||||
|
||||
@@ -16,7 +16,7 @@ use std::{
|
||||
|
||||
use chrono::Utc;
|
||||
use galaxy_core::channel::ChannelState;
|
||||
use galaxyui::{Entity, ModelContext, SingletonEntity};
|
||||
use galaxyui::{Entity, ModelContext, SingletonEntity, View, ViewContext};
|
||||
use serde::Serialize;
|
||||
use serde_json::{json, Value};
|
||||
use settings::Setting;
|
||||
@@ -123,6 +123,30 @@ where
|
||||
);
|
||||
}
|
||||
|
||||
/// Views must log without updating AISettings: its singleton cannot be read
|
||||
/// while it is temporarily borrowed by a model update.
|
||||
pub(crate) fn log_view_event<V: View>(ctx: &mut ViewContext<V>, record: RemoteLogRecord) {
|
||||
let Some(config) = RemoteLogConfig::from_settings(AISettings::as_ref(ctx)) else {
|
||||
return;
|
||||
};
|
||||
let payload = RemoteLogPayload {
|
||||
level: record.level.as_str(),
|
||||
message: record.message,
|
||||
service: REMOTE_LOG_SERVICE,
|
||||
context: enrich_context(record.context),
|
||||
};
|
||||
#[cfg(all(feature = "remote_logs_local_file", not(target_family = "wasm")))]
|
||||
write_local_log(&payload);
|
||||
let _ = ctx.spawn(
|
||||
async move { send_remote_log(config, payload).await },
|
||||
|_, result, _| {
|
||||
if let Err(error) = result {
|
||||
log::warn!("[remote-logging] Failed to send remote AI log: {error}");
|
||||
}
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(all(feature = "remote_logs_local_file", not(target_family = "wasm")))]
|
||||
fn write_local_log(payload: &RemoteLogPayload) {
|
||||
let state = LOCAL_LOG_FILE.get_or_init(|| Mutex::new(None));
|
||||
|
||||
Reference in New Issue
Block a user