remove old okta login authentication refresh

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
+1 -1
View File
@@ -143,7 +143,7 @@ ACP: controller.rs → ResponseStream → acp_output_stream (session-owned lifec
- `convert.rs` — Bedrock request construction and prompt-caching behavior - `convert.rs` — Bedrock request construction and prompt-caching behavior
- `client.rs` — AWS SDK client construction, runtime creation, and independent background streaming calls - `client.rs` — AWS SDK client construction, runtime creation, and independent background streaming calls
- `models.rs` — Model registry and cross-region inference prefix logic - `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`) - `diagnostic.rs` — Debug logging (enabled via `GALAXY_BEDROCK_DIAGNOSTICS=1`)
- `external_config.rs` — Fallback config from Claude Code/OpenCode settings - `external_config.rs` — Fallback config from Claude Code/OpenCode settings
@@ -188,7 +188,9 @@ impl ResponseStream {
// Fall back to Bedrock // Fall back to Bedrock
if *settings.bedrock_enabled.value() { if *settings.bedrock_enabled.value() {
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); 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_refresh_strategy(),
api_key_manager.aws_credentials_state(), api_key_manager.aws_credentials_state(),
); );
+1
View File
@@ -183,6 +183,7 @@ impl CrosscheckReviewer {
if *settings.bedrock_enabled.value() { if *settings.bedrock_enabled.value() {
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings) let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
.with_model_route(model_id, settings.bedrock_models.value())
.with_managed_credentials( .with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(), &api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(), api_key_manager.aws_credentials_state(),
+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::custom_model_routers::{self, CustomModelRouter, ModelConfigError};
use super::execution_profiles::profiles::AIExecutionProfilesModel; use super::execution_profiles::profiles::AIExecutionProfilesModel;
use crate::ai::acp::{acp_launch_fingerprint, acp_provider_selection_identity}; 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::auth_manager::{AuthManager, AuthManagerEvent};
use crate::auth::AuthStateProvider; use crate::auth::AuthStateProvider;
use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind}; use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind};
@@ -616,6 +618,7 @@ pub struct LLMPreferences {
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: bool, chatgpt_subscription_models_refresh_in_flight: bool,
bedrock_models_refresh_in_flight: bool, bedrock_models_refresh_in_flight: bool,
bedrock_verification_generation: u64,
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
bedrock_verified_config: Option<super::provider::client::BedrockClientConfig>, bedrock_verified_config: Option<super::provider::client::BedrockClientConfig>,
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
@@ -796,6 +799,7 @@ impl LLMPreferences {
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: false, chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false, bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
bedrock_verified_config: None, bedrock_verified_config: None,
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
@@ -844,7 +848,9 @@ impl LLMPreferences {
return; return;
} }
self.bedrock_models_refresh_in_flight = true; self.bedrock_models_refresh_in_flight = true;
let generation = self.begin_bedrock_verification();
let config = bedrock_discovery_config(settings); let config = bedrock_discovery_config(settings);
let revision = super::provider::authentication::connection_revision(ctx);
let requested_config = config.clone(); let requested_config = config.clone();
self.bedrock_verified_config = None; self.bedrock_verified_config = None;
self.bedrock_verified_model_ids.clear(); self.bedrock_verified_model_ids.clear();
@@ -860,6 +866,11 @@ impl LLMPreferences {
discovery.models discovery.models
}); });
me.bedrock_models_refresh_in_flight = false; 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() { if !*AISettings::as_ref(ctx).bedrock_enabled.value() {
return; return;
} }
@@ -888,26 +899,37 @@ impl LLMPreferences {
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs); ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
} }
Err(error) => { 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"))] #[cfg(not(target_family = "wasm"))]
pub(crate) fn record_bedrock_verification( pub(crate) fn record_bedrock_verification(
&mut self, &mut self,
generation: u64,
config: super::provider::client::BedrockClientConfig, config: super::provider::client::BedrockClientConfig,
models: &[BedrockModelConfig], models: &[BedrockModelConfig],
ctx: &mut ModelContext<Self>, ctx: &mut ModelContext<Self>,
) { ) -> bool {
if self.bedrock_verification_generation != generation {
return false;
}
self.bedrock_verified_model_ids = self.bedrock_verified_model_ids =
models.iter().map(|model| model.model_id.clone()).collect(); models.iter().map(|model| model.model_id.clone()).collect();
self.bedrock_verified_config = Some(config); self.bedrock_verified_config = Some(config);
self.inject_bedrock_models(ctx); self.inject_bedrock_models(ctx);
self.ensure_default_model_present(); self.ensure_default_model_present();
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs); ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
true
} }
#[cfg(not(target_family = "wasm"))] #[cfg(not(target_family = "wasm"))]
@@ -945,48 +967,14 @@ impl LLMPreferences {
let region = settings.bedrock_region.value().clone(); let region = settings.bedrock_region.value().clone();
let cross_region = *settings.bedrock_cross_region_inference.value(); let cross_region = *settings.bedrock_cross_region_inference.value();
// Check if user wants only 1-hour cache models for model in discovered_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 {
if !model.enabled || !self.bedrock_model_is_verified(&model.model_id, ctx) { if !model.enabled || !self.bedrock_model_is_verified(&model.model_id, ctx) {
continue; continue;
} }
let metadata = model_metadata(&model.model_id); 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!( log::debug!(
"[bedrock] Excluding {} from Agent Mode because its catalog metadata marks it as incompatible with streaming Converse", "[bedrock] Excluding {} from Agent Mode because its catalog metadata marks it as incompatible with streaming Converse",
model.model_id model.model_id
@@ -1046,7 +1034,6 @@ impl LLMPreferences {
// Default agent mode to ANTHROPIC_MODEL from external config if set, // Default agent mode to ANTHROPIC_MODEL from external config if set,
// otherwise Claude Opus 4.6, falling back to the first available model. // 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(); let external_config_for_default = ExternalBedrockConfig::load();
if let Some(id) = external_config_for_default if let Some(id) = external_config_for_default
@@ -1082,7 +1069,11 @@ impl LLMPreferences {
.map(|m| m.id.clone()) .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. // 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()) .or_else(|| self.models_by_feature.coding.choices.first())
.map(|m| m.id.clone()) .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. // Default CLI agent to the first available model.
if let Some(ref mut cli) = self.models_by_feature.cli_agent { 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; cli.default_id = id;
} }
} }
+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 { fn bedrock_model(model_id: &str) -> BedrockModelConfig {
BedrockModelConfig { BedrockModelConfig {
enabled: true, enabled: true,
@@ -166,6 +211,7 @@ fn bedrock_model(model_id: &str) -> BedrockModelConfig {
model_id: model_id.to_string(), model_id: model_id.to_string(),
display_name: model_id.to_string(), display_name: model_id.to_string(),
vision_supported: false, vision_supported: false,
..Default::default()
} }
} }
@@ -235,6 +281,7 @@ fn empty_preferences() -> LLMPreferences {
catalog_refresh_generations: HashMap::new(), catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false, chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false, bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None, bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(), bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(), acp_selections: HashMap::new(),
@@ -451,6 +498,7 @@ fn acp_models_are_injected_only_while_acp_is_enabled() {
catalog_refresh_generations: HashMap::new(), catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false, chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false, bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None, bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(), bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(), acp_selections: HashMap::new(),
+17
View File
@@ -22,10 +22,27 @@ impl BedrockClientConfig {
secret_access_key: settings.bedrock_secret_access_key.value().clone(), secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None, session_token: None,
credentials_expiry: None, credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(), 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( pub(crate) fn with_managed_credentials(
mut self, mut self,
strategy: &AwsCredentialsRefreshStrategy, strategy: &AwsCredentialsRefreshStrategy,
@@ -13,6 +13,7 @@ fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig {
secret_access_key: "configured-secret".to_owned(), secret_access_key: "configured-secret".to_owned(),
session_token: None, session_token: None,
credentials_expiry: None, credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: true, 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)); 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 { fn test_credentials() -> AwsCredentials {
AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None) AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None)
} }
+7 -2
View File
@@ -28,6 +28,7 @@ fn strip_context_marker(model_id: &str) -> String {
pub struct BedrockClient { pub struct BedrockClient {
runtime_client: BedrockRuntimeClient, runtime_client: BedrockRuntimeClient,
region: String, region: String,
invocation_model_id: Option<String>,
} }
#[derive(Debug, Clone, PartialEq)] #[derive(Debug, Clone, PartialEq)]
@@ -39,6 +40,7 @@ pub struct BedrockClientConfig {
pub secret_access_key: String, pub secret_access_key: String,
pub session_token: Option<String>, pub session_token: Option<String>,
pub credentials_expiry: Option<SystemTime>, pub credentials_expiry: Option<SystemTime>,
pub invocation_model_id: Option<String>,
pub cross_region_inference: bool, pub cross_region_inference: bool,
} }
@@ -155,6 +157,7 @@ impl BedrockClient {
Ok(Self { Ok(Self {
runtime_client, runtime_client,
region, region,
invocation_model_id: config.invocation_model_id,
}) })
} }
@@ -185,14 +188,16 @@ impl BedrockClient {
tool_result_archive: Vec<ConversationMessage>, tool_result_archive: Vec<ConversationMessage>,
) -> Result<LegacyResponseStream, BedrockError> { ) -> Result<LegacyResponseStream, BedrockError> {
let base_model_id = strip_context_marker(model_id); 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) apply_cross_region_prefix(&base_model_id, &self.region)
} else { } else {
base_model_id base_model_id
}; };
let external_config = ExternalBedrockConfig::load(); 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 { if !caching_config.enabled {
log::info!("[bedrock] Prompt caching disabled (DISABLE_PROMPT_CACHING=1)"); log::info!("[bedrock] Prompt caching disabled (DISABLE_PROMPT_CACHING=1)");
+16
View File
@@ -38,6 +38,22 @@ impl CachingConfig {
extended_ttl_requested: external.enable_prompt_caching_1h, 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 { pub struct ConvertedRequest {
+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 aws_sdk_bedrockruntime::types::{ContentBlock, ConversationRole, SystemContentBlock, Tool};
use serde_json::json; use serde_json::json;
+168 -25
View File
@@ -1,17 +1,25 @@
//! AWS Bedrock control-plane discovery. //! AWS Bedrock control-plane discovery.
//! //!
//! The foundation-model catalog is only a candidate list. Every candidate is //! AWS supplies the model inventory and invocation routes. Bundled metadata
//! checked with `GetFoundationModelAvailability` before it is offered to the //! enriches known capabilities but never determines inventory membership.
//! user or persisted in Galaxy settings.
use std::collections::{BTreeMap, HashSet};
use std::future::Future;
use aws_config::BehaviorVersion; 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_bedrock::Client;
use aws_sdk_bedrockruntime::config::Region; use aws_sdk_bedrockruntime::config::Region;
use futures::{stream, StreamExt}; 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 serde_json::json;
use super::client::{BedrockClientConfig, BedrockError}; use super::client::{BedrockClientConfig, BedrockError};
use super::models::foundation_model_id;
use crate::ai::remote_logging::{RemoteLogLevel, RemoteLogRecord}; use crate::ai::remote_logging::{RemoteLogLevel, RemoteLogRecord};
use crate::settings::ai::BedrockModelConfig; use crate::settings::ai::BedrockModelConfig;
@@ -29,9 +37,66 @@ pub(crate) async fn discover_available_models(
.await .await
.map_err(|error| error.to_string())?; .map_err(|error| error.to_string())?;
let client = Client::new(&aws_config); let client = Client::new(&aws_config);
// The bundled catalog defines candidates; AWS defines account access. let inventory = client.list_foundation_models().send().await
// Never add arbitrary entries from ListFoundationModels to this catalog. .map_err(|error| format!("Unable to list Bedrock models: {}", sanitized_error(&error)))?;
let candidates = catalog_models().to_vec(); 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 sweep_id = uuid::Uuid::new_v4().to_string();
let region = aws_config let region = aws_config
.region() .region()
@@ -39,30 +104,30 @@ pub(crate) async fn discover_available_models(
.unwrap_or("UNKNOWN") .unwrap_or("UNKNOWN")
.to_owned(); .to_owned();
log::info!( log::info!(
"[bedrock] Starting availability sweep {sweep_id}: {} catalog models, region={region}", "[bedrock] Starting availability sweep {sweep_id}: {} live candidates, region={region}",
candidates.len() candidates.len()
); );
let checks = candidates.into_iter().map(|metadata| { let checks = candidates.into_iter().map(|mut model| {
let client = client.clone(); let client = client.clone();
let sweep_id = sweep_id.clone(); let sweep_id = sweep_id.clone();
let region = region.clone(); let region = region.clone();
async move { async move {
let model_id = metadata.model_id.as_str(); let model_id = model.model_id.clone();
let availability = match client let availability = match client
.get_foundation_model_availability() .get_foundation_model_availability()
.model_id(model_id) .model_id(&model_id)
.send() .send()
.await .await
{ {
Ok(availability) => availability, Ok(availability) => availability,
Err(error) => { 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}"); log::warn!("[bedrock] Availability sweep {sweep_id}: model={model_id}, region={region}, decision=excluded, error={error}");
return (None, RemoteLogRecord { return (None, RemoteLogRecord {
level: RemoteLogLevel::Warn, level: RemoteLogLevel::Warn,
message: "Bedrock model availability check failed".to_owned(), 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}), 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()); 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 entitlement = availability.entitlement_availability().as_str();
let region_status = availability.region_availability().as_str(); let region_status = availability.region_availability().as_str();
let access_allowed = model_availability_is_usable(agreement, authorization, entitlement, region_status); 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 allowed = access_allowed && runtime_supported;
let decision = if allowed { "allowed" } else { "excluded" }; 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")); 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(), 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}), 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); } model.access_verified = allowed;
(allowed.then_some(model), record, true)
(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)
} }
}); });
let results = stream::iter(checks) let results = stream::iter(checks)
.buffer_unordered(AVAILABILITY_CHECK_CONCURRENCY) .buffer_unordered(AVAILABILITY_CHECK_CONCURRENCY)
.collect::<Vec<_>>() .collect::<Vec<_>>()
.await; .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<_>>(); let mut models = models.into_iter().flatten().collect::<Vec<_>>();
models.sort_by(|left, right| left.display_name.cmp(&right.display_name)); models.sort_by(|left, right| left.display_name.cmp(&right.display_name));
log::info!( log::info!(
@@ -106,11 +168,84 @@ pub(crate) async fn discover_available_models(
records.push(RemoteLogRecord { records.push(RemoteLogRecord {
level: RemoteLogLevel::Info, level: RemoteLogLevel::Info,
message: "Bedrock availability sweep completed".to_owned(), 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 }) 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( fn model_availability_is_usable(
agreement_status: Option<&str>, agreement_status: Option<&str>,
authorization_status: &str, authorization_status: &str,
@@ -133,6 +268,14 @@ pub(crate) fn preserve_model_toggles(
.find(|previous| previous.model_id == model.model_id) .find(|previous| previous.model_id == model.model_id)
{ {
model.enabled = previous.enabled; 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 discovered
+140
View File
@@ -1,4 +1,143 @@
use super::model_availability_is_usable; 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] #[test]
fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() { 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, vision_supported: false,
enabled: false, enabled: false,
access_verified: true, access_verified: true,
..Default::default()
}; };
let saved = serde_json::to_value(&model).unwrap(); let saved = serde_json::to_value(&model).unwrap();
assert!(saved.get("access_verified").is_none()); assert!(saved.get("access_verified").is_none());
+1
View File
@@ -269,6 +269,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(), secret_access_key: String::new(),
session_token: None, session_token: None,
credentials_expiry: None, credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false, cross_region_inference: false,
}) })
} }
+1
View File
@@ -146,6 +146,7 @@ fn parse_claude_code_model_map(
model_id: arn, model_id: arn,
display_name, display_name,
vision_supported: true, vision_supported: true,
..Default::default()
} }
}) })
.collect() .collect()
+1
View File
@@ -24,6 +24,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
secret_access_key: String::new(), secret_access_key: String::new(),
session_token: None, session_token: None,
credentials_expiry: None, credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false, cross_region_inference: false,
}) })
} }
+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 { pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") { if model_id.starts_with("arn:") {
return model_id.to_string(); return model_id.to_string();
+4 -3
View File
@@ -346,14 +346,15 @@ pub(crate) async fn provider_runtime_for_request(
}, },
crate::ai::provider::ProviderConfig::Bedrock(config) => { crate::ai::provider::ProviderConfig::Bedrock(config) => {
let max_output_tokens = Some(bedrock_max_output_tokens(&model)); 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 = 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?; let client = BedrockClient::from_config(config).await?;
Arc::new(BedrockRuntime::from_aws_client( Arc::new(BedrockRuntime::from_aws_client(
client.runtime_client(), client.runtime_client(),
BedrockRigConfig { BedrockRigConfig {
model, model: invocation_model,
region: client.region().to_string(), region: client.region().to_string(),
cross_region_inference, cross_region_inference,
prompt_caching: caching_config.enabled, prompt_caching: caching_config.enabled,
+24
View File
@@ -834,10 +834,34 @@ pub struct BedrockModelConfig {
#[serde(default)] #[serde(default)]
#[schemars(description = "Whether the model supports image/vision input.")] #[schemars(description = "Whether the model supports image/vision input.")]
pub vision_supported: bool, pub vision_supported: bool,
#[serde(default)]
pub invocation_model_id: String,
#[serde(default)]
pub provider_name: String,
#[serde(default)]
pub response_streaming_supported: Option<bool>,
#[serde(default)]
pub runtime_supported: Option<bool>,
} }
impl settings_value::SettingsValue for BedrockModelConfig {} 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 { fn default_context_size() -> u32 {
200_000 200_000
} }
+17
View File
@@ -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) { fn add_ai_enablement_dependencies_for_test(app: &mut App) {
app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(|_| AuthStateProvider::new_for_test());
app.add_singleton_model(UserWorkspaces::default_mock); 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(), model_id: "test-model".to_string(),
display_name: "Test model".to_string(), display_name: "Test model".to_string(),
vision_supported: false, vision_supported: false,
..Default::default()
}], }],
ctx, ctx,
) )
+39 -13
View File
@@ -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>) { fn begin_discovery(&mut self, ctx: &mut ViewContext<Self>) {
self.cancel_acp_setup(); self.cancel_acp_setup();
self.step = ProviderSetupStep::Discover; self.step = ProviderSetupStep::Discover;
@@ -1145,34 +1159,45 @@ impl ProviderSetupView {
return; return;
} }
ProviderSetupProviderType::Bedrock => { ProviderSetupProviderType::Bedrock => {
let config = crate::ai::provider::client::BedrockClientConfig { let config = self.bedrock_client_config();
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 requested_config = config.clone(); 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( ctx.spawn(
async move { async move {
crate::ai::provider::discovery::discover_available_models(config).await 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) => { Ok(discovery) => {
for record in discovery.records { for record in discovery.records {
crate::ai::remote_logging::log_view_event(ctx, record); crate::ai::remote_logging::log_view_event(ctx, record);
} }
let models = discovery.models; let models = discovery.models;
LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| { let accepted = LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| {
preferences.record_bedrock_verification( preferences.record_bedrock_verification(
verification_generation,
requested_config, requested_config,
&models, &models,
ctx, 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 = me.draft_bedrock.models =
crate::ai::provider::discovery::preserve_model_toggles( crate::ai::provider::discovery::preserve_model_toggles(
&me.draft_bedrock.models, &me.draft_bedrock.models,
@@ -1190,6 +1215,7 @@ impl ProviderSetupView {
me.update_next_button(ctx); me.update_next_button(ctx);
ctx.notify(); ctx.notify();
} }
}
}, },
); );
return; return;