diff --git a/AGENTS.md b/AGENTS.md index a86543e9..ee8875b3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -143,7 +143,7 @@ ACP: controller.rs → ResponseStream → acp_output_stream (session-owned lifec - `convert.rs` — Bedrock request construction and prompt-caching behavior - `client.rs` — AWS SDK client construction, runtime creation, and independent background streaming calls - `models.rs` — Model registry and cross-region inference prefix logic -- `discovery.rs` — Checks bundled `galaxy_bedrock_model_catalog` candidates with AWS `GetFoundationModelAvailability`; agreement, authorization, entitlement, and regional availability must all pass. Cached access is never restored as authorization; per-model enablement survives successful refreshes. Sweeps run once at startup when enabled and once per explicit connection Test; settings saves, reloads, and model toggles must never trigger another sweep. Live verification belongs to `LLMPreferences` and is scoped to the exact connection configuration. +- `discovery.rs` — Discovers the live inventory with AWS `ListFoundationModels` and paginated `ListInferenceProfiles`, resolving profile-only model details with `GetFoundationModel`. The bundled `galaxy_bedrock_model_catalog` is optional capability/token-budget enrichment, never an inventory allowlist. Invocation uses AWS-returned model IDs/profile ARNs rather than guessed prefixes. `GetFoundationModelAvailability` still requires agreement, authorization, entitlement, and regional availability; cached access is never restored as authorization. Saved model enablement and names survive refreshes, and failed discovery must not overwrite preferences with an empty catalog. Sweeps run once at startup when enabled and once per explicit connection Test; settings saves, credential reloads, login completion, and model toggles must never trigger another sweep. Live verification belongs to `LLMPreferences` and is scoped to the exact connection configuration. Prompt-cache preferences must never hide models; unknown cache capabilities use no caching. - `diagnostic.rs` — Debug logging (enabled via `GALAXY_BEDROCK_DIAGNOSTICS=1`) - `external_config.rs` — Fallback config from Claude Code/OpenCode settings diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 7e7340c1..ea77f7d9 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -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(), ); diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 73ebb556..c1aa4d82 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -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(), diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 8779a857..5fb6c1a6 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -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, #[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, - ) { + ) -> 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; } } diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index e8546ff8..c0a83ce8 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -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(), diff --git a/app/src/ai/provider/authentication.rs b/app/src/ai/provider/authentication.rs index 6431588e..6f555ef3 100644 --- a/app/src/ai/provider/authentication.rs +++ b/app/src/ai/provider/authentication.rs @@ -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, diff --git a/app/src/ai/provider/authentication_tests.rs b/app/src/ai/provider/authentication_tests.rs index 3b7d2ab1..7b463707 100644 --- a/app/src/ai/provider/authentication_tests.rs +++ b/app/src/ai/provider/authentication_tests.rs @@ -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) } diff --git a/app/src/ai/provider/client.rs b/app/src/ai/provider/client.rs index 4976e593..82fe4266 100644 --- a/app/src/ai/provider/client.rs +++ b/app/src/ai/provider/client.rs @@ -28,6 +28,7 @@ fn strip_context_marker(model_id: &str) -> String { pub struct BedrockClient { runtime_client: BedrockRuntimeClient, region: String, + invocation_model_id: Option, } #[derive(Debug, Clone, PartialEq)] @@ -39,6 +40,7 @@ pub struct BedrockClientConfig { pub secret_access_key: String, pub session_token: Option, pub credentials_expiry: Option, + pub invocation_model_id: Option, 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, ) -> Result { 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)"); diff --git a/app/src/ai/provider/convert.rs b/app/src/ai/provider/convert.rs index dd3b86b8..b57e3af3 100644 --- a/app/src/ai/provider/convert.rs +++ b/app/src/ai/provider/convert.rs @@ -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 { diff --git a/app/src/ai/provider/convert_tests.rs b/app/src/ai/provider/convert_tests.rs index 25011584..b44b8cb7 100644 --- a/app/src/ai/provider/convert_tests.rs +++ b/app/src/ai/provider/convert_tests.rs @@ -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; diff --git a/app/src/ai/provider/discovery.rs b/app/src/ai/provider/discovery.rs index 2fac4e1b..86cae58f 100644 --- a/app/src/ai/provider/discovery.rs +++ b/app/src/ai/provider/discovery.rs @@ -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 = 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::>() .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::>(); 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(mut fetch: F) -> Result, String> +where + F: FnMut(Option) -> Fut, + Fut: Future>, +{ + 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, + profiles: &[InferenceProfileSummary], + cross_region: bool, +) -> Vec { + let mut routes = profiles.iter() + .filter(|profile| profile.status() == &InferenceProfileStatus::Active) + .collect::>(); + 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 diff --git a/app/src/ai/provider/discovery_tests.rs b/app/src/ai/provider/discovery_tests.rs index 560f5f45..1c9ed080 100644 --- a/app/src/ai/provider/discovery_tests.rs +++ b/app/src/ai/provider/discovery_tests.rs @@ -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()); diff --git a/app/src/ai/provider/e2e_tests.rs b/app/src/ai/provider/e2e_tests.rs index 434739dc..51019c74 100644 --- a/app/src/ai/provider/e2e_tests.rs +++ b/app/src/ai/provider/e2e_tests.rs @@ -269,6 +269,7 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, credentials_expiry: None, + invocation_model_id: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/external_config.rs b/app/src/ai/provider/external_config.rs index 71bc8145..d9b5b2fa 100644 --- a/app/src/ai/provider/external_config.rs +++ b/app/src/ai/provider/external_config.rs @@ -146,6 +146,7 @@ fn parse_claude_code_model_map( model_id: arn, display_name, vision_supported: true, + ..Default::default() } }) .collect() diff --git a/app/src/ai/provider/integration_tests.rs b/app/src/ai/provider/integration_tests.rs index 5a52394a..82ff8c69 100644 --- a/app/src/ai/provider/integration_tests.rs +++ b/app/src/ai/provider/integration_tests.rs @@ -24,6 +24,7 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, credentials_expiry: None, + invocation_model_id: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/models.rs b/app/src/ai/provider/models.rs index 9c9c69be..df3d777c 100644 --- a/app/src/ai/provider/models.rs +++ b/app/src/ai/provider/models.rs @@ -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(); diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 85140b4a..61172595 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -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, diff --git a/app/src/settings/ai.rs b/app/src/settings/ai.rs index be2b2c1e..04474c65 100644 --- a/app/src/settings/ai.rs +++ b/app/src/settings/ai.rs @@ -834,10 +834,34 @@ pub struct BedrockModelConfig { #[serde(default)] #[schemars(description = "Whether the model supports image/vision input.")] pub vision_supported: bool, + #[serde(default)] + pub invocation_model_id: String, + #[serde(default)] + pub provider_name: String, + #[serde(default)] + pub response_streaming_supported: Option, + #[serde(default)] + pub runtime_supported: Option, } impl settings_value::SettingsValue for BedrockModelConfig {} +impl Default for BedrockModelConfig { + fn default() -> Self { + Self { + enabled: true, + access_verified: false, + model_id: String::new(), + display_name: String::new(), + vision_supported: false, + invocation_model_id: String::new(), + provider_name: String::new(), + response_streaming_supported: None, + runtime_supported: None, + } + } +} + fn default_context_size() -> u32 { 200_000 } diff --git a/app/src/settings/ai_tests.rs b/app/src/settings/ai_tests.rs index f98f61fc..97cd1775 100644 --- a/app/src/settings/ai_tests.rs +++ b/app/src/settings/ai_tests.rs @@ -32,6 +32,22 @@ fn create_test_request_limit_info( } } +#[test] +fn legacy_bedrock_model_settings_default_live_capabilities_without_authorization() { + let model: BedrockModelConfig = serde_json::from_value(serde_json::json!({ + "model_id": "future.new-model", + "display_name": "Saved model", + "enabled": false, + "vision_supported": true + })).unwrap(); + assert_eq!(model.invocation_model_id, ""); + assert_eq!(model.provider_name, ""); + assert_eq!(model.response_streaming_supported, None); + assert_eq!(model.runtime_supported, None); + assert!(!model.access_verified); + assert!(!model.enabled); +} + fn add_ai_enablement_dependencies_for_test(app: &mut App) { app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(UserWorkspaces::default_mock); @@ -446,6 +462,7 @@ fn orchestration_is_enabled_when_ai_is_enabled() { model_id: "test-model".to_string(), display_name: "Test model".to_string(), vision_supported: false, + ..Default::default() }], ctx, ) diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 7d089565..9ae55646 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -1130,6 +1130,20 @@ impl ProviderSetupView { } } + fn bedrock_client_config(&self) -> crate::ai::provider::client::BedrockClientConfig { + crate::ai::provider::client::BedrockClientConfig { + auth_method: self.draft_bedrock.auth_method, + profile: self.draft_bedrock.profile.clone(), + region: self.draft_bedrock.region.clone(), + access_key_id: self.draft_bedrock.access_key_id.clone(), + secret_access_key: self.draft_bedrock.secret_access_key.clone(), + session_token: None, + credentials_expiry: None, + invocation_model_id: None, + cross_region_inference: self.draft_bedrock.cross_region_inference, + } + } + fn begin_discovery(&mut self, ctx: &mut ViewContext) { self.cancel_acp_setup(); self.step = ProviderSetupStep::Discover; @@ -1145,34 +1159,45 @@ impl ProviderSetupView { return; } ProviderSetupProviderType::Bedrock => { - let config = crate::ai::provider::client::BedrockClientConfig { - auth_method: self.draft_bedrock.auth_method, - profile: self.draft_bedrock.profile.clone(), - region: self.draft_bedrock.region.clone(), - access_key_id: self.draft_bedrock.access_key_id.clone(), - secret_access_key: self.draft_bedrock.secret_access_key.clone(), - session_token: None, - credentials_expiry: None, - cross_region_inference: self.draft_bedrock.cross_region_inference, - }; + let config = self.bedrock_client_config(); let requested_config = config.clone(); + let generation = self.acp_discovery_generation; + let verification_generation = LLMPreferences::handle(ctx) + .update(ctx, |preferences, _| preferences.begin_bedrock_verification()); ctx.spawn( async move { crate::ai::provider::discovery::discover_available_models(config).await }, - move |me, result, ctx| match result { + move |me, result, ctx| { + if generation != me.acp_discovery_generation + || me.provider_type != ProviderSetupProviderType::Bedrock + || me.step != ProviderSetupStep::Discover + || me.bedrock_client_config() != requested_config + { + return; + } + match result { Ok(discovery) => { for record in discovery.records { crate::ai::remote_logging::log_view_event(ctx, record); } let models = discovery.models; - LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| { + let accepted = LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| { preferences.record_bedrock_verification( + verification_generation, requested_config, &models, ctx, - ); + ) }); + if !accepted { + me.discovery_state = DiscoveryState::Failed( + "A newer Bedrock connection test superseded this result. Retry the test.".to_owned(), + ); + me.update_next_button(ctx); + ctx.notify(); + return; + } me.draft_bedrock.models = crate::ai::provider::discovery::preserve_model_toggles( &me.draft_bedrock.models, @@ -1190,6 +1215,7 @@ impl ProviderSetupView { me.update_next_button(ctx); ctx.notify(); } + } }, ); return;