public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix provider model inventories and Bedrock verification

Browse Source
This commit is contained in:
rkw6086
2026-10-02 16:41:16 -05:00
parent 2796fdd114
commit 188489c890
13 changed files with 1016 additions and 344 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/provider/client.rs
+2 -2
View File
@@ -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);
}
app/src/ai/provider/discovery.rs
+85 -74
View File
@@ -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;
app/src/ai/provider/discovery_tests.rs
+21
View File
@@ -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(
app/src/ai/provider/external_config.rs
+10 -2
View File
@@ -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()