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

remove old okta login authentication refresh

Browse Source
This commit is contained in:
Josh Woodcock
2026-10-14 13:51:35 -05:00
parent 04cff91509
commit 1820561f2a
20 changed files with 566 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/blocklist/controller/response_stream.rs
+3 -1
View File
@@ -188,7 +188,9 @@ impl ResponseStream {
// Fall back to Bedrock
if *settings.bedrock_enabled.value() {
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let config = BedrockClientConfig::from_settings(settings).with_managed_credentials(
let config = BedrockClientConfig::from_settings(settings)
.with_model_route(model_id, settings.bedrock_models.value())
.with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
);
app/src/ai/crosscheck/reviewer.rs
+1
View File
@@ -183,6 +183,7 @@ impl CrosscheckReviewer {
if *settings.bedrock_enabled.value() {
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
.with_model_route(model_id, settings.bedrock_models.value())
.with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
app/src/ai/llms.rs
+41 -44
View File
@@ -27,6 +27,8 @@ use warp_multi_agent_api as api;
use super::custom_model_routers::{self, CustomModelRouter, ModelConfigError};
use super::execution_profiles::profiles::AIExecutionProfilesModel;
use crate::ai::acp::{acp_launch_fingerprint, acp_provider_selection_identity};
#[cfg(not(target_family = "wasm"))]
use crate::ai::provider::external_config::ExternalBedrockConfig;
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
use crate::auth::AuthStateProvider;
use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind};
@@ -616,6 +618,7 @@ pub struct LLMPreferences {
#[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: bool,
bedrock_models_refresh_in_flight: bool,
bedrock_verification_generation: u64,
#[cfg(not(target_family = "wasm"))]
bedrock_verified_config: Option<super::provider::client::BedrockClientConfig>,
#[cfg(not(target_family = "wasm"))]
@@ -796,6 +799,7 @@ impl LLMPreferences {
#[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
#[cfg(not(target_family = "wasm"))]
bedrock_verified_config: None,
#[cfg(not(target_family = "wasm"))]
@@ -844,7 +848,9 @@ impl LLMPreferences {
return;
}
self.bedrock_models_refresh_in_flight = true;
let generation = self.begin_bedrock_verification();
let config = bedrock_discovery_config(settings);
let revision = super::provider::authentication::connection_revision(ctx);
let requested_config = config.clone();
self.bedrock_verified_config = None;
self.bedrock_verified_model_ids.clear();
@@ -860,6 +866,11 @@ impl LLMPreferences {
discovery.models
});
me.bedrock_models_refresh_in_flight = false;
if me.bedrock_verification_generation != generation
|| super::provider::authentication::connection_revision(ctx) != revision
{
return;
}
if !*AISettings::as_ref(ctx).bedrock_enabled.value() {
return;
}
@@ -888,26 +899,37 @@ impl LLMPreferences {
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
}
Err(error) => {
log::debug!("[bedrock] Startup model discovery unavailable: {error}");
log::warn!("[bedrock] Startup model discovery unavailable: {error}");
}
}
},
);
}
#[cfg(not(target_family = "wasm"))]
pub(crate) fn begin_bedrock_verification(&mut self) -> u64 {
self.bedrock_verification_generation += 1;
self.bedrock_verification_generation
}
#[cfg(not(target_family = "wasm"))]
pub(crate) fn record_bedrock_verification(
&mut self,
generation: u64,
config: super::provider::client::BedrockClientConfig,
models: &[BedrockModelConfig],
ctx: &mut ModelContext<Self>,
) {
) -> bool {
if self.bedrock_verification_generation != generation {
return false;
}
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);
true
}
#[cfg(not(target_family = "wasm"))]
@@ -945,48 +967,14 @@ impl LLMPreferences {
let region = settings.bedrock_region.value().clone();
let cross_region = *settings.bedrock_cross_region_inference.value();
// Check if user wants only 1-hour cache models
use crate::ai::provider::external_config::ExternalBedrockConfig;
let external_config = ExternalBedrockConfig::load();
let require_1h_cache = external_config.enable_prompt_caching_1h;
let mut effective = discovered_models;
// Filter out models that don't support 1-hour caching if required
if require_1h_cache {
effective.retain(|model| {
// 1-hour caching is supported by Claude 4.5+ models
// Opus 4.5+, Sonnet 4.5+, Haiku 4.5+
let supports_1h = model.model_id.contains("4-5")
|| model.model_id.contains("4.5")
|| model.model_id.contains("-4-6") // Opus/Sonnet 4.6+ also support 1h
|| model.model_id.contains("4.6")
|| model.model_id.contains("-4-7")
|| model.model_id.contains("4.7")
|| model.model_id.contains("-4-8")
|| model.model_id.contains("4.8");
if !supports_1h {
log::info!(
"[bedrock] Filtering out model {} - does not support 1-hour cache (ENABLE_PROMPT_CACHING_1H=1)",
model.model_id
);
}
supports_1h
});
if effective.is_empty() {
log::warn!("[bedrock] No models left after filtering for 1-hour cache support!");
}
}
let effective = effective;
for model in effective {
for model in discovered_models {
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)) {
if model.runtime_supported == Some(false)
|| metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false))
{
log::debug!(
"[bedrock] Excluding {} from Agent Mode because its catalog metadata marks it as incompatible with streaming Converse",
model.model_id
@@ -1046,7 +1034,6 @@ impl LLMPreferences {
// Default agent mode to ANTHROPIC_MODEL from external config if set,
// otherwise Claude Opus 4.6, falling back to the first available model.
// Note: external_config already loaded above for filtering
let external_config_for_default = ExternalBedrockConfig::load();
if let Some(id) = external_config_for_default
@@ -1082,7 +1069,11 @@ impl LLMPreferences {
.map(|m| m.id.clone())
})
{
self.models_by_feature.agent_mode.default_id = id;
if !self.models_by_feature.agent_mode.choices.iter().any(|model| {
model.id == self.models_by_feature.agent_mode.default_id
}) {
self.models_by_feature.agent_mode.default_id = id;
}
}
// Default coding to Claude Sonnet 4.6, falling back to the first available model.
@@ -1095,12 +1086,18 @@ impl LLMPreferences {
.or_else(|| self.models_by_feature.coding.choices.first())
.map(|m| m.id.clone())
{
self.models_by_feature.coding.default_id = id;
if !self.models_by_feature.coding.choices.iter().any(|model| {
model.id == self.models_by_feature.coding.default_id
}) {
self.models_by_feature.coding.default_id = id;
}
}
// Default CLI agent to the first available model.
if let Some(ref mut cli) = self.models_by_feature.cli_agent {
if let Some(id) = cli.choices.first().map(|m| m.id.clone()) {
if let Some(id) = cli.choices.first().filter(|_| {
!cli.choices.iter().any(|model| model.id == cli.default_id)
}).map(|m| m.id.clone()) {
cli.default_id = id;
}
}
app/src/ai/llms_tests.rs
+48
View File
@@ -159,6 +159,51 @@ fn openai_model(model_id: &str) -> OpenAIModelConfig {
}
}
#[test]
fn live_bedrock_models_are_not_hidden_by_missing_catalog_metadata() {
App::test((), |mut app| async move {
initialize_settings_for_tests(&mut app);
let id = "future.new-live-model";
AISettings::handle(&app).update(&mut app, |settings, ctx| {
settings.bedrock_enabled.set_value(true, ctx).unwrap();
settings.bedrock_models.set_value(
vec![bedrock_model("future.other-model"), bedrock_model(id)], ctx,
).unwrap();
});
let mut preferences = empty_preferences();
preferences.models_by_feature.agent_mode.default_id = id.into();
preferences.models_by_feature.coding.default_id = id.into();
app.read(|ctx| {
preferences.bedrock_verified_config =
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
preferences.bedrock_verified_model_ids.insert(id.to_owned());
preferences.bedrock_verified_model_ids.insert("future.other-model".to_owned());
preferences.inject_bedrock_models(ctx);
});
assert!(preferences.models_by_feature.agent_mode.choices.iter().any(|model| model.id.as_str() == id));
assert_eq!(preferences.models_by_feature.agent_mode.default_id.as_str(), id);
assert_eq!(preferences.models_by_feature.coding.default_id.as_str(), id);
});
}
#[test]
fn newer_bedrock_discovery_invalidates_older_results() {
App::test((), |mut app| async move {
initialize_settings_for_tests(&mut app);
let preferences = app.add_model(|_| empty_preferences());
let models = [bedrock_model("future.new-live-model")];
preferences.update(&mut app, |preferences, ctx| {
let first = preferences.begin_bedrock_verification();
let latest = preferences.begin_bedrock_verification();
let config = super::bedrock_discovery_config(AISettings::as_ref(ctx));
assert!(!preferences.record_bedrock_verification(first, config.clone(), &models, ctx));
assert!(preferences.bedrock_verified_model_ids.is_empty());
assert!(preferences.record_bedrock_verification(latest, config, &models, ctx));
assert!(preferences.bedrock_verified_model_ids.contains("future.new-live-model"));
});
});
}
fn bedrock_model(model_id: &str) -> BedrockModelConfig {
BedrockModelConfig {
enabled: true,
@@ -166,6 +211,7 @@ fn bedrock_model(model_id: &str) -> BedrockModelConfig {
model_id: model_id.to_string(),
display_name: model_id.to_string(),
vision_supported: false,
..Default::default()
}
}
@@ -235,6 +281,7 @@ fn empty_preferences() -> LLMPreferences {
catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(),
@@ -451,6 +498,7 @@ fn acp_models_are_injected_only_while_acp_is_enabled() {
catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(),
app/src/ai/provider/authentication.rs
+17
View File
@@ -22,10 +22,27 @@ impl BedrockClientConfig {
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
}
}
pub(crate) fn with_model_route(
mut self,
model_id: &str,
models: &[crate::settings::BedrockModelConfig],
) -> Self {
self.invocation_model_id = models.iter().find(|model| {
model.model_id == model_id
|| model.invocation_model_id == model_id
|| super::models::foundation_model_id(&model.model_id)
== super::models::foundation_model_id(model_id)
}).and_then(|model| {
(!model.invocation_model_id.is_empty()).then(|| model.invocation_model_id.clone())
});
self
}
pub(crate) fn with_managed_credentials(
mut self,
strategy: &AwsCredentialsRefreshStrategy,
app/src/ai/provider/authentication_tests.rs
+17
View File
@@ -13,6 +13,7 @@ fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig {
secret_access_key: "configured-secret".to_owned(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: true,
}
}
@@ -144,6 +145,22 @@ fn credential_expiry_is_checked_at_the_boundary() {
assert!(!credentials_are_current(Some(now - Duration::from_secs(1)), now));
}
#[test]
fn model_routes_use_discovered_ids_and_match_legacy_cross_region_selections() {
let model = crate::settings::BedrockModelConfig {
model_id: "future.new-model".to_owned(),
invocation_model_id: "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.future.new-model".to_owned(),
..Default::default()
};
for selected in ["future.new-model", "us.future.new-model"] {
let resolved = config(BedrockAuthMethod::Profile)
.with_model_route(selected, &[model.clone()]);
assert_eq!(resolved.invocation_model_id.as_deref(), Some(model.invocation_model_id.as_str()));
}
let resolved = config(BedrockAuthMethod::Profile).with_model_route("unrelated.model", &[model]);
assert_eq!(resolved.invocation_model_id, None);
}
fn test_credentials() -> AwsCredentials {
AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None)
}
app/src/ai/provider/client.rs
+7 -2
View File
@@ -28,6 +28,7 @@ fn strip_context_marker(model_id: &str) -> String {
pub struct BedrockClient {
runtime_client: BedrockRuntimeClient,
region: String,
invocation_model_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -39,6 +40,7 @@ pub struct BedrockClientConfig {
pub secret_access_key: String,
pub session_token: Option<String>,
pub credentials_expiry: Option<SystemTime>,
pub invocation_model_id: Option<String>,
pub cross_region_inference: bool,
}
@@ -155,6 +157,7 @@ impl BedrockClient {
Ok(Self {
runtime_client,
region,
invocation_model_id: config.invocation_model_id,
})
}
@@ -185,14 +188,16 @@ impl BedrockClient {
tool_result_archive: Vec<ConversationMessage>,
) -> Result<LegacyResponseStream, BedrockError> {
let base_model_id = strip_context_marker(model_id);
let effective_model_id = if cross_region_inference {
let effective_model_id = if let Some(model) = &self.invocation_model_id {
model.clone()
} else if cross_region_inference {
apply_cross_region_prefix(&base_model_id, &self.region)
} else {
base_model_id
};
let external_config = ExternalBedrockConfig::load();
let caching_config = CachingConfig::from_external_config(&external_config);
let caching_config = CachingConfig::from_external_config(&external_config).for_model(model_id);
if !caching_config.enabled {
log::info!("[bedrock] Prompt caching disabled (DISABLE_PROMPT_CACHING=1)");
app/src/ai/provider/convert.rs
+16
View File
@@ -38,6 +38,22 @@ impl CachingConfig {
extended_ttl_requested: external.enable_prompt_caching_1h,
}
}
pub fn for_model(mut self, model_id: &str) -> Self {
let model = galaxy_bedrock_model_catalog::model_metadata(model_id);
let known_claude_4 = model.is_some_and(|metadata| {
metadata.provider_name == "Anthropic"
&& metadata.supports_agent_runtime() == Some(true)
&& (metadata.model_id.starts_with("anthropic.claude-opus-4")
|| metadata.model_id.starts_with("anthropic.claude-sonnet-4")
|| metadata.model_id.starts_with("anthropic.claude-haiku-4"))
});
self.enabled &= known_claude_4;
self.extended_ttl_requested &= self.enabled && model.is_some_and(|metadata| {
["4-5", "4-6", "4-7", "4-8"].iter().any(|version| metadata.model_id.contains(version))
});
self
}
}
pub struct ConvertedRequest {
app/src/ai/provider/convert_tests.rs
+12
View File
@@ -1,3 +1,15 @@
#[test]
fn cache_preferences_do_not_enable_unsupported_or_unknown_model_caching() {
let external = super::external_config::ExternalBedrockConfig {
enable_prompt_caching_1h: true,
..Default::default()
};
for id in ["future.new-model-v100", "openai.gpt-5.6-sol"] {
let config = super::convert::CachingConfig::from_external_config(&external).for_model(id);
assert!(!config.enabled);
assert!(!config.extended_ttl_requested);
}
}
use aws_sdk_bedrockruntime::types::{ContentBlock, ConversationRole, SystemContentBlock, Tool};
use serde_json::json;
app/src/ai/provider/discovery.rs
+168 -25
View File
@@ -1,17 +1,25 @@
//! AWS Bedrock control-plane discovery.
//!
//! The foundation-model catalog is only a candidate list. Every candidate is
//! checked with `GetFoundationModelAvailability` before it is offered to the
//! user or persisted in Galaxy settings.
//! AWS supplies the model inventory and invocation routes. Bundled metadata
//! enriches known capabilities but never determines inventory membership.
use std::collections::{BTreeMap, HashSet};
use std::future::Future;
use aws_config::BehaviorVersion;
use aws_sdk_bedrock::operation::list_inference_profiles::ListInferenceProfilesOutput;
use aws_sdk_bedrock::types::{
FoundationModelSummary, InferenceProfileStatus, InferenceProfileSummary,
InferenceProfileType,
};
use aws_sdk_bedrock::Client;
use aws_sdk_bedrockruntime::config::Region;
use futures::{stream, StreamExt};
use galaxy_bedrock_model_catalog::models as catalog_models;
use galaxy_bedrock_model_catalog::model_metadata;
use serde_json::json;
use super::client::{BedrockClientConfig, BedrockError};
use super::models::foundation_model_id;
use crate::ai::remote_logging::{RemoteLogLevel, RemoteLogRecord};
use crate::settings::ai::BedrockModelConfig;
@@ -29,9 +37,66 @@ pub(crate) async fn discover_available_models(
.await
.map_err(|error| error.to_string())?;
let client = Client::new(&aws_config);
// 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 inventory = client.list_foundation_models().send().await
.map_err(|error| format!("Unable to list Bedrock models: {}", sanitized_error(&error)))?;
let mut foundation_models = inventory.model_summaries.unwrap_or_default();
let profiles = if config.cross_region_inference {
collect_profile_pages(|token| {
let client = client.clone();
async move {
client.list_inference_profiles()
.set_next_token(token)
.send().await
.map_err(|error| format!("Unable to list Bedrock inference profiles: {}", sanitized_error(&error)))
}
}).await?
} else {
Vec::new()
};
let mut known_ids: HashSet<String> = foundation_models.iter()
.map(|model| model.model_id().to_owned()).collect();
// A profile can introduce a model before it appears in the regional listing.
for profile in profiles.iter().filter(|profile| profile.status() == &InferenceProfileStatus::Active) {
for model in profile.models() {
let Some(arn) = model.model_arn() else { continue };
if !arn.contains(":foundation-model/") {
continue;
}
let id = foundation_model_id(arn);
if !known_ids.insert(id.to_owned()) {
continue;
}
let details = match client.get_foundation_model().model_identifier(id).send().await {
Ok(output) => match output.model_details {
Some(details) => details,
None => {
log::warn!("[bedrock] AWS returned no metadata for profile model {id}");
continue;
}
},
Err(error) => {
let error = sanitized_error(&error);
log::warn!("[bedrock] Unable to describe profile model {id}: {error}");
continue;
}
};
let summary = FoundationModelSummary::builder()
.model_id(details.model_id)
.model_arn(details.model_arn)
.set_model_name(details.model_name)
.set_provider_name(details.provider_name)
.set_input_modalities(details.input_modalities)
.set_output_modalities(details.output_modalities)
.set_response_streaming_supported(details.response_streaming_supported)
.set_inference_types_supported(details.inference_types_supported)
.set_model_lifecycle(details.model_lifecycle)
.build().map_err(|error| format!("Invalid Bedrock model metadata: {error}"))?;
foundation_models.push(summary);
}
}
let inventory_count = foundation_models.len();
let candidates = models_from_inventory(foundation_models, &profiles, config.cross_region_inference);
let candidate_count = candidates.len();
let sweep_id = uuid::Uuid::new_v4().to_string();
let region = aws_config
.region()
@@ -39,30 +104,30 @@ pub(crate) async fn discover_available_models(
.unwrap_or("UNKNOWN")
.to_owned();
log::info!(
"[bedrock] Starting availability sweep {sweep_id}: {} catalog models, region={region}",
"[bedrock] Starting availability sweep {sweep_id}: {} live candidates, region={region}",
candidates.len()
);
let checks = candidates.into_iter().map(|metadata| {
let checks = candidates.into_iter().map(|mut model| {
let client = client.clone();
let sweep_id = sweep_id.clone();
let region = region.clone();
async move {
let model_id = metadata.model_id.as_str();
let model_id = model.model_id.clone();
let availability = match client
.get_foundation_model_availability()
.model_id(model_id)
.model_id(&model_id)
.send()
.await
{
Ok(availability) => availability,
Err(error) => {
let error = crate::ai::remote_logging::sanitize_error(format!("{error:?}"));
let error = sanitized_error(&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}),
});
}, false);
}
};
let agreement = availability.agreement_availability().map(|value| value.status().as_str());
@@ -70,7 +135,7 @@ pub(crate) async fn discover_available_models(
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 runtime_supported = model.runtime_supported != 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"));
@@ -79,23 +144,20 @@ pub(crate) async fn discover_available_models(
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); }
(Some(BedrockModelConfig {
model_id: model_id.to_owned(),
display_name: metadata.model_name.clone(),
vision_supported: metadata.supports_vision_input(),
enabled: true,
access_verified: true,
}), record)
model.access_verified = allowed;
(allowed.then_some(model), record, true)
}
});
let results = stream::iter(checks)
.buffer_unordered(AVAILABILITY_CHECK_CONCURRENCY)
.collect::<Vec<_>>()
.await;
if candidate_count > 0 && !results.iter().any(|(_, _, checked)| *checked) {
return Err("All Bedrock availability checks failed; saved model preferences are unchanged".to_owned());
}
let (models, mut records): (Vec<_>, Vec<_>) = results.into_iter().unzip();
let (models, mut records): (Vec<_>, Vec<_>) = results.into_iter()
.map(|(model, record, _)| (model, record)).unzip();
let mut models = models.into_iter().flatten().collect::<Vec<_>>();
models.sort_by(|left, right| left.display_name.cmp(&right.display_name));
log::info!(
@@ -106,11 +168,84 @@ pub(crate) async fn discover_available_models(
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()}),
context: json!({"event": "bedrock_availability_sweep_completed", "sweep_id": sweep_id, "aws_region": region, "inventory_count": inventory_count, "candidate_count": candidate_count, "checked_count": records.len(), "allowed_count": models.len()}),
});
Ok(BedrockDiscovery { models, records })
}
fn sanitized_error(error: &impl std::fmt::Debug) -> String {
crate::ai::remote_logging::sanitize_error(format!("{error:?}"))
}
async fn collect_profile_pages<F, Fut>(mut fetch: F) -> Result<Vec<InferenceProfileSummary>, String>
where
F: FnMut(Option<String>) -> Fut,
Fut: Future<Output = Result<ListInferenceProfilesOutput, String>>,
{
let mut profiles = Vec::new();
let mut token = None;
let mut seen = HashSet::new();
loop {
let page = fetch(token).await?;
profiles.extend(page.inference_profile_summaries.unwrap_or_default());
token = page.next_token.filter(|token| !token.is_empty());
let Some(next) = &token else { return Ok(profiles) };
if !seen.insert(next.clone()) {
return Err("AWS repeated an inference-profile pagination token".to_owned());
}
}
}
fn models_from_inventory(
foundation_models: Vec<FoundationModelSummary>,
profiles: &[InferenceProfileSummary],
cross_region: bool,
) -> Vec<BedrockModelConfig> {
let mut routes = profiles.iter()
.filter(|profile| profile.status() == &InferenceProfileStatus::Active)
.collect::<Vec<_>>();
routes.sort_by_key(|profile| (
profile.r#type() != &InferenceProfileType::SystemDefined,
profile.inference_profile_id().starts_with("global."),
profile.inference_profile_id().to_owned(),
));
let mut models = BTreeMap::new();
for summary in foundation_models {
let id = summary.model_id();
let metadata = model_metadata(id);
let text_output = summary.output_modalities().iter().any(|modality| modality.as_str() == "TEXT");
let incompatible = (!summary.output_modalities().is_empty() && !text_output)
|| summary.response_streaming_supported() == Some(false)
|| metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false));
if incompatible {
continue;
}
let profile = cross_region.then(|| routes.iter().find(|profile| {
!profile.models().is_empty() && profile.models().iter().all(|model| {
model.model_arn().is_some_and(|arn| foundation_model_id(arn) == id)
})
})).flatten();
let invocation = if let Some(profile) = profile {
profile.inference_profile_arn().to_owned()
} else if summary.inference_types_supported().iter().any(|kind| kind.as_str() == "ON_DEMAND") {
id.to_owned()
} else {
continue;
};
models.entry(id.to_owned()).or_insert_with(|| BedrockModelConfig {
model_id: id.to_owned(),
invocation_model_id: invocation,
display_name: summary.model_name().unwrap_or(id).to_owned(),
provider_name: summary.provider_name().unwrap_or_default().to_owned(),
vision_supported: summary.input_modalities().iter().any(|modality| modality.as_str() == "IMAGE"),
response_streaming_supported: summary.response_streaming_supported(),
runtime_supported: (text_output && summary.response_streaming_supported() == Some(true)).then_some(true),
..Default::default()
});
}
models.into_values().collect()
}
fn model_availability_is_usable(
agreement_status: Option<&str>,
authorization_status: &str,
@@ -133,6 +268,14 @@ pub(crate) fn preserve_model_toggles(
.find(|previous| previous.model_id == model.model_id)
{
model.enabled = previous.enabled;
model.display_name.clone_from(&previous.display_name);
}
}
for previous in existing {
if !discovered.iter().any(|model| model.model_id == previous.model_id) {
let mut preserved = previous.clone();
preserved.access_verified = false;
discovered.push(preserved);
}
}
discovered
app/src/ai/provider/discovery_tests.rs
+140
View File
@@ -1,4 +1,143 @@
use super::model_availability_is_usable;
use aws_sdk_bedrock::operation::list_inference_profiles::ListInferenceProfilesOutput;
use aws_sdk_bedrock::types::{
FoundationModelSummary, InferenceProfileModel, InferenceProfileStatus,
InferenceProfileSummary, InferenceProfileType, InferenceType, ModelModality,
};
fn foundation(id: &str, on_demand: bool) -> FoundationModelSummary {
FoundationModelSummary::builder()
.model_id(id)
.model_arn(format!("arn:aws:bedrock:us-east-1::foundation-model/{id}"))
.model_name("New live model")
.provider_name("New provider")
.input_modalities(ModelModality::Text)
.output_modalities(ModelModality::Text)
.response_streaming_supported(true)
.set_inference_types_supported(on_demand.then(|| vec![InferenceType::OnDemand]))
.build()
.unwrap()
}
fn inference_profile(id: &str, model: &str) -> InferenceProfileSummary {
InferenceProfileSummary::builder()
.inference_profile_id(id)
.inference_profile_name(id)
.inference_profile_arn(format!("arn:aws:bedrock:us-east-1:123456789012:inference-profile/{id}"))
.status(InferenceProfileStatus::Active)
.r#type(InferenceProfileType::SystemDefined)
.models(InferenceProfileModel::builder()
.model_arn(format!("arn:aws:bedrock:us-east-1::foundation-model/{model}"))
.build())
.build()
.unwrap()
}
#[test]
fn new_live_model_does_not_need_a_bundled_catalog_entry() {
let id = "future.new-model-v100";
assert!(galaxy_bedrock_model_catalog::model_metadata(id).is_none());
let profile = inference_profile("us.future.new-model-v100", id);
let models = super::models_from_inventory(
vec![foundation(id, false)], &[profile.clone()], true,
);
assert_eq!(models.len(), 1);
assert_eq!(models[0].model_id, id);
assert_eq!(models[0].invocation_model_id, profile.inference_profile_arn());
assert_eq!(models[0].response_streaming_supported, Some(true));
assert_eq!(models[0].provider_name, "New provider");
}
#[test]
fn cross_region_disabled_requires_an_actual_on_demand_route() {
let id = "future.profile-only";
let profile = inference_profile("us.future.profile-only", id);
assert!(super::models_from_inventory(
vec![foundation(id, false)], &[profile.clone()], false,
).is_empty());
let models = super::models_from_inventory(
vec![foundation(id, true)], &[profile], false,
);
assert_eq!(models.len(), 1);
assert_eq!(models[0].invocation_model_id, id);
}
#[test]
fn inventory_deduplicates_routes_and_rejects_known_incompatible_models() {
let id = "future.new-model-v100";
let profile = inference_profile("us.future.new-model-v100", id);
let models = super::models_from_inventory(
vec![foundation(id, false), foundation(id, false)],
&[profile.clone(), profile],
true,
);
assert_eq!(models.len(), 1);
assert!(super::models_from_inventory(
vec![foundation("openai.gpt-5.5", true)], &[], true,
).is_empty());
let mut not_streaming = foundation(id, true);
not_streaming.response_streaming_supported = Some(false);
assert!(super::models_from_inventory(vec![not_streaming], &[], true).is_empty());
let mut image_model = foundation(id, true);
image_model.output_modalities = Some(vec![ModelModality::Image]);
assert!(super::models_from_inventory(vec![image_model], &[], true).is_empty());
}
#[tokio::test]
async fn inference_profile_discovery_follows_every_page() {
let mut calls = Vec::new();
let profiles = super::collect_profile_pages(|token| {
calls.push(token.clone());
let page = match token.as_deref() {
None => ListInferenceProfilesOutput::builder()
.inference_profile_summaries(inference_profile("us.future.first", "future.first"))
.next_token("page-2").build(),
Some("page-2") => ListInferenceProfilesOutput::builder()
.inference_profile_summaries(inference_profile("us.future.second", "future.second"))
.build(),
Some(token) => panic!("Unexpected token: {token}"),
};
std::future::ready(Ok(page))
}).await.unwrap();
assert_eq!(profiles.len(), 2);
assert_eq!(calls, [None, Some("page-2".to_owned())]);
}
#[tokio::test]
async fn failed_or_repeating_profile_pages_do_not_return_partial_success() {
let result = super::collect_profile_pages(|_| {
std::future::ready(Err("network unavailable".to_owned()))
}).await;
assert_eq!(result.unwrap_err(), "network unavailable");
let result = super::collect_profile_pages(|_| {
std::future::ready(Ok(ListInferenceProfilesOutput::builder().next_token("repeated").build()))
}).await;
assert!(result.is_err());
}
#[test]
fn transiently_missing_models_keep_preferences_without_restoring_authorization() {
let mut saved = crate::settings::BedrockModelConfig {
enabled: false,
access_verified: true,
model_id: "future.existing".to_owned(),
display_name: "My model name".to_owned(),
..Default::default()
};
let preserved = super::preserve_model_toggles(&[saved.clone()], Vec::new());
assert_eq!(preserved.len(), 1);
assert!(!preserved[0].enabled);
assert!(!preserved[0].access_verified);
saved.invocation_model_id = "old-route".to_owned();
let mut discovered = saved.clone();
discovered.enabled = true;
discovered.display_name = "AWS name".to_owned();
discovered.invocation_model_id = "new-route".to_owned();
let preserved = super::preserve_model_toggles(&[saved], vec![discovered]);
assert!(!preserved[0].enabled);
assert_eq!(preserved[0].display_name, "My model name");
assert_eq!(preserved[0].invocation_model_id, "new-route");
}
#[test]
fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() {
@@ -8,6 +147,7 @@ fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() {
vision_supported: false,
enabled: false,
access_verified: true,
..Default::default()
};
let saved = serde_json::to_value(&model).unwrap();
assert!(saved.get("access_verified").is_none());
app/src/ai/provider/e2e_tests.rs
+1
View File
@@ -269,6 +269,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false,
})
}
app/src/ai/provider/external_config.rs
+1
View File
@@ -146,6 +146,7 @@ fn parse_claude_code_model_map(
model_id: arn,
display_name,
vision_supported: true,
..Default::default()
}
})
.collect()
app/src/ai/provider/integration_tests.rs
+1
View File
@@ -24,6 +24,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false,
})
}
app/src/ai/provider/models.rs
+8
View File
@@ -1,3 +1,11 @@
pub(crate) fn foundation_model_id(model_id: &str) -> &str {
let id = model_id.rsplit('/').next().unwrap_or(model_id);
let id = id.strip_suffix("[1m]").or_else(|| id.strip_suffix("[1M]")).unwrap_or(id);
["global.", "us.", "eu.", "apac.", "jp.", "au."]
.iter()
.find_map(|prefix| id.strip_prefix(prefix))
.unwrap_or(id)
}
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") {
return model_id.to_string();
app/src/ai/runtime/rig.rs
+4 -3
View File
@@ -346,14 +346,15 @@ pub(crate) async fn provider_runtime_for_request(
},
crate::ai::provider::ProviderConfig::Bedrock(config) => {
let max_output_tokens = Some(bedrock_max_output_tokens(&model));
let cross_region_inference = config.cross_region_inference;
let cross_region_inference = config.invocation_model_id.is_none() && config.cross_region_inference;
let caching_config =
CachingConfig::from_external_config(&ExternalBedrockConfig::load());
CachingConfig::from_external_config(&ExternalBedrockConfig::load()).for_model(&model);
let invocation_model = config.invocation_model_id.clone().unwrap_or(model);
let client = BedrockClient::from_config(config).await?;
Arc::new(BedrockRuntime::from_aws_client(
client.runtime_client(),
BedrockRigConfig {
model,
model: invocation_model,
region: client.region().to_string(),
cross_region_inference,
prompt_caching: caching_config.enabled,